Removing resource limits for the controller container on the head pod

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -135,14 +135,10 @@ class Config(BaseReplicatedJob.Config):
 
         Attributes:
             inner: The wrapped TPUReplicatedJob configuration.
-            pathways_head_cpu: CPU request for pathways-head container.
-            pathways_head_mem: Memory request for pathways-head container.
         """
 
         inner: Required[TPUReplicatedJob.Config] = REQUIRED
         pathways_xla_flags: list[str] = []
-        pathways_head_cpu: Optional[str] = None
-        pathways_head_mem: Optional[str] = None
 
     @classmethod
     def define_flags(cls, fv):
@@ -159,24 +155,6 @@ def define_flags(cls, fv):
             "Example: 'xla_tpu_x=24,megascale_y=true'",
             **common_kwargs,
         )
-        flags.DEFINE_string(
-            "pathways_head_cpu",
-            None,
-            "CPU request for pathways-head container in cores. Default is 1 core.",
-            **common_kwargs,
-        )
-        flags.DEFINE_string(
-            "pathways_head_mem",
-            None,
-            "Memory request for pathways-head container in GiB. Default is 16GiB",
-            **common_kwargs,
-        )
-
-    @classmethod
-    def set_defaults(cls, fv):
-        super().set_defaults(fv)
-        fv.set_default("pathways_head_cpu", fv.pathways_head_cpu or "32")
-        fv.set_default("pathways_head_mem", fv.pathways_head_mem or "128")
 
     @classmethod
     def default_config(cls):
@@ -287,14 +265,10 @@ def _build_pathways_head_container(self) -> dict:
 
         head_container["env"] = env_list
 
-        cpu_req = f"{float(self.config.pathways_head_cpu) * 1000}m"
-        mem_req = f"{self.config.pathways_head_mem}Gi"
-        resources = {
-            "requests": {"cpu": cpu_req, "memory": mem_req},
-            "limits": {"cpu": cpu_req, "memory": mem_req},
+        head_container["resources"] = {
+            "requests": {},
+            "limits": {},
         }
-        head_container["resources"] = resources
-
         return head_container
 
     def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]: