Support extra_docker_run_args

ctuning · Feb 5, 2024 · bf6ad7e · bf6ad7e
1 parent 045e188
commit bf6ad7e
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 0 deletions.
diff --git a/cm-mlops/automation/script/module_misc.py b/cm-mlops/automation/script/module_misc.py
@@ -1725,6 +1725,8 @@ def docker(i):
 
         shm_size = i.get('docker_shm_size', docker_settings.get('shm_size', ''))
 
+        extra_run_args = i.get('docker_extra_run_args', docker_settings.get('extra_run_args', ''))
+
         if detached == '':
             detached = docker_settings.get('detached', '')
 
@@ -1808,6 +1810,9 @@ def docker(i):
         if shm_size != '':
             cm_docker_input['shm_size'] = shm_size
 
+        if extra_run_args != '':
+            cm_docker_input['extra_run_args'] = extra_run_args
+
         print ('')
 
 

diff --git a/cm-mlops/script/build-mlperf-inference-server-nvidia/_cm.yaml b/cm-mlops/script/build-mlperf-inference-server-nvidia/_cm.yaml
@@ -214,6 +214,7 @@ docker:
   skip_run_cmd: 'no'
   all_gpus: 'yes'
   shm_size: '32gb'
+  extra_run_args: ' --ulimit memlock=-1'
   docker_os: ubuntu
   docker_real_run: False
   interactive: True

diff --git a/cm-mlops/script/run-docker-container/_cm.json b/cm-mlops/script/run-docker-container/_cm.json
@@ -33,6 +33,7 @@
     "mounts": "CM_DOCKER_VOLUME_MOUNTS",
     "port_maps": "CM_DOCKER_PORT_MAPS",
     "shm_size": "CM_DOCKER_SHM_SIZE",
+    "extra_run_args": "CM_DOCKER_EXTRA_RUN_ARGS",
     "device": "CM_DOCKER_ADD_DEVICE",
     "cache": "CM_DOCKER_CACHE",
     "all_gpus": "CM_DOCKER_ADD_ALL_GPUS"

diff --git a/cm-mlops/script/run-docker-container/customize.py b/cm-mlops/script/run-docker-container/customize.py
@@ -99,6 +99,9 @@ def postprocess(i):
     if env.get('CM_DOCKER_SHM_SIZE', '') != '':
         run_opts += " --shm-size={}".format(env['CM_DOCKER_SHM_SIZE'])
 
+    if env.get('CM_DOCKER_EXTRA_RUN_ARGS', '') != '':
+        run_opts += env['CM_DOCKER_EXTRA_RUN_ARGS']
+
     if env.get('CM_DOCKER_PORT_MAPS', []):
         for ports in env['CM_DOCKER_PORT_MAPS']:
             port_map_cmds.append(ports)