def _set_mms_configs()

in src/sagemaker_sklearn_container/serving_mms.py [0:0]

• 30 lines of code
• 7 McCabe index (conditional complexity)

def _set_mms_configs(is_multi_model, handler):
    """Set environment variables for MMS to parse during server initialization. These env vars are used to
    propagate the config.properties.tmp file used during MxNet Model Server initialization.
    'SAGEMAKER_MMS_MODEL_STORE' has to be set to the model location during single model inference because MMS
    is initialized with the model. In multi-model mode, MMS is started with no models loaded.
    Note: Ideally, instead of relying on env vars, this should be written directly to a config file.
    """
    max_content_length = os.getenv("MAX_CONTENT_LENGTH", DEFAULT_MAX_CONTENT_LEN)
    if int(max_content_length) > MAX_CONTENT_LEN_LIMIT:
        # Cap at 20mb
        max_content_length = MAX_CONTENT_LEN_LIMIT

    max_workers = multiprocessing.cpu_count()
    max_job_queue_size = 2 * max_workers

    # Max heap size = (max workers + max job queue size) * max payload size * 1.2 (20% buffer) + 128 (base amount)
    max_heap_size = ceil((max_workers + max_job_queue_size) * (int(max_content_length) / 1024 ** 2) * 1.2) + 128

    os.environ["SAGEMAKER_MMS_MODEL_STORE"] = '/'
    os.environ["SAGEMAKER_MMS_LOAD_MODELS"] = ''
    os.environ["SAGEMAKER_MMS_DEFAULT_HANDLER"] = handler

    # Users can define port
    _set_default_if_not_exist("SAGEMAKER_BIND_TO_PORT", str(PORT))

    # Multi Model Server configs, exposed to users as env vars
    _set_default_if_not_exist("SAGEMAKER_NUM_MODEL_WORKERS", MMS_NUM_MODEL_WORKERS_INIT)
    _set_default_if_not_exist("SAGEMAKER_MODEL_JOB_QUEUE_SIZE", MMS_MODEL_JOB_QUEUE_SIZE_DEFAULT)
    _set_default_if_not_exist("SAGEMAKER_MAX_REQUEST_SIZE", max_content_length)

    # JVM configurations for MMS, exposed to users as env vars
    _set_default_if_not_exist("SAGEMAKER_MAX_HEAP_SIZE", str(max_heap_size) + 'm')
    _set_default_if_not_exist("SAGEMAKER_MAX_DIRECT_MEMORY_SIZE", os.environ["SAGEMAKER_MAX_HEAP_SIZE"])

    MMS_CONFIG_FILE_PATH = get_mms_config_file_path()

    # TODO: Revert config.properties.tmp to config.properties and add back in vmargs
    # set with environment variables after MMS implements parsing environment variables
    # for vmargs, update MMS section of final/Dockerfile.cpu to match, and remove the
    # following code.
    try:
        with open(MMS_CONFIG_FILE_PATH + '.tmp', 'r') as f:
            with open(MMS_CONFIG_FILE_PATH, 'w+') as g:
                g.write("vmargs=-XX:-UseLargePages"
                        + " -XX:+UseParNewGC"
                        + " -XX:MaxMetaspaceSize=32M"
                        + " -XX:InitiatingHeapOccupancyPercent=25"
                        + " -Xms" + os.environ["SAGEMAKER_MAX_HEAP_SIZE"]
                        + " -Xmx" + os.environ["SAGEMAKER_MAX_HEAP_SIZE"]
                        + " -XX:MaxDirectMemorySize=" + os.environ["SAGEMAKER_MAX_DIRECT_MEMORY_SIZE"] + "\n")
                g.write(f.read())
    except Exception:
        pass