Add docker_split_mlc_run_cmd to split Dockerfile MLC run flow into cacheable layers (mlcommons#924)

Author: Copilot

automation/script/docker.py

@@ -25,10 +25,10 @@ def dockerfile(self_module, input_params):
     # Step 2: Process Dockerfile-related configurations
     env = input_params.get('env', {})
 
-    host_only_env_keys = [ "HOME", "USER" ]
+    host_only_env_keys = ["HOME", "USER"]
     for key in host_only_env_keys:
         if key in env:
-            del(env[key])
+            del (env[key])
 
     state_data = input_params.get('state', {})
     constant_vars = input_params.get('const', {})
@@ -226,10 +226,10 @@ def docker_run(self_module, i):
     show_time = i.get('show_time', False)
     logger = self_module.logger
     env = i.get('env', {})
-    host_only_env_keys = [ "HOME", "USER" ]
+    host_only_env_keys = ["HOME", "USER"]
     for key in host_only_env_keys:
         if key in env:
-            del(env[key])
+            del (env[key])
 
     self_module.env = env
 

automation/script/docker_utils.py

@@ -142,6 +142,7 @@ def prepare_docker_inputs(input_params, docker_settings,
         "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", "env", "build_env"
     ]
+    keys += ["cache", "split_mlc_run_cmd"]
 
     if run_stage:
         keys += [

script/build-dockerfile/README.md

@@ -74,6 +74,7 @@ mlcr build,dockerfile
 | `--run_cmd` |  |  | `` |
 | `--run_cmd_extra` |  |  | `` |
 | `--script_tags` |  |  | `` |
+| `--split_mlc_run_cmd` | Split dependent MLC run commands into separate Dockerfile RUN layers (effective with cache enabled) |  | `` |
 | `--skip_mlc_sys_upgrade` |  |  | `` |
 | `--user` |  |  | `` |
 ### Generic Script Inputs

script/build-dockerfile/customize.py

@@ -427,10 +427,17 @@ def preprocess(i):
                 env['MLC_DOCKER_RUN_CMD']
 
     logger.info(env['MLC_DOCKER_RUN_CMD'])
+    split_mlc_run_cmd = is_true(env.get('MLC_DOCKER_SPLIT_MLC_RUN_CMD', ''))
+    if split_mlc_run_cmd and is_false(env.get('MLC_DOCKER_CACHE', "True")):
+        logger.warning(
+            "Ignoring docker_split_mlc_run_cmd since docker cache is disabled.")
+        split_mlc_run_cmd = False
+
     fake_run = env.get("MLC_DOCKER_FAKE_RUN_OPTION",
                        " --fake_run") + dockerfile_env_input_string
+    enable_fake_deps = env.get('MLC_DOCKER_FAKE_DEPS')
     fake_run = fake_run + \
-        " --fake_deps" if env.get('MLC_DOCKER_FAKE_DEPS') else fake_run
+        " --fake_deps" if enable_fake_deps else fake_run
 
     x = 'RUN ' + _run_secret_prefix + env['MLC_DOCKER_RUN_CMD']
 
@@ -441,8 +448,9 @@ def preprocess(i):
         if run_cmd_extra != '':
             x += ' ' + run_cmd_extra
 
-    if env.get('MLC_DOCKER_RUN_SCRIPT_TAGS', '') != '' and is_true(env.get(
-            'MLC_DOCKER_ADD_DEPENDENT_SCRIPTS_RUN_COMMANDS', '')):
+    if env.get('MLC_DOCKER_RUN_SCRIPT_TAGS', '') != '' and (
+            split_mlc_run_cmd or is_true(
+                env.get('MLC_DOCKER_ADD_DEPENDENT_SCRIPTS_RUN_COMMANDS', ''))):
         mlc_input = {'action': 'run',
                      'automation': 'script',
                      'tags': f"""{env['MLC_DOCKER_RUN_SCRIPT_TAGS']}""",
@@ -452,11 +460,11 @@ def preprocess(i):
                      'fake_run': True,
                      'fake_deps': True
                      }
-        r = self_module.mlc.access(mlc_input)
+        r = automation.action_object.access(mlc_input)
         if r['return'] > 0:
             return r
         print_deps = r['new_state']['print_deps']
-        fake_run_str = " --fake_run" if env.get('MLC_DOCKER_FAKE_DEPS') else ""
+        fake_run_str = " --fake_run" if enable_fake_deps else ""
         cmds = ["RUN " + dep for dep in print_deps]
         for cmd in cmds:
             f.write(cmd + fake_run_str + EOL)

script/build-dockerfile/meta.yaml

@@ -60,6 +60,7 @@ input_mapping:
   run_cmd: MLC_DOCKER_RUN_CMD
   run_cmd_extra: MLC_DOCKER_RUN_CMD_EXTRA
   script_tags: MLC_DOCKER_RUN_SCRIPT_TAGS
+  split_mlc_run_cmd: MLC_DOCKER_SPLIT_MLC_RUN_CMD
   skip_mlc_sys_upgrade: MLC_DOCKER_SKIP_MLC_SYS_UPGRADE
   user: MLC_DOCKER_USER