easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/utils/config_util.py

@@ -0,0 +1,652 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""Functions for reading and updating configuration files.
+
+Such as Hyper parameter tuning or automatic feature expanding.
+"""
+
+import datetime
+import json
+import logging
+import os
+import re
+import sys
+
+import numpy as np
+import six
+import tensorflow as tf
+from google.protobuf import json_format
+from google.protobuf import text_format
+from tensorflow.python.lib.io import file_io
+
+from easy_rec.python.protos import pipeline_pb2
+from easy_rec.python.protos.feature_config_pb2 import FeatureConfig
+from easy_rec.python.utils import pai_util
+from easy_rec.python.utils.hive_utils import HiveUtils
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+def search_pipeline_config(directory):
+  dir_list = []
+  for root, dirs, files in tf.gfile.Walk(directory):
+    for f in files:
+      _, ext = os.path.splitext(f)
+      if ext == '.config':
+        dir_list.append(os.path.join(root, f))
+  if len(dir_list) == 0:
+    raise ValueError('config is not found in directory %s' % directory)
+  elif len(dir_list) > 1:
+    raise ValueError('config saved model found in directory %s' % directory)
+  logging.info('use pipeline config: %s' % dir_list[0])
+  return dir_list[0]
+
+
+def get_configs_from_pipeline_file(pipeline_config_path, auto_expand=True):
+  """Reads config from a file containing pipeline_pb2.EasyRecConfig.
+
+  Args:
+    pipeline_config_path: Path to pipeline_pb2.EasyRecConfig text
+      proto.
+
+  Returns:
+    Dictionary of configuration objects. Keys are `model`, `train_config`,
+      `train_input_config`, `eval_config`, `eval_input_config`. Value are the
+      corresponding config objects.
+  """
+  if isinstance(pipeline_config_path, pipeline_pb2.EasyRecConfig):
+    return pipeline_config_path
+
+  assert tf.gfile.Exists(
+      pipeline_config_path
+  ), 'pipeline_config_path [%s] not exists' % pipeline_config_path
+
+  pipeline_config = pipeline_pb2.EasyRecConfig()
+  with tf.gfile.GFile(pipeline_config_path, 'r') as f:
+    config_str = f.read()
+    if pipeline_config_path.endswith('.config'):
+      text_format.Merge(config_str, pipeline_config)
+    elif pipeline_config_path.endswith('.json'):
+      json_format.Parse(config_str, pipeline_config)
+    else:
+      assert False, 'invalid file format(%s), currently support formats: .config(prototxt) .json' % pipeline_config_path
+
+  if auto_expand:
+    return auto_expand_share_feature_configs(pipeline_config)
+  else:
+    return pipeline_config
+
+
+def auto_expand_share_feature_configs(pipeline_config):
+  feature_configs = get_compatible_feature_configs(pipeline_config)
+  for share_config in feature_configs:
+    if len(share_config.shared_names) == 0:
+      continue
+
+    # auto expand all shared_names
+    input_names = []
+    for input_name in share_config.shared_names:
+      if pipeline_config.data_config.auto_expand_input_fields:
+        input_names.extend(auto_expand_names(input_name))
+      else:
+        input_names.append(input_name)
+
+    # make a clean copy
+    while len(share_config.shared_names) > 0:
+      share_config.shared_names.pop()
+
+    fea_config = FeatureConfig()
+    fea_config.CopyFrom(share_config)
+    while len(fea_config.input_names) > 0:
+      fea_config.input_names.pop()
+
+    # generate for each item in input_name
+    for tmp_name in input_names:
+      tmp_config = FeatureConfig()
+      tmp_config.CopyFrom(fea_config)
+      tmp_config.input_names.append(tmp_name)
+      if pipeline_config.feature_configs:
+        pipeline_config.feature_configs.append(tmp_config)
+      else:
+        pipeline_config.feature_config.features.append(tmp_config)
+  return pipeline_config
+
+
+# def auto_expand_names(input_name):
+#   """Auto expand field[1-3] to field1, field2, field3.
+#
+#   Args:
+#     input_name: a string pattern like field[1-3]
+#
+#   Returns:
+#     a string list of the expanded names
+#   """
+#   match_obj = re.match(r'([a-zA-Z_]+)\[([0-9]+)-([0-9]+)\]', input_name)
+#   if match_obj:
+#     prefix = match_obj.group(1)
+#     sid = int(match_obj.group(2))
+#     eid = int(match_obj.group(3)) + 1
+#     input_name = ['%s%d' % (prefix, tid) for tid in range(sid, eid)]
+#   else:
+#     input_name = [input_name]
+#   return input_name
+
+
+def auto_expand_names(input_name):
+  """Auto expand field[1-3] to field1, field2, field3.
+
+  Args:
+    input_name: a string pattern like field[1-3]
+
+  Returns:
+    a string list of the expanded names
+  Todo:
+    could be extended to support more complicated patterns
+  """
+  flag = 1
+  if input_name.endswith(']'):
+    match_obj = re.match(r'([a-zA-Z_]+)\[([0-9]+)-([0-9]+)\]', input_name)
+  else:
+    flag = 2
+    match_obj = re.match(r'([a-zA-Z_]+)\[([0-9]+)-([0-9]+)\]([a-zA-Z_]+)',
+                         input_name)
+  if match_obj:
+    prefix = match_obj.group(1)
+    sid = int(match_obj.group(2))
+    eid = int(match_obj.group(3)) + 1
+    if flag == 2:
+      endfix = match_obj.group(4)
+      input_name = ['%s%d%s' % (prefix, tid, endfix) for tid in range(sid, eid)]
+    else:
+      input_name = ['%s%d' % (prefix, tid) for tid in range(sid, eid)]
+
+  else:
+    input_name = [input_name]
+  return input_name
+
+
+def create_pipeline_proto_from_configs(configs):
+  """Creates a pipeline_pb2.EasyRecConfig from configs dictionary.
+
+  This function performs the inverse operation of
+  create_configs_from_pipeline_proto().
+
+  Args:
+    configs: Dictionary of configs. See get_configs_from_pipeline_file().
+
+  Returns:
+    A fully populated pipeline_pb2.EasyRecConfig.
+  """
+  pipeline_config = pipeline_pb2.EasyRecConfig()
+  pipeline_config.model.CopyFrom(configs['model'])
+  pipeline_config.train_config.CopyFrom(configs['train_config'])
+  pipeline_config.train_input_reader.CopyFrom(configs['train_input_config'])
+  pipeline_config.eval_config.CopyFrom(configs['eval_config'])
+  pipeline_config.eval_input_reader.CopyFrom(configs['eval_input_config'])
+  if 'graph_rewriter_config' in configs:
+    pipeline_config.graph_rewriter.CopyFrom(configs['graph_rewriter_config'])
+  return pipeline_config
+
+
+def save_pipeline_config(pipeline_config,
+                         directory,
+                         filename='pipeline.config'):
+  """Saves a pipeline config text file to disk.
+
+  Args:
+    pipeline_config: A pipeline_pb2.TrainEvalPipelineConfig.
+    directory: The model directory into which the pipeline config file will be
+      saved.
+    filename: pipelineconfig filename
+  """
+  if not file_io.file_exists(directory):
+    file_io.recursive_create_dir(directory)
+  pipeline_config_path = os.path.join(directory, filename)
+  # as_utf8=True to make sure pbtxt is human readable when string contains chinese
+  save_message(pipeline_config, pipeline_config_path)
+
+
+def _get_basic_types():
+  dtypes = [
+      bool, int, str, float,
+      type(u''), np.float16, np.float32, np.float64, np.char, np.byte, np.uint8,
+      np.int8, np.int16, np.uint16, np.uint32, np.int32, np.uint64, np.int64,
+      bool, str
+  ]
+  if six.PY2:
+    dtypes.append(long)  # noqa: F821
+
+  return dtypes
+
+
+def edit_config(pipeline_config, edit_config_json):
+  """Update params specified by automl.
+
+  Args:
+    pipeline_config: EasyRecConfig
+    edit_config_json: edit config json
+  """
+
+  def _type_convert(proto, val, parent=None):
+    if type(val) != type(proto):
+      try:
+        if isinstance(proto, bool):
+          assert val in ['True', 'true', 'False', 'false']
+          val = val in ['True', 'true']
+        else:
+          val = type(proto)(val)
+      except ValueError as ex:
+        if parent is None:
+          raise ex
+        assert isinstance(proto, int)
+        val = getattr(parent, val)
+        assert isinstance(val, int)
+    return val
+
+  def _get_attr(obj, attr, only_last=False):
+    # only_last means we only return the last element in paths array
+    attr_toks = [x.strip() for x in attr.split('.') if x != '']
+    paths = []
+    objs = [obj]
+    nobjs = []
+    for key in attr_toks:
+      # clear old paths to clear new paths
+      paths = []
+      for obj in objs:
+        if '[' in key:
+          pos = key.find('[')
+          name, cond = key[:pos], key[pos + 1:]
+          cond = cond[:-1]
+          update_objs = getattr(obj, name)
+          # select all update_objs
+          if cond == ':':
+            for tid, update_obj in enumerate(update_objs):
+              paths.append((obj, update_obj, None, tid))
+              nobjs.append(update_obj)
+            continue
+
+          # select by range update_objs[1:10]
+          if ':' in cond:
+            colon_pos = cond.find(':')
+            sid = cond[:colon_pos]
+            if len(sid) == 0:
+              sid = 0
+            else:
+              sid = int(sid)
+            eid = cond[(colon_pos + 1):]
+            if len(eid) == 0:
+              eid = len(update_objs)
+            else:
+              eid = int(eid)
+            for tid, update_obj in enumerate(update_objs[sid:eid]):
+              paths.append((obj, update_obj, None, tid + sid))
+              nobjs.append(update_obj)
+            continue
+
+          # for simple index update_objs[0]
+          try:
+            obj_id = int(cond)
+            obj = update_objs[obj_id]
+            paths.append((obj, update_objs, None, obj_id))
+            nobjs.append(obj)
+            continue
+          except ValueError:
+            pass
+
+          # for complex conditions a[optimizer.lr=20]
+          op_func_map = {
+              '>=': lambda x, y: x >= y,
+              '<=': lambda x, y: x <= y,
+              '<': lambda x, y: x < y,
+              '>': lambda x, y: x > y,
+              '=': lambda x, y: x == y
+          }
+          cond_key = None
+          cond_val = None
+          op_func = None
+          for op in ['>=', '<=', '>', '<', '=']:
+            tmp_pos = cond.rfind(op)
+            if tmp_pos != -1:
+              cond_key = cond[:tmp_pos]
+              cond_val = cond[(tmp_pos + len(op)):]
+              op_func = op_func_map[op]
+              break
+
+          assert cond_key is not None, 'invalid cond: %s' % cond
+          assert cond_val is not None, 'invalid cond: %s' % cond
+
+          for tid, update_obj in enumerate(update_objs):
+            tmp, tmp_parent, _, _ = _get_attr(
+                update_obj, cond_key, only_last=True)
+
+            cond_val = _type_convert(tmp, cond_val, tmp_parent)
+
+            if op_func(tmp, cond_val):
+              obj_id = tid
+              paths.append((update_obj, update_objs, None, obj_id))
+              nobjs.append(update_obj)
+        else:
+          sub_obj = getattr(obj, key)
+          paths.append((sub_obj, obj, key, -1))
+          nobjs.append(sub_obj)
+      # exchange to prepare for parsing next token
+      objs = nobjs
+      nobjs = []
+    if only_last:
+      return paths[-1]
+    else:
+      return paths
+
+  for param_keys in edit_config_json:
+    # multiple keys/vals combination
+    param_vals = edit_config_json[param_keys]
+    param_vals = [x.strip() for x in str(param_vals).split(';')]
+    param_keys = [x.strip() for x in str(param_keys).split(';')]
+    for param_key, param_val in zip(param_keys, param_vals):
+      update_obj = pipeline_config
+      tmp_paths = _get_attr(update_obj, param_key)
+      # update a set of objs
+      for tmp_val, tmp_obj, tmp_name, tmp_id in tmp_paths:
+        # list and dict are not basic types, must be handle separately
+        basic_types = _get_basic_types()
+        if type(tmp_val) in basic_types:
+          # simple type cast
+          tmp_val = _type_convert(tmp_val, param_val, tmp_obj)
+          if tmp_name is None:
+            tmp_obj[tmp_id] = tmp_val
+          else:
+            setattr(tmp_obj, tmp_name, tmp_val)
+        elif 'Scalar' in str(type(tmp_val)) and 'ClearField' in dir(tmp_obj):
+          tmp_obj.ClearField(tmp_name)
+          text_format.Parse('%s:%s' % (tmp_name, param_val), tmp_obj)
+        else:
+          tmp_val.Clear()
+          param_val = param_val.strip()
+          if param_val.startswith('{') and param_val.endswith('}'):
+            param_val = param_val[1:-1]
+          text_format.Parse(param_val, tmp_val)
+
+  return pipeline_config
+
+
+def save_message(protobuf_message, filename):
+  """Saves a pipeline config text file to disk.
+
+  Args:
+    protobuf_message: A pipeline_pb2.TrainEvalPipelineConfig.
+    filename: pipeline config filename
+  """
+  directory, _ = os.path.split(filename)
+  if not file_io.file_exists(directory):
+    file_io.recursive_create_dir(directory)
+  # as_utf8=True to make sure pbtxt is human readable when string contains chinese
+  config_text = text_format.MessageToString(protobuf_message, as_utf8=True)
+  with tf.gfile.Open(filename, 'wb') as f:
+    logging.info('Writing protobuf message file to %s', filename)
+    f.write(config_text)
+
+
+def add_boundaries_to_config(pipeline_config, tables):
+  import common_io
+
+  feature_boundaries_info = {}
+  reader = common_io.table.TableReader(tables, selected_cols='feature,json')
+  while True:
+    try:
+      record = reader.read()
+      raw_info = json.loads(record[0][1])
+      bin_info = []
+      for info in raw_info['bin']['norm'][:-1]:
+        split_point = float(info['value'].split(',')[1][:-1])
+        bin_info.append(split_point)
+      feature_boundaries_info[record[0][0]] = bin_info
+    except common_io.exception.OutOfRangeException:
+      reader.close()
+      break
+
+  logging.info('feature boundaries: %s' % feature_boundaries_info)
+  feature_configs = get_compatible_feature_configs(pipeline_config)
+  for feature_config in feature_configs:
+    feature_name = feature_config.input_names[0]
+    if feature_name in feature_boundaries_info:
+      if feature_config.feature_type != feature_config.SequenceFeature:
+        logging.info(
+            'feature = {0}, type = {1}, will turn to RawFeature.'.format(
+                feature_name, feature_config.feature_type))
+        feature_config.feature_type = feature_config.RawFeature
+      feature_config.hash_bucket_size = 0
+      feature_config.ClearField('boundaries')
+      feature_config.boundaries.extend(feature_boundaries_info[feature_name])
+      logging.info('edited %s' % feature_name)
+
+
+def get_compatible_feature_configs(pipeline_config):
+  if pipeline_config.feature_configs:
+    feature_configs = pipeline_config.feature_configs
+  else:
+    feature_configs = pipeline_config.feature_config.features
+  return feature_configs
+
+
+def parse_time(time_data):
+  """Parse time string to timestamp.
+
+  Args:
+    time_data: could be two formats: '%Y%m%d %H:%M:%S' or '%s'
+  Return:
+    timestamp: int
+  """
+  if isinstance(time_data, str) or isinstance(time_data, type(u'')):
+    if len(time_data) == 17:
+      return int(
+          datetime.datetime.strptime(time_data,
+                                     '%Y%m%d %H:%M:%S').strftime('%s'))
+    elif len(time_data) == 10:
+      return int(time_data)
+    else:
+      assert 'invalid time string: %s' % time_data
+  else:
+    return int(time_data)
+
+
+def search_fg_json(directory):
+  dir_list = []
+  for root, dirs, files in tf.gfile.Walk(directory):
+    for f in files:
+      _, ext = os.path.splitext(f)
+      if ext == '.json':
+        dir_list.append(os.path.join(root, f))
+  if len(dir_list) == 0:
+    return None
+  elif len(dir_list) > 1:
+    raise ValueError('fg.json found in directory %s' % directory)
+  logging.info('use fg.json: %s' % dir_list[0])
+  return dir_list[0]
+
+
+def get_input_name_from_fg_json(fg_json):
+  if not fg_json:
+    return []
+  input_names = []
+  for fea in fg_json['features']:
+    if 'feature_name' in fea:
+      if 'stub_type' in fea and fea['stub_type']:
+        continue
+      input_names.append(fea['feature_name'])
+    elif 'sequence_name' in fea:
+      sequence_name = fea['sequence_name']
+      for seq_fea in fea['features']:
+        assert 'feature_name' in seq_fea
+        if 'stub_type' in seq_fea and seq_fea['stub_type']:
+          continue
+        feature_name = seq_fea['feature_name']
+        input_names.append(sequence_name + '__' + feature_name)
+  return input_names
+
+
+def get_train_input_path(pipeline_config):
+  input_name = pipeline_config.WhichOneof('train_path')
+  return getattr(pipeline_config, input_name)
+
+
+def get_eval_input_path(pipeline_config):
+  input_name = pipeline_config.WhichOneof('eval_path')
+  return getattr(pipeline_config, input_name)
+
+
+def get_model_dir_path(pipeline_config):
+  model_dir = pipeline_config.model_dir
+  return model_dir
+
+
+def set_train_input_path(pipeline_config, train_input_path):
+  if pipeline_config.WhichOneof('train_path') == 'hive_train_input':
+    if isinstance(train_input_path, list):
+      assert len(
+          train_input_path
+      ) <= 1, 'only support one hive_train_input.table_name when hive input'
+      pipeline_config.hive_train_input.table_name = train_input_path[0]
+    else:
+      assert len(
+          train_input_path.split(',')
+      ) <= 1, 'only support one hive_train_input.table_name when hive input'
+      pipeline_config.hive_train_input.table_name = train_input_path
+    logging.info('update hive_train_input.table_name to %s' %
+                 pipeline_config.hive_train_input.table_name)
+
+  elif pipeline_config.WhichOneof('train_path') == 'kafka_train_input':
+    if isinstance(train_input_path, list):
+      pipeline_config.kafka_train_input = ','.join(train_input_path)
+    else:
+      pipeline_config.kafka_train_input = train_input_path
+  elif pipeline_config.WhichOneof('train_path') == 'parquet_train_input':
+    if isinstance(train_input_path, list):
+      pipeline_config.parquet_train_input = ','.join(train_input_path)
+    else:
+      pipeline_config.parquet_train_input = train_input_path
+  else:
+    if isinstance(train_input_path, list):
+      pipeline_config.train_input_path = ','.join(train_input_path)
+    else:
+      pipeline_config.train_input_path = train_input_path
+    logging.info('update train_input_path to %s' %
+                 pipeline_config.train_input_path)
+  return pipeline_config
+
+
+def set_eval_input_path(pipeline_config, eval_input_path):
+  if pipeline_config.WhichOneof('eval_path') == 'hive_eval_input':
+    if isinstance(eval_input_path, list):
+      assert len(
+          eval_input_path
+      ) <= 1, 'only support one hive_eval_input.table_name when hive input'
+      pipeline_config.hive_eval_input.table_name = eval_input_path[0]
+    else:
+      assert len(
+          eval_input_path.split(',')
+      ) <= 1, 'only support one hive_eval_input.table_name when hive input'
+      pipeline_config.hive_eval_input.table_name = eval_input_path
+    logging.info('update hive_eval_input.table_name to %s' %
+                 pipeline_config.hive_eval_input.table_name)
+  elif pipeline_config.WhichOneof('eval_path') == 'parquet_eval_input':
+    if isinstance(eval_input_path, list):
+      pipeline_config.parquet_eval_input = ','.join(eval_input_path)
+    else:
+      pipeline_config.parquet_eval_input = eval_input_path
+  elif pipeline_config.WhichOneof('eval_path') == 'kafka_eval_input':
+    if isinstance(eval_input_path, list):
+      pipeline_config.kafka_eval_input = ','.join(eval_input_path)
+    else:
+      pipeline_config.kafka_eval_input = eval_input_path
+  else:
+    if isinstance(eval_input_path, list):
+      pipeline_config.eval_input_path = ','.join(eval_input_path)
+    else:
+      pipeline_config.eval_input_path = eval_input_path
+    logging.info('update eval_input_path to %s' %
+                 pipeline_config.eval_input_path)
+  return pipeline_config
+
+
+def process_data_path(data_path, hive_util):
+  if data_path.startswith('hdfs://'):
+    return data_path
+  if re.match(r'(.*)\.(.*)', data_path):
+    hdfs_path = hive_util.get_table_location(data_path)
+    assert hdfs_path, "Can't find hdfs path of %s" % data_path
+    logging.info('update %s to %s' % (data_path, hdfs_path))
+    return hdfs_path
+  return data_path
+
+
+def process_neg_sampler_data_path(pipeline_config):
+  # replace neg_sampler hive table => hdfs path
+  if pai_util.is_on_pai():
+    return
+  if not pipeline_config.data_config.HasField('sampler'):
+    return
+  # not using hive, so not need to process it
+  if pipeline_config.WhichOneof('train_path') != 'hive_train_input':
+    return
+  hive_util = HiveUtils(
+      data_config=pipeline_config.data_config,
+      hive_config=pipeline_config.hive_train_input)
+  sampler_type = pipeline_config.data_config.WhichOneof('sampler')
+  sampler_config = getattr(pipeline_config.data_config, sampler_type)
+  if hasattr(sampler_config, 'input_path'):
+    sampler_config.input_path = process_data_path(sampler_config.input_path,
+                                                  hive_util)
+  if hasattr(sampler_config, 'user_input_path'):
+    sampler_config.user_input_path = process_data_path(
+        sampler_config.user_input_path, hive_util)
+  if hasattr(sampler_config, 'item_input_path'):
+    sampler_config.item_input_path = process_data_path(
+        sampler_config.item_input_path, hive_util)
+  if hasattr(sampler_config, 'pos_edge_input_path'):
+    sampler_config.pos_edge_input_path = process_data_path(
+        sampler_config.pos_edge_input_path, hive_util)
+  if hasattr(sampler_config, 'hard_neg_edge_input_path'):
+    sampler_config.hard_neg_edge_input_path = process_data_path(
+        sampler_config.hard_neg_edge_input_path, hive_util)
+
+
+def parse_extra_config_param(extra_args, edit_config_json):
+  arg_num = len(extra_args)
+  arg_id = 0
+  while arg_id < arg_num:
+    if extra_args[arg_id].startswith('--data_config.') or    \
+       extra_args[arg_id].startswith('--train_config.') or   \
+       extra_args[arg_id].startswith('--feature_config.') or \
+       extra_args[arg_id].startswith('--model_config.') or   \
+       extra_args[arg_id].startswith('--export_config.') or  \
+       extra_args[arg_id].startswith('--eval_config.'):
+      tmp_arg = extra_args[arg_id][2:]
+      if '=' in tmp_arg:
+        sep_pos = tmp_arg.find('=')
+        k = tmp_arg[:sep_pos]
+        v = tmp_arg[(sep_pos + 1):]
+        v = v.strip(' "\'')
+        edit_config_json[k] = v
+        arg_id += 1
+      elif arg_id + 1 < len(extra_args):
+        edit_config_json[tmp_arg] = extra_args[arg_id + 1].strip(' "\'')
+        arg_id += 2
+      else:
+        logging.error('missing value for arg: %s' % extra_args[arg_id])
+        sys.exit(1)
+    else:
+      logging.error('unknown args: %s' % extra_args[arg_id])
+      sys.exit(1)
+
+
+def process_multi_file_input_path(sampler_config_input_path):
+
+  if '*' in sampler_config_input_path:
+    input_path = ','.join(
+        file_path
+        for file_path in tf.gfile.Glob(sampler_config_input_path.split(',')))
+  else:
+    input_path = sampler_config_input_path
+
+  return input_path

