feat: add get_operator_logs to pytorch job

Author: Roja Reddy Sareddy

src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py

@@ -24,7 +24,7 @@
 PLURAL = "hyperpodpytorchjobs"
 KIND = "HyperPodPyTorchJob"
 TRAINING_OPERATOR_NAMESPACE = "aws-hyperpod"
-TRAINING_OPERATOR_POD_PREFIX = "hp-training-operator-hp-training-controller-manager-"
+TRAINING_OPERATOR_LABEL = "hp-training-control-plane"
 
 
 class HyperPodPytorchJob(_HyperPodPytorchJob):
@@ -249,16 +249,16 @@ def get_operator_logs(cls, since_hours: float):
                 f"No pod found in namespace {TRAINING_OPERATOR_NAMESPACE}"
             )
 
-        # Find the training operator pod
+        # Find the training operator pod by label
         operator_pod = None
         for pod in pods.items:
-            if pod.metadata.name.startswith(TRAINING_OPERATOR_POD_PREFIX):
+            if pod.metadata.labels and TRAINING_OPERATOR_LABEL in pod.metadata.labels:
                 operator_pod = pod
                 break
 
         if not operator_pod:
             raise Exception(
-                f"No training operator pod found with prefix {TRAINING_OPERATOR_POD_PREFIX}"
+                f"No training operator pod found with label {TRAINING_OPERATOR_LABEL}"
             )
 
         pod_name = operator_pod.metadata.name

test/unit_tests/training/test_hyperpod_pytorch_job.py

@@ -286,12 +286,14 @@ def test_get_logs_from_pod_with_container_name(
     @patch("kubernetes.client.CoreV1Api")
     @patch.object(HyperPodPytorchJob, "verify_kube_config")
     def test_get_operator_logs(self, mock_verify_config, mock_core_api):
-        # Mock multiple pods, including the training operator pod
+        # Mock multiple pods, including the training operator pod with correct label
         mock_other_pod = MagicMock()
         mock_other_pod.metadata.name = "other-pod-123"
+        mock_other_pod.metadata.labels = {"app": "other"}
 
         mock_operator_pod = MagicMock()
-        mock_operator_pod.metadata.name = "hp-training-operator-hp-training-controller-manager-abc123"
+        mock_operator_pod.metadata.name = "training-operator-pod-abc123"
+        mock_operator_pod.metadata.labels = {"hp-training-control-plane": "true"}
 
         mock_core_api.return_value.list_namespaced_pod.return_value.items = [mock_other_pod, mock_operator_pod]
         mock_core_api.return_value.read_namespaced_pod_log.return_value = "training operator logs"
@@ -300,7 +302,7 @@ def test_get_operator_logs(self, mock_verify_config, mock_core_api):
 
         self.assertEqual(result, "training operator logs")
         mock_core_api.return_value.read_namespaced_pod_log.assert_called_once_with(
-            name="hp-training-operator-hp-training-controller-manager-abc123",
+            name="training-operator-pod-abc123",
             namespace="aws-hyperpod",
             timestamps=True,
             since_seconds=9000,