feat: add get_operator_logs to pytorch job (#218)

* feat: add get_operator_logs to pytorch job

* feat: add get_operator_logs to pytorch job

* feat: add get_operator_logs to pytorch job

* feat: add get_operator_logs to pytorch job

---------

Co-authored-by: Roja Reddy Sareddy <[email protected]>

Author: rsareddy0329

src/sagemaker/hyperpod/cli/commands/training.py

@@ -354,3 +354,21 @@ def pytorch_get_logs(job_name: str, pod_name: str, namespace: str):
 
     except Exception as e:
         raise click.UsageError(f"Failed to list jobs: {str(e)}")
+
+
+@click.command("hyp-pytorch-job")
+@click.option(
+    "--since-hours",
+    type=click.FLOAT,
+    required=True,
+    help="Required. The time frame to get logs for.",
+)
+@_hyperpod_telemetry_emitter(Feature.HYPERPOD_CLI, "get_pytorch_operator_logs")
+def pytorch_get_operator_logs(
+    since_hours: float,
+):
+    """
+    Get operator logs for pytorch training jobs.
+    """
+    logs = HyperPodPytorchJob.get_operator_logs(since_hours=since_hours)
+    click.echo(logs)

src/sagemaker/hyperpod/cli/hyp_cli.py

@@ -16,6 +16,7 @@
     pytorch_delete,
     pytorch_list_pods,
     pytorch_get_logs,
+    pytorch_get_operator_logs,
 )
 from sagemaker.hyperpod.cli.commands.inference import (
     js_create,
@@ -139,6 +140,7 @@ def get_operator_logs():
 get_logs.add_command(js_get_logs)
 get_logs.add_command(custom_get_logs)
 
+get_operator_logs.add_command(pytorch_get_operator_logs)
 get_operator_logs.add_command(js_get_operator_logs)
 get_operator_logs.add_command(custom_get_operator_logs)
 

src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py

@@ -23,6 +23,8 @@
 API_VERSION = "v1"
 PLURAL = "hyperpodpytorchjobs"
 KIND = "HyperPodPyTorchJob"
+TRAINING_OPERATOR_NAMESPACE = "aws-hyperpod"
+TRAINING_OPERATOR_LABEL = "hp-training-control-plane"
 
 
 class HyperPodPytorchJob(_HyperPodPytorchJob):
@@ -233,6 +235,40 @@ def get_logs_from_pod(self, pod_name: str, container: Optional[str] = None) -> s
             logger.error(f"Failed to get logs from pod {pod_name}!")
             handle_exception(e, self.metadata.name, self.metadata.namespace)
 
+    @classmethod
+    @_hyperpod_telemetry_emitter(Feature.HYPERPOD, "get_operator_logs_pytorchjob")
+    def get_operator_logs(cls, since_hours: float):
+        cls.verify_kube_config()
+
+        v1 = client.CoreV1Api()
+
+        # Get pods with the training operator label directly
+        pods = v1.list_namespaced_pod(
+            namespace=TRAINING_OPERATOR_NAMESPACE,
+            label_selector=TRAINING_OPERATOR_LABEL
+        )
+
+        if not pods.items:
+            raise Exception(
+                f"No training operator pod found with label {TRAINING_OPERATOR_LABEL}"
+            )
+
+        # Use the first pod found
+        operator_pod = pods.items[0]
+        pod_name = operator_pod.metadata.name
+
+        try:
+            logs = v1.read_namespaced_pod_log(
+                name=pod_name,
+                namespace=TRAINING_OPERATOR_NAMESPACE,
+                timestamps=True,
+                since_seconds=int(3600 * since_hours),
+            )
+        except Exception as e:
+            handle_exception(e, pod_name, TRAINING_OPERATOR_NAMESPACE)
+
+        return logs
+
 
 def _load_hp_job(response: dict) -> HyperPodPytorchJob:
 

test/integration_tests/training/cli/test_cli_training.py

@@ -239,4 +239,9 @@ def test_delete_job(self, test_job_name):
         assert list_result.returncode == 0
 
         # The job name should no longer be in the output
-        assert test_job_name not in list_result.stdout
+        assert test_job_name not in list_result.stdout
+
+def test_pytorch_get_operator_logs():
+    """Test getting operator logs via CLI"""
+    result = execute_command(["hyp", "get-operator-logs", "hyp-pytorch-job", "--since-hours", "1"])
+    assert result.returncode == 0

test/integration_tests/training/sdk/test_sdk_training.py

@@ -112,3 +112,8 @@ def test_delete_job(self, pytorch_job):
         jobs = HyperPodPytorchJob.list()
         job_names = [job.metadata.name for job in jobs]
         assert pytorch_job.metadata.name not in job_names
+
+def test_get_operator_logs():
+    """Test getting operator logs"""
+    logs = HyperPodPytorchJob.get_operator_logs(since_hours=1)
+    assert logs

test/unit_tests/cli/test_training.py

@@ -6,6 +6,7 @@
     pytorch_create,
     list_jobs,
     pytorch_describe,
+    pytorch_get_operator_logs,
 )
 from hyperpod_pytorch_job_template.v1_1.model import ALLOWED_TOPOLOGY_LABELS
 import sys
@@ -827,3 +828,12 @@ def test_none_topology_labels(self):
         )
         self.assertIsNone(config.preferred_topology)
         self.assertIsNone(config.required_topology)
+
+@patch('sagemaker.hyperpod.cli.commands.training.HyperPodPytorchJob')
+def test_pytorch_get_operator_logs(mock_hp):
+    mock_hp.get_operator_logs.return_value = "operator logs"
+    runner = CliRunner()
+    result = runner.invoke(pytorch_get_operator_logs, ['--since-hours', '2'])
+    assert result.exit_code == 0
+    assert 'operator logs' in result.output
+    mock_hp.get_operator_logs.assert_called_once_with(since_hours=2.0)

test/unit_tests/training/test_hyperpod_pytorch_job.py

@@ -283,6 +283,31 @@ def test_get_logs_from_pod_with_container_name(
         )
         self.assertEqual(result, "test logs")
 
+    @patch("kubernetes.client.CoreV1Api")
+    @patch.object(HyperPodPytorchJob, "verify_kube_config")
+    def test_get_operator_logs(self, mock_verify_config, mock_core_api):
+        # Mock only the training operator pod (since we're using label selector)
+        mock_operator_pod = MagicMock()
+        mock_operator_pod.metadata.name = "training-operator-pod-abc123"
+        
+        mock_core_api.return_value.list_namespaced_pod.return_value.items = [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")
+        # Verify label selector is used
+        mock_core_api.return_value.list_namespaced_pod.assert_called_once_with(
+            namespace="aws-hyperpod",
+            label_selector="hp-training-control-plane"
+        )
+        mock_core_api.return_value.read_namespaced_pod_log.assert_called_once_with(
+            name="training-operator-pod-abc123",
+            namespace="aws-hyperpod",
+            timestamps=True,
+            since_seconds=9000,
+        )
+
 
 class TestLoadHpJob(unittest.TestCase):
     """Test the _load_hp_job function"""
@@ -350,4 +375,4 @@ def test_load_hp_job_list_empty(self):
         result = _load_hp_job_list(response)
 
         self.assertEqual(len(result), 0)
-        self.assertEqual(result, [])
+        self.assertEqual(result, [])