easy_rec/python/utils/constant.py

@@ -0,0 +1,43 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os
+
+SAMPLE_WEIGHT = 'SAMPLE_WEIGHT'
+
+DENSE_UPDATE_VARIABLES = 'DENSE_UPDATE_VARIABLES'
+
+SPARSE_UPDATE_VARIABLES = 'SPARSE_UPDATE_VARIABLES'
+ENABLE_AVX_STR_SPLIT = 'ENABLE_AVX_STR_SPLIT'
+
+# Environment variables to control whether to sort
+# feature columns by name, by default sort is not
+# enabled. The flag is set for backward compatibility.
+SORT_COL_BY_NAME = 'SORT_COL_BY_NAME'
+
+# arithmetic_optimization causes significant slow training
+# of a test case:
+#   train_eval_test.TrainEvalTest.test_train_parquet
+NO_ARITHMETRIC_OPTI = 'NO_ARITHMETRIC_OPTI'
+
+# shard embedding var_name collection
+EmbeddingParallel = 'EmbeddingParallel'
+
+# environ variable to force embedding placement on cpu
+EmbeddingOnCPU = 'place_embedding_on_cpu'
+
+# clear ps counter queue at start to reuse existing ps
+ClearPsCounterQueue = 'CLEAR_PS_COUNTER_QUEUE'
+
+
+def enable_avx_str_split():
+  os.environ[ENABLE_AVX_STR_SPLIT] = '1'
+
+
+def has_avx_str_split():
+  return ENABLE_AVX_STR_SPLIT in os.environ and os.environ[
+      ENABLE_AVX_STR_SPLIT] == '1'
+
+
+def disable_avx_str_split():
+  del os.environ[ENABLE_AVX_STR_SPLIT]