bring back dshm volume for A3 Mega and A3 High (#512)

Author: gcie

src/xpk/core/workload_decorators/tcpx_decorator.py

@@ -131,6 +131,9 @@ def add_volumes(job_manifest: dict):
   })
   volumes.append({'name': 'sys', 'hostPath': {'path': '/sys'}})
   volumes.append({'name': 'proc-sys', 'hostPath': {'path': '/proc/sys'}})
+  volumes.append(
+      {'name': 'dshm', 'emptyDir': {'medium': 'Memory', 'sizeLimit': '128Gi'}}
+  )
 
 
 def add_tcpx_daemon_container(job_manifest):
@@ -177,3 +180,6 @@ def update_gpu_containers(job_manifest):
       volumeMounts.append(
           {'name': 'libraries', 'mountPath': '/usr/local/nvidia/lib64'}
       )
+      container['volumeMounts'].append(
+          {'name': 'dshm', 'mountPath': '/dev/shm'}
+      )

src/xpk/core/workload_decorators/tcpxo_decorator.py

@@ -15,6 +15,7 @@
 """
 
 import yaml
+
 from ...utils.yaml import literal_string
 
 # Component version
@@ -141,6 +142,9 @@ def add_volumes(job_manifest):
       'name': 'aperture-devices',
       'hostPath': {'path': '/dev/aperture_devices'},
   })
+  volumes.append(
+      {'name': 'dshm', 'emptyDir': {'medium': 'Memory', 'sizeLimit': '128Gi'}}
+  )
 
 
 def add_tcpxo_daemon_container(job_manifest):
@@ -189,3 +193,6 @@ def update_gpu_containers(job_manifest):
       container['volumeMounts'].append(
           {'name': 'libraries', 'mountPath': '/usr/local/nvidia'}
       )
+      container['volumeMounts'].append(
+          {'name': 'dshm', 'mountPath': '/dev/shm'}
+      )