Configure output-uploader as a sidecar container (#767)

apple · Oct 21, 2024 · ac63eef · ac63eef
1 parent b3f1d3a
commit ac63eef
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/axlearn/cloud/gcp/job.py b/axlearn/cloud/gcp/job.py
@@ -572,6 +572,9 @@ def _build_uploader_container(self) -> Nested[Any]:
         return dict(
             name="output-uploader",
             image="google/cloud-sdk:alpine",
+            # https://kubernetes.io/docs/concepts/workloads/pods/sidecar-containers/#pod-sidecar-containers
+            # SideCar container is an init container with restartPolicy as "Always".
+            restartPolicy="Always",
             command=["/bin/sh", "-c"],
             args=[sync_command],
             resources=resources,
@@ -726,7 +729,8 @@ def _build_pod(self) -> Nested[Any]:
                     **selector,
                 },
                 tolerations=tolerations,
-                containers=[self._build_container(), self._build_uploader_container()],
+                containers=[self._build_container()],
+                initContainers=[self._build_uploader_container()],
                 serviceAccountName=cfg.service_account,
                 volumes=volumes,
             ),

diff --git a/axlearn/cloud/gcp/job_test.py b/axlearn/cloud/gcp/job_test.py
@@ -398,7 +398,7 @@ def test_build_pod(
                 )
                 self.assertEqual("spot", labels.get("bastion-tier", None))
 
-            self.assertEqual(len(pod_spec["containers"]), 2)
+            self.assertEqual(len(pod_spec["containers"]), 1)
 
             # Verify worker container specs
             container = pod_spec["containers"][0]
@@ -459,9 +459,12 @@ def test_build_pod(
             )
 
             # Verify uploader container specs
-            uploader_container = pod_spec["containers"][1]
+            self.assertEqual(len(pod_spec["initContainers"]), 1)
+
+            uploader_container = pod_spec["initContainers"][0]
             self.assertEqual(uploader_container["name"], "output-uploader")
             self.assertEqual(uploader_container["image"], "google/cloud-sdk:alpine")
+            self.assertEqual(uploader_container["restartPolicy"], "Always")
             self.assertIn("volumeMounts", uploader_container)
 
             volume_mounts = uploader_container["volumeMounts"]