Environment Variable Not Passed into Docker Image When Registering Model

[orriduck]

Describe the bug
I am trying to register a sagemaker model with specific environment variable and docker image, while when I am trying to deploy that, the environment variable doesn't goes in.

To reproduce

My part of DockerFile

ENV PATH=/root/.local/bin:$PATH

...

# COPY config.yml .
COPY /code .

FROM tensorflow/tensorflow:2.3.1

RUN apt-get clean && \
    apt-get update -y && \
    apt-get install -y ca-certificates nginx 

# When setting up the PATH, we need to make sure to add the folder
# that we are going to be using for code. This is the folder used by
# SageMaker to find and run our code.
ENV PATH="/opt/ml/code:/root/.local/bin:${PATH}"
RUN mkdir -p /opt/ml/code
WORKDIR /opt/ml/code

# Set some environment variables. PYTHONUNBUFFERED keeps Python from buffering our standard
# output stream, which means that logs can be delivered to the user quickly. PYTHONDONTWRITEBYTECODE
# keeps Python from writing the .pyc files which are unnecessary in this case. We also update
# PATH so that the train and serve programs are found when the container is invoked.
ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE

My serve file related snippet for this

#!/usr/bin/env python

# This file implements the scoring service shell. You don't necessarily need to modify it for various
# algorithms. It starts nginx and gunicorn with the correct configurations and then simply waits until
# gunicorn exits.
#
# The flask server is specified to be the app object in wsgi.py

import multiprocessing
import os
import signal
import subprocess
import sys
import gc

cpu_count = multiprocessing.cpu_count()

# new settings

# Parameter                Environment Variable              Default Value
# ---------                --------------------              -------------
# number of workers        MODEL_SERVER_WORKERS              min(1+(CPU core/4), 6)
# timeout                  MODEL_SERVER_TIMEOUT              240 seconds


print("[ENVIRONMENT]:", os.environ)

model_server_timeout = os.environ.get('MODEL_SERVER_TIMEOUT', 240)
model_server_workers = int(os.environ.get('MODEL_SERVER_WORKERS', min(1 + int(cpu_count / 4), 6) ))
print(f"[HOST_WORKER_SETUP] model_server_timeout: {model_server_timeout}s, model_server_workers: {model_server_workers}")

def sigterm_handler(nginx_pid, gunicorn_pid):
    try:
        os.kill(nginx_pid, signal.SIGQUIT)
        gc.collect()
    except OSError:
        pass

    try:
        os.kill(gunicorn_pid, signal.SIGTERM)
        gc.collect()
    except OSError:
        pass

    sys.exit(0)


def start_server():
    print('Starting the inference server with {} workers.'.format(model_server_workers))

    # link the log streams to stdout/err so they will be logged to the container logs
    subprocess.check_call(['ln', '-sf', '/dev/stdout', '/var/log/nginx/access.log'])
    subprocess.check_call(['ln', '-sf', '/dev/stderr', '/var/log/nginx/error.log'])

    nginx = subprocess.Popen(['nginx', '-c', '/opt/ml/code/nginx.conf'])
    gunicorn = subprocess.Popen(['gunicorn',
                                 '--timeout', str(model_server_timeout),
                                 '-k', 'gevent',
                                 '-b', 'unix:/tmp/gunicorn.sock',
                                 '-w', str(model_server_workers),
                                 'wsgi:app'])

    signal.signal(signal.SIGTERM, lambda a, b: sigterm_handler(nginx.pid, gunicorn.pid))
...

I am doing register model in this way

sagemaker_session = sagemaker.session.Session()
sm_client = boto3.client("sagemaker")

sagemaker_model = sagemaker.model.Model(
    image_uri=INFERENCE_IMAGE_URI,
    model_data=MODEL_ARTIFACT_DATA,
    role=ROLE,
    name=MODEL_NAME,
    sagemaker_session=sagemaker_session,
    env={
        "MODEL_SERVER_WORKERS": 1
    }
)

registered_model_package_desc = \
    f"Inherit from registered model package arn {INHERIT_MODEL_PACKAGE_ARN}" \
    if INHERIT_MODEL_PACKAGE_ARN is not None else None
    
fetched_model_group_flag = find_model_group(sm_client, MODEL_PACKAGE_GROUP).__len__() > 0
if fetched_model_group_flag is False:
    response = sm_client.create_model_package_group(
        ModelPackageGroupName=MODEL_PACKAGE_GROUP
)

registered_model_package = sagemaker_model.register(
    content_types=REQUEST_CONTENT_TYPES,
    response_types=RESPONSE_CONTENT_TYPES,
    inference_instances=INFERENCE_INSTANCES,
    transform_instances=TRANSFORM_INSTANCES,
    model_package_group_name=MODEL_PACKAGE_GROUP,
    approval_status=MODEL_APPROVAL_STATE,
    description=registered_model_package_desc
)
print("Registered Model Package:", registered_model_package.model_package_arn)

Expected behavior
There should be Environment variable shows up overwrite the MODEL_SERVER_WORKERS

Screenshots or logs
If applicable, add screenshots or logs to help explain your problem.

System information
A description of your system. Please provide:

• SageMaker Python SDK version: 2.24 (SageMaker Studio)
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): (SageMaker Studio)
• Framework version: (SageMaker Studio)
• Python version: (SageMaker Studio - Data Science)
• CPU or GPU: CPU
• Custom Docker image (Y/N): Y

Additional context
Add any other context about the problem here.

[akrishna1995]

Hi @orriduck - I see there has been a fix marked in #2405 linked to this issue. Can you help confirm if this issue still persists? Closing this issue for now, feel free to reopen if the issue still persists with the latest sagemaker.