task(RHOAIENG-33283): Change to zipping working_dir

Author: kryanbeane

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -303,13 +303,18 @@ def _build_submitter_pod_template(
         """
         Build submitterPodTemplate with ConfigMap volume mount for local files.
 
+        If files contain working_dir.zip, an init container will unzip it before
+        the main submitter container runs.
+
         Args:
             files: Dict of file_name -> file_content
             configmap_name: Name of the ConfigMap containing the files
 
         Returns:
             submitterPodTemplate specification
         """
+        from codeflare_sdk.ray.rayjobs.runtime_env import UNZIP_PATH
+
         # Image has to be hard coded for the job submitter
         image = get_ray_image_for_python_version()
         if (
@@ -324,16 +329,26 @@ def _build_submitter_pod_template(
         for file_name in files.keys():
             config_map_items.append({"key": file_name, "path": file_name})
 
+        # Check if we need to unzip working_dir
+        has_working_dir_zip = "working_dir.zip" in files
+
+        # Base volume mounts for main container
+        volume_mounts = [{"name": "ray-job-files", "mountPath": MOUNT_PATH}]
+
+        # If we have a zip file, we need shared volume for unzipped content
+        if has_working_dir_zip:
+            volume_mounts.append(
+                {"name": "unzipped-working-dir", "mountPath": UNZIP_PATH}
+            )
+
         submitter_pod_template = {
             "spec": {
                 "restartPolicy": "Never",
                 "containers": [
                     {
                         "name": "ray-job-submitter",
                         "image": image,
-                        "volumeMounts": [
-                            {"name": "ray-job-files", "mountPath": MOUNT_PATH}
-                        ],
+                        "volumeMounts": volume_mounts,
                     }
                 ],
                 "volumes": [
@@ -348,6 +363,33 @@ def _build_submitter_pod_template(
             }
         }
 
+        # Add init container and volume for unzipping if needed
+        if has_working_dir_zip:
+            # Add emptyDir volume for unzipped content
+            submitter_pod_template["spec"]["volumes"].append(
+                {"name": "unzipped-working-dir", "emptyDir": {}}
+            )
+
+            # Add init container to unzip before KubeRay's submitter runs
+            submitter_pod_template["spec"]["initContainers"] = [
+                {
+                    "name": "unzip-working-dir",
+                    "image": image,
+                    "command": ["/bin/sh", "-c"],
+                    "args": [
+                        # Decode base64 zip, save to temp file, extract, cleanup
+                        f"mkdir -p {UNZIP_PATH} && "
+                        f"python3 -m base64 -d {MOUNT_PATH}/working_dir.zip > /tmp/working_dir.zip && "
+                        f"python3 -m zipfile -e /tmp/working_dir.zip {UNZIP_PATH}/ && "
+                        f"rm /tmp/working_dir.zip && "
+                        f"echo 'Successfully unzipped working_dir to {UNZIP_PATH}' && "
+                        f"ls -la {UNZIP_PATH}"
+                    ],
+                    "volumeMounts": volume_mounts,
+                }
+            ]
+            logger.info(f"Added init container to unzip working_dir to {UNZIP_PATH}")
+
         logger.info(
             f"Built submitterPodTemplate with {len(files)} files mounted at {MOUNT_PATH}, using image: {image}"
         )