Feature: Add support for listing Batch jobs by share identifier (#5585)

Co-authored-by: houtampl <houtampl@amazon.com>

Author: ampleh

sagemaker-train/src/sagemaker/train/aws_batch/training_queue.py

@@ -165,7 +165,38 @@ def list_jobs(
             for job_result in job_result_dict.get("jobSummaryList", []):
                 if "jobArn" in job_result and "jobName" in job_result:
                     jobs_to_return.append(
-                        TrainingQueuedJob(job_result["jobArn"], job_result["jobName"])
+                        TrainingQueuedJob(job_result["jobArn"], job_result["jobName"], job_result.get("shareIdentifier", None))
+                    )
+                else:
+                    logging.warning("Missing JobArn or JobName in Batch ListJobs API")
+                    continue
+        return jobs_to_return
+
+    def list_jobs_by_share(
+        self,
+        status: Optional[str] = JOB_STATUS_RUNNING,
+        share_identifier: Optional[str] = None,
+    ) -> List[TrainingQueuedJob]:
+        """List Batch jobs according to status and share.
+
+        Args:
+            status: Batch job status.
+            share_identifier: Batch fairshare share identifier.
+
+        Returns: A list of QueuedJob.
+
+        """
+        filters = None
+        if share_identifier:
+            filters = [{"name": "SHARE_IDENTIFIER", "values": [share_identifier]}]
+
+        jobs_to_return = []
+        next_token = None
+        for job_result_dict in _list_service_job(self.queue_name, status, filters, next_token):
+            for job_result in job_result_dict.get("jobSummaryList", []):
+                if "jobArn" in job_result and "jobName" in job_result:
+                    jobs_to_return.append(
+                        TrainingQueuedJob(job_result["jobArn"], job_result["jobName"], job_result.get("shareIdentifier", None))
                     )
                 else:
                     logging.warning("Missing JobArn or JobName in Batch ListJobs API")

sagemaker-train/src/sagemaker/train/aws_batch/training_queued_job.py

@@ -45,9 +45,10 @@ class TrainingQueuedJob:
     With this class, customers are able to attach the latest training job to a ModelTrainer.
     """
 
-    def __init__(self, job_arn: str, job_name: str):
+    def __init__(self, job_arn: str, job_name: str, share_identifier: Optional[str] = None):
         self.job_arn = job_arn
         self.job_name = job_name
+        self.share_identifier = share_identifier
         self._no_training_job_status = {"SUBMITTED", "PENDING", "RUNNABLE"}
 
     def get_model_trainer(self) -> ModelTrainer:

sagemaker-train/tests/unit/train/aws_batch/conftest.py

@@ -141,6 +141,24 @@
     "nextToken": None,
 }
 
+LIST_SERVICE_JOB_BY_SHARE_RESP_WITH_JOBS = {
+    "jobSummaryList": [
+        {
+            "jobName": JOB_NAME,
+            "jobArn": JOB_ARN,
+            "jobId": JOB_ID,
+            "shareIdentifier": SHARE_IDENTIFIER,
+        },
+        {
+            "jobName": "another-job",
+            "jobArn": "arn:aws:batch:us-west-2:123456789012:job/another-id",
+            "jobId": "another-id",
+            "shareIdentifier": "another-share-identifier",
+        },
+    ],
+    "nextToken": None,
+}
+
 LIST_SERVICE_JOB_RESP_WITH_NEXT_TOKEN = {
     "jobSummaryList": [
         {"jobName": JOB_NAME, "jobArn": JOB_ARN, "jobId": JOB_ID},

sagemaker-train/tests/unit/train/aws_batch/test_training_queue.py

@@ -29,6 +29,7 @@
     DEFAULT_SAGEMAKER_TRAINING_RETRY_CONFIG,
     SUBMIT_SERVICE_JOB_RESP,
     LIST_SERVICE_JOB_RESP_WITH_JOBS,
+    LIST_SERVICE_JOB_BY_SHARE_RESP_WITH_JOBS,
     LIST_SERVICE_JOB_RESP_EMPTY,
     TRAINING_JOB_PAYLOAD,
 )
@@ -279,14 +280,14 @@ def test_list_jobs_with_name_filter(self, mock_list_service_job):
 
         # Verify list_service_job was called
         mock_list_service_job.assert_called_once()
-        
+
         # Get the call arguments - list_service_job is called with positional args:
         # list_service_job(queue_name, status, filters, next_token)
         call_args = mock_list_service_job.call_args[0]
-        
+
         # The 3rd positional argument (index 2) is filters
         filters = call_args[2] if len(call_args) > 2 else None
-        
+
         # Verify filters contain the job name
         assert filters is not None, "Filters should be passed to list_service_job"
         assert filters[0]["name"] == "JOB_NAME", "JOB_NAME filter should be present"
@@ -303,6 +304,56 @@ def test_list_jobs_empty(self, mock_list_service_job):
         assert len(jobs) == 0
 
 
+class TestTrainingQueueListByShare:
+    """Tests for TrainingQueue.list_jobs_by_share method"""
+
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
+    def test_list_jobs_by_share_default(self, mock_list_service_job):
+        """Test list_jobs_by_share with default parameters"""
+        mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_BY_SHARE_RESP_WITH_JOBS])
+
+        queue = TrainingQueue(JOB_QUEUE)
+        jobs = queue.list_jobs_by_share()
+
+        assert len(jobs) == 2
+        assert jobs[0].share_identifier == SHARE_IDENTIFIER
+
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
+    def test_list_jobs_by_share_with_share_filter(self, mock_list_service_job):
+        """Test list_jobs_by_share with job name filter"""
+        mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_BY_SHARE_RESP_WITH_JOBS])
+
+        queue = TrainingQueue(JOB_QUEUE)
+        jobs = queue.list_jobs_by_share(share_identifier=SHARE_IDENTIFIER)
+
+        # Verify list_service_job was called
+        mock_list_service_job.assert_called_once()
+
+        # Get the call arguments - list_service_job is called with positional args:
+        # list_service_job(queue_name, status, filters, next_token)
+        call_args = mock_list_service_job.call_args[0]
+
+        # The 3rd positional argument (index 2) is filters
+        filters = call_args[2] if len(call_args) > 2 else None
+
+        # Verify filters contain the share identifier
+        assert filters is not None, "Filters should be passed to list_service_job"
+        assert filters[0]["name"] == "SHARE_IDENTIFIER", "SHARE_IDENTIFIER filter should be present"
+        assert filters[0]["values"] == [
+            SHARE_IDENTIFIER
+        ], "Filter values should contain the share identifier"
+
+    @patch("sagemaker.train.aws_batch.training_queue._list_service_job")
+    def test_list_jobs_by_share_empty(self, mock_list_service_job):
+        """Test list_jobs_by_share returns empty list"""
+        mock_list_service_job.return_value = iter([LIST_SERVICE_JOB_RESP_EMPTY])
+
+        queue = TrainingQueue(JOB_QUEUE)
+        jobs = queue.list_jobs_by_share()
+
+        assert len(jobs) == 0
+
+
 class TestTrainingQueueGet:
     """Tests for TrainingQueue.get_job method"""