add nccl volume to the base podtemplate (flyteorg#8)

Signed-off-by: byhsu <[email protected]>
ByronHsu · May 18, 2023 · f523ab3 · f523ab3
1 parent 443b62d
commit f523ab3
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/flytekit/lnkd/pod_template.py b/flytekit/lnkd/pod_template.py
@@ -8,7 +8,7 @@
 JOB_USER_ID = 1001
 JOB_GROUP_ID = 1001
 NFS_USER_ID = 1001 # 2147483006
-NFS_INIT_IMAGE = "container-image-registry.corp.linkedin.com:8083/temp/ghcr.io/flyteorg/li-flyte/li-flyte-base-image:0.0.31.2"
+NFS_INIT_IMAGE = "container-image-registry.corp.linkedin.com:8083/lps-image/linkedin/li-flyte/li-flyte-base-image:0.0.29"
 NFS_WORKING_DIR = "/home/jobuser" # "/home/nfsuser"
 @dataclass
 class LIPodTemplate(PodTemplate):
@@ -140,14 +140,24 @@ def _build_base_pod_spec(self) -> V1PodSpec:
         pod_spec=V1PodSpec(
             containers=[V1Container(
                 name=PRIMARY_CONTAINER_DEFAULT_NAME,
+                volume_mounts=[V1VolumeMount(
+                    name="nccl-shared-memory",
+                    mount_path="/dev/shm",
+                )],
                 security_context=V1SecurityContext(
                     run_as_user=JOB_USER_ID,
                     run_as_group=JOB_GROUP_ID
-                )
+                ),
             )],
             security_context=V1PodSecurityContext(
                 run_as_user=JOB_USER_ID,
                 run_as_group=JOB_GROUP_ID
-            )
+            ),
+            volumes=[V1Volume(
+                name="nccl-shared-memory",
+                empty_dir=V1EmptyDirVolumeSource(
+                    medium="Memory",
+                )
+            )],
         )
         return pod_spec