Support docker_build_env, fixes #424

automation/script/docker.py

@@ -36,7 +36,6 @@ def dockerfile(self_module, input_params):
     state_data = input_params.get('state', {})
     constant_vars = input_params.get('const', {})
     constant_state = input_params.get('const_state', {})
-    dockerfile_env = input_params.get('dockerfile_env', {})
     tag_values = input_params.get('tags', '').split(",")
     variation_tags = [tag[1:] for tag in tag_values if tag.startswith("_")]
 
@@ -56,7 +55,6 @@ def dockerfile(self_module, input_params):
             'script_variation_tags': variation_tags
         }
         docker_settings = metadata.get('docker', {})
-        docker_settings['dockerfile_env'] = dockerfile_env
         state_data['docker'] = docker_settings
         add_deps_recursive = input_params.get('add_deps_recursive', {})
 
@@ -94,8 +92,6 @@ def dockerfile(self_module, input_params):
 
         # Set Docker-specific configurations
         docker_settings = state_data['docker']
-        docker_settings['dockerfile_env'] = dockerfile_env
-        dockerfile_env['MLC_RUN_STATE_DOCKER'] = True
 
         if not docker_settings.get('run', True) and not input_params.get(
                 'docker_run_override', False):
@@ -171,13 +167,18 @@ def dockerfile(self_module, input_params):
         if input_params.get('docker_push_image') in [True, 'True', 'yes']:
             env['MLC_DOCKER_PUSH_IMAGE'] = 'yes'
 
+        dockerfile_env = docker_inputs.get('env', {})
+        dockerfile_build_env = docker_inputs.get('build_env', {})
+
+        dockerfile_env['MLC_RUN_STATE_DOCKER'] = True
         # Generate Dockerfile
         mlc_docker_input = {
             'action': 'run', 'automation': 'script', 'tags': 'build,dockerfile',
             'fake_run_option': " " if docker_inputs.get('real_run') else " --fake_run",
             'comments': comments, 'run_cmd': f"{run_command_string} --quiet",
             'script_tags': input_params.get('tags'), 'env': env,
             'dockerfile_env': dockerfile_env,
+            'dockerfile_build_env': dockerfile_build_env,
             'quiet': True, 'real_run': True
         }
 

automation/script/docker_utils.py

@@ -134,7 +134,7 @@ def prepare_docker_inputs(input_params, docker_settings,
     keys = [
         "mlc_repo", "mlc_repo_branch", "base_image", "os", "os_version",
         "mlc_repos", "skip_mlc_sys_upgrade", "extra_sys_deps", "image_name",
-        "gh_token", "fake_run_deps", "run_final_cmds", "real_run", "copy_files", "path", "user"
+        "gh_token", "fake_run_deps", "run_final_cmds", "real_run", "copy_files", "path", "user", "env", "build_env"
     ]
 
     if run_stage:

script/app-mlperf-inference/meta.yaml

@@ -386,9 +386,14 @@ variations:
     env:
       MLC_NVIDIA_MITTEN_FROM_SRC: 'yes'
     docker:
+      os_version: "24.04"
+      user: 'ubuntu'
       build_deps:
         - tags: detect,os
       image_name: mlperf-inference-nvidia-v5.0-common
+      build_env:
+        ENV: release
+
     update_meta_if_env:
       - enable_if_env:
           MLC_HOST_PLATFORM_FLAVOR:
@@ -1309,7 +1314,7 @@ variations:
       mounts:
         - "${{ MLC_ML_MODEL_FILE_WITH_PATH }}:${{ MLC_ML_MODEL_FILE_WITH_PATH }}"
         - "${{ DLRM_DATA_PATH }}:${{ DLRM_DATA_PATH }}"
-      dockerfile_env:
+      env:
         MLC_ML_MODEL_FILE_WITH_PATH: "on"
 
 

script/build-dockerfile/customize.py

@@ -246,13 +246,15 @@ def preprocess(i):
 
     for key, value in config['ENV'].items():
         f.write('ENV ' + key + "=\"" + value + "\"" + EOL)
+
+    dockerfile_build_env = env.get('MLC_DOCKERFILE_BUILD_ENV', {})
+    for key in dockerfile_build_env:
+        value = dockerfile_build_env[key]
+        f.write('ENV ' + key + "=\"" + value + "\"" + EOL)
+
     for cmd in config['RUN_CMDS']:
         f.write('RUN ' + cmd + EOL)
 
-    if env.get('MLC_MLPERF_IMPLEMENTATION', '') == "nvidia" and env.get(
-            'MLC_MLPERF_INFERENCE_VERSION', '') == "5.0":
-        f.write('ENV ' + 'ENV' + "=\"" + 'release' + "\"" + EOL)
-
     f.write(EOL + '# Setup docker user' + EOL)
     docker_user = get_value(env, config, 'USER', 'MLC_DOCKER_USER')
     docker_group = get_value(env, config, 'GROUP', 'MLC_DOCKER_GROUP')
@@ -261,7 +263,8 @@ def preprocess(i):
             'MLC_DOCKER_USE_DEFAULT_USER', '') == '':
         env['MLC_DOCKER_USE_DEFAULT_USER'] = 'yes'
 
-    if docker_user and not is_true(env.get('MLC_DOCKER_USE_DEFAULT_USER', '')):
+    if docker_user and not is_true(
+            env.get('MLC_DOCKER_USE_DEFAULT_USER', '') and docker_user != 'root'):
 
         f.write('RUN groupadd -g $GID -o ' + docker_group + EOL)
 
@@ -271,15 +274,13 @@ def preprocess(i):
         user_shell = json.loads(shell)
         f.write(f"""RUN (id -u {docker_user} > /dev/null 2>&1 && usermod -u $UID {docker_user}) || useradd """ + DOCKER_USER_ID + DOCKER_GROUP + ' --create-home --shell ' + user_shell[0] + ' '
                 + docker_user + EOL)
-        f.write(f'RUN usermod -aG sudo {docker_user}' + EOL)
+        # f.write(f'RUN usermod -aG sudo {docker_user}' + EOL)
 
         f.write(
             # create the file with both lines and a trailing newline
             f"RUN printf '{docker_user} ALL=(ALL) NOPASSWD: ALL\\n"
             f"Defaults:{docker_user} !requiretty\\n' "
-            f"> /etc/sudoers.d/{docker_user} \\\n"
-            # lock down permissions
-            f"    && chmod 0440 /etc/sudoers.d/{docker_user}{EOL}"
+            f">> /etc/sudoers " + EOL
         )
 
         f.write('USER ' + docker_user + ":" + docker_group + EOL)

script/build-dockerfile/meta.yaml

@@ -37,6 +37,7 @@ input_mapping:
   docker_os: MLC_DOCKER_OS
   docker_os_version: MLC_DOCKER_OS_VERSION
   dockerfile_env: MLC_DOCKERFILE_ENV
+  dockerfile_build_env: MLC_DOCKERFILE_BUILD_ENV
   extra_sys_deps: MLC_DOCKER_EXTRA_SYS_DEPS
   fake_docker_deps: MLC_DOCKER_FAKE_DEPS
   fake_run_option: MLC_DOCKER_FAKE_RUN_OPTION