feat: add get_operator_logs to pytorch job

Author: Roja Reddy Sareddy

src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py

@@ -24,6 +24,7 @@
 PLURAL = "hyperpodpytorchjobs"
 KIND = "HyperPodPyTorchJob"
 TRAINING_OPERATOR_NAMESPACE = "aws-hyperpod"
+TRAINING_OPERATOR_POD_PREFIX = "hp-training-operator-hp-training-controller-manager-"
 
 
 class HyperPodPytorchJob(_HyperPodPytorchJob):
@@ -248,9 +249,19 @@ def get_operator_logs(cls, since_hours: float):
                 f"No pod found in namespace {TRAINING_OPERATOR_NAMESPACE}"
             )
 
-        # Get logs from first pod
-        first_pod = pods.items[0]
-        pod_name = first_pod.metadata.name
+        # Find the training operator pod
+        operator_pod = None
+        for pod in pods.items:
+            if pod.metadata.name.startswith(TRAINING_OPERATOR_POD_PREFIX):
+                operator_pod = pod
+                break
+
+        if not operator_pod:
+            raise Exception(
+                f"No training operator pod found with prefix {TRAINING_OPERATOR_POD_PREFIX}"
+            )
+
+        pod_name = operator_pod.metadata.name
 
         try:
             logs = v1.read_namespaced_pod_log(

test/unit_tests/training/test_hyperpod_pytorch_job.py

@@ -286,16 +286,21 @@ 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_pod = MagicMock()
-        mock_pod.metadata.name = "training-operator-pod"
-        mock_core_api.return_value.list_namespaced_pod.return_value.items = [mock_pod]
+        # Mock multiple pods, including the training operator pod
+        mock_other_pod = MagicMock()
+        mock_other_pod.metadata.name = "other-pod-123"
+        
+        mock_operator_pod = MagicMock()
+        mock_operator_pod.metadata.name = "hp-training-operator-hp-training-controller-manager-abc123"
+        
+        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"
 
         result = HyperPodPytorchJob.get_operator_logs(2.5)
 
         self.assertEqual(result, "training operator logs")
         mock_core_api.return_value.read_namespaced_pod_log.assert_called_once_with(
-            name="training-operator-pod",
+            name="hp-training-operator-hp-training-controller-manager-abc123",
             namespace="aws-hyperpod",
             timestamps=True,
             since_seconds=9000,