Added debugpy for remote debugging experiment

Author: lukebaumann

Dockerfile

@@ -100,6 +100,8 @@ FROM base AS pathways-tpu
 
 ARG EXTRAS=
 
+RUN uv pip install debugpy && uv cache clean
+EXPOSE 5678
 RUN uv pip install --prerelease=allow .[core,pathways-tpu] && uv cache clean
 RUN if [ -n "$EXTRAS" ]; then uv pip install .[$EXTRAS] && uv cache clean; fi
 COPY . .

axlearn/cloud/gcp/jobset_utils.py

@@ -471,6 +471,7 @@ def _build_container(self) -> Nested[Any]:
                 dict(containerPort=8080),  # Port for MXLA coordinator.
                 dict(containerPort=8431),  # Port to export TPU runtime metrics.
                 dict(containerPort=self._load_balancer.target_port),  # Port for load balancer.
+                dict(containerPort=5678),  # Port for debugger.
             ],
             securityContext=dict(privileged=True),
             # TODO(markblee): Improve SIGTERM behavior for command.