Support docker_base_image

ctuning · Jan 30, 2024 · c922cb9 · c922cb9
1 parent 927908d
commit c922cb9
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 7 deletions.
diff --git a/cm-mlops/automation/script/module_misc.py b/cm-mlops/automation/script/module_misc.py
@@ -1391,9 +1391,15 @@ def dockerfile(i):
         run_cmd  = r['run_cmd_string']
 
 
-
+        docker_base_image = i.get('docker_base_image', docker_settings.get('base_image'))
         docker_os = i.get('docker_os', docker_settings.get('docker_os', 'ubuntu'))
         docker_os_version = i.get('docker_os_version', docker_settings.get('docker_os_version', '22.04'))
+        if not docker_base_image:
+            dockerfilename_suffix = docker_os +'_'+docker_os_version
+        else:
+            dockerfilename_suffix = docker_base_image.split("/")
+            dockerfilename_suffix = dockerfilename_suffix[len(dockerfilename_suffix) - 1]
+
         fake_run_deps = i.get('fake_run_deps', docker_settings.get('fake_run_deps', False))
         docker_run_final_cmds = docker_settings.get('docker_run_final_cmds', [])
 
@@ -1409,7 +1415,7 @@ def dockerfile(i):
 
         env['CM_DOCKER_PRE_RUN_COMMANDS'] = docker_run_final_cmds
 
-        dockerfile_path = os.path.join(script_path,'dockerfiles', docker_os +'_'+docker_os_version +'.Dockerfile')
+        dockerfile_path = os.path.join(script_path,'dockerfiles', dockerfilename_suffix +'.Dockerfile')
         if i.get('print_deps'):
             cm_input = {'action': 'run',
                     'automation': 'script',
@@ -1434,6 +1440,7 @@ def dockerfile(i):
                             'automation': 'script',
                             'tags': 'build,dockerfile',
                             'cm_repo': cm_repo,
+                            'docker_base_image': docker_base_image,
                             'docker_os': docker_os,
                             'docker_os_version': docker_os_version,
                             'file_path': dockerfile_path,
@@ -1659,9 +1666,18 @@ def docker(i):
 
         mount_string = "" if len(mounts)==0 else ",".join(mounts)
 
+        docker_base_image = i.get('docker_base_image', docker_settings.get('base_image'))
+        docker_os = i.get('docker_os', docker_settings.get('docker_os', 'ubuntu'))
+        docker_os_version = i.get('docker_os_version', docker_settings.get('docker_os_version', '22.04'))
+        if not docker_base_image:
+            dockerfilename_suffix = docker_os +'_'+docker_os_version
+        else:
+            dockerfilename_suffix = docker_base_image.split("/")
+            dockerfilename_suffix = dockerfilename_suffix[len(dockerfilename_suffix) - 1]
+
         cm_repo=i.get('docker_cm_repo', 'mlcommons@ck')
 
-        dockerfile_path = os.path.join(script_path,'dockerfiles', _os +'_'+version +'.Dockerfile')
+        dockerfile_path = os.path.join(script_path,'dockerfiles', dockerfilename_suffix +'.Dockerfile')
 
         docker_skip_run_cmd = i.get('docker_skip_run_cmd', docker_settings.get('skip_run_cmd', False)) #skips docker run cmd and gives an interactive shell to the user
 
@@ -1712,15 +1728,16 @@ def docker(i):
                            'automation': 'script',
                            'tags': 'run,docker,container',
                            'recreate': 'yes',
+                           'docker_base_image': docker_base_image,
                            'docker_os': _os,
+                           'docker_os_version': version,
                            'cm_repo': cm_repo,
                            'env': env,
                            'image_repo': image_repo,
                            'interactive': interactive,
                            'mounts': mounts,
                            'image_name': 'cm-script-'+script_alias,
 #                            'image_tag': script_alias,
-                           'docker_os_version': version,
                            'detached': detached,
                            'script_tags': f'{tag_string}',
                            'run_cmd': run_cmd if docker_skip_run_cmd not in [ 'yes', True, 'True' ] else 'echo "cm version"',

diff --git a/cm-mlops/script/build-dockerfile/_cm.json b/cm-mlops/script/build-dockerfile/_cm.json
@@ -10,6 +10,7 @@
     "cm_repo": "CM_MLOPS_REPO",
     "docker_os": "CM_DOCKER_OS",
     "docker_os_version": "CM_DOCKER_OS_VERSION",
+    "docker_base_image": "CM_DOCKER_IMAGE_BASE",
     "fake_run_option": "CM_DOCKER_FAKE_RUN_OPTION",
     "file_path": "CM_DOCKERFILE_WITH_PATH",
     "gh_token": "CM_GH_TOKEN",

diff --git a/cm-mlops/script/build-mlperf-inference-server-nvidia/_cm.yaml b/cm-mlops/script/build-mlperf-inference-server-nvidia/_cm.yaml
@@ -197,6 +197,9 @@ versions:
 docker:
   skip_run_cmd: 'no'
   all_gpus: 'yes'
+  docker_os: ubuntu
+  docker_os_version: '20.04'
+  base_image: nvcr.io/nvidia/mlperf/mlperf-inference:mlpinf-v3.1-cuda12.2-cudnn8.9-x86_64-ubuntu20.04-public
   docker_input_mapping:
     imagenet_path: IMAGENET_PATH
     results_dir: RESULTS_DIR
@@ -213,6 +216,3 @@ docker:
    - "${{ CM_TENSORRT_TAR_FILE_PATH }}:${{ CM_TENSORRT_TAR_FILE_PATH }}"
    - "${{ CUDA_RUN_FILE_LOCAL_PATH }}:${{ CUDA_RUN_FILE_LOCAL_PATH }}"
    - "${{ MLPERF_SCRATCH_PATH }}:${{ MLPERF_SCRATCH_PATH }}"
-  pre_run_cmds:
-    - cm pull repo mlcommons@ck
-    - cm run script --tags=get,dataset,original,imagenet,_full --imagenet_path=/data/imagenet-val
diff --git a/cm-mlops/script/run-docker-container/_cm.json b/cm-mlops/script/run-docker-container/_cm.json
@@ -22,6 +22,7 @@
     "image_tag": "CM_DOCKER_IMAGE_TAG",
     "docker_os": "CM_DOCKER_OS",
     "docker_os_version": "CM_DOCKER_OS_VERSION",
+    "docker_image_base": "CM_DOCKER_IMAGE_BASE",
     "script_tags": "CM_DOCKER_RUN_SCRIPT_TAGS",
     "run_cmd_extra": "CM_DOCKER_RUN_CMD_EXTRA",
     "real_run": "CM_REAL_RUN",