Function set_env_for_DeepRec

modelzoo/dlrm/train.py:915–933

Set some ENV for these DeepRec's features enabled by ENV. More Detail information is shown in https://deeprec.readthedocs.io/zh/latest/index.html. START_STATISTIC_STEP & STOP_STATISTIC_STEP: On CPU platform, DeepRec supports memory optimization in both stand-alone and distribute

()

# This func is used to set ENV and enable these features.
# A triple quotes comment is used to introduce these features and play an emphasizing role.
def set_env_for_DeepRec():
    '''
    Set some ENV for these DeepRec's features enabled by ENV.
    More Detail information is shown in https://deeprec.readthedocs.io/zh/latest/index.html.
    START_STATISTIC_STEP & STOP_STATISTIC_STEP: On CPU platform, DeepRec supports memory optimization
        in both stand-alone and distributed trainging. It's default to open, and the
        default start and stop steps of collection is 1000 and 1100. Reduce the initial
        cold start time by the following settings.
    MALLOC_CONF: On CPU platform, DeepRec can use memory optimization with the jemalloc library.
        Please preload libjemalloc.so by `LD_PRELOAD=./libjemalloc.so.2 python ...`
    '''
    os.environ['START_STATISTIC_STEP'] = '100'
    os.environ['STOP_STATISTIC_STEP'] = '110'
    os.environ['MALLOC_CONF'] = \
        'background_thread:true,metadata_thp:auto,dirty_decay_ms:20000,muzzy_decay_ms:20000'
    if args.group_embedding == "collective":
        tf.config.experimental.enable_distributed_strategy(strategy="collective")
        if args.smartstaged and not args.tf:
            os.environ["TF_GPU_THREAD_COUNT"] = "16"

if __name__ == '__main__':
    parser = get_arg_parser()