fix: wait=True respects sagemaker_session for processing and training jobs (5765)

Author: aviruthen

sagemaker-core/src/sagemaker/core/helper/session_helper.py

@@ -1625,6 +1625,46 @@ def wait_for_optimization_job(self, job, poll=5):
         _check_job_status(job, desc, "OptimizationJobStatus")
         return desc
 
+    def wait_for_processing_job(self, job, poll=5):
+        """Wait for an Amazon SageMaker Processing job to complete.
+
+        Args:
+            job (str): Name of the processing job to wait for.
+            poll (int): Polling interval in seconds (default: 5).
+
+        Returns:
+            (dict): Return value from the ``DescribeProcessingJob`` API.
+
+        Raises:
+            exceptions.CapacityError: If the processing job fails with CapacityError.
+            exceptions.UnexpectedStatusException: If the processing job fails.
+        """
+        desc = _wait_until(
+            lambda: _processing_job_status(self.sagemaker_client, job), poll
+        )
+        _check_job_status(job, desc, "ProcessingJobStatus")
+        return desc
+
+    def wait_for_training_job(self, job, poll=5):
+        """Wait for an Amazon SageMaker Training job to complete.
+
+        Args:
+            job (str): Name of the training job to wait for.
+            poll (int): Polling interval in seconds (default: 5).
+
+        Returns:
+            (dict): Return value from the ``DescribeTrainingJob`` API.
+
+        Raises:
+            exceptions.CapacityError: If the training job fails with CapacityError.
+            exceptions.UnexpectedStatusException: If the training job fails.
+        """
+        desc = _wait_until(
+            lambda: _training_job_status(self.sagemaker_client, job), poll
+        )
+        _check_job_status(job, desc, "TrainingJobStatus")
+        return desc
+
     def update_inference_component(
         self, inference_component_name, specification=None, runtime_config=None, wait=True
     ):
@@ -2896,6 +2936,70 @@ def _optimization_job_status(sagemaker_client, job_name):
     return desc
 
 
+def _processing_job_status(sagemaker_client, job_name):
+    """Check the status of a processing job.
+
+    Args:
+        sagemaker_client: The boto3 SageMaker client.
+        job_name (str): The name of the processing job.
+
+    Returns:
+        dict: The processing job description if complete, None if still in progress.
+    """
+    status_codes = {
+        "Completed": "!",
+        "InProgress": ".",
+        "Failed": "*",
+        "Stopped": "s",
+        "Stopping": "_",
+    }
+    in_progress_statuses = ["InProgress", "Stopping", "Starting"]
+
+    desc = sagemaker_client.describe_processing_job(ProcessingJobName=job_name)
+    status = desc["ProcessingJobStatus"]
+
+    status = _STATUS_CODE_TABLE.get(status, status)
+    print(status_codes.get(status, "?"), end="")
+    sys.stdout.flush()
+
+    if status in in_progress_statuses:
+        return None
+
+    return desc
+
+
+def _training_job_status(sagemaker_client, job_name):
+    """Check the status of a training job.
+
+    Args:
+        sagemaker_client: The boto3 SageMaker client.
+        job_name (str): The name of the training job.
+
+    Returns:
+        dict: The training job description if complete, None if still in progress.
+    """
+    status_codes = {
+        "Completed": "!",
+        "InProgress": ".",
+        "Failed": "*",
+        "Stopped": "s",
+        "Stopping": "_",
+    }
+    in_progress_statuses = ["InProgress", "Stopping", "Starting"]
+
+    desc = sagemaker_client.describe_training_job(TrainingJobName=job_name)
+    status = desc["TrainingJobStatus"]
+
+    status = _STATUS_CODE_TABLE.get(status, status)
+    print(status_codes.get(status, "?"), end="")
+    sys.stdout.flush()
+
+    if status in in_progress_statuses:
+        return None
+
+    return desc
+
+
 def container_def(
     image_uri,
     model_data_url=None,

sagemaker-core/src/sagemaker/core/processing.py

@@ -296,7 +296,9 @@ def run(
         if not isinstance(self.sagemaker_session, PipelineSession):
             self.jobs.append(self.latest_job)
             if wait:
-                self.latest_job.wait(logs=logs)
+                self.sagemaker_session.wait_for_processing_job(
+                    self.latest_job.processing_job_name
+                )
 
     def _extend_processing_args(self, inputs, outputs, **kwargs):  # pylint: disable=W0613
         """Extend inputs and outputs based on extra parameters"""
@@ -846,7 +848,9 @@ def run(
         if not isinstance(self.sagemaker_session, PipelineSession):
             self.jobs.append(self.latest_job)
             if wait:
-                self.latest_job.wait(logs=logs)
+                self.sagemaker_session.wait_for_processing_job(
+                    self.latest_job.processing_job_name
+                )
 
     def _include_code_in_inputs(self, inputs, code, kms_key=None):
         """Converts code to appropriate input and includes in input list.

sagemaker-train/src/sagemaker/train/common_utils/trainer_wait.py

@@ -15,6 +15,34 @@
 from sagemaker.train.common_utils.mlflow_metrics_util import _MLflowMetricsUtil
 
 
+def _refresh_training_job(training_job, sagemaker_session=None):
+    """Refresh a training job using the session-aware client if available.
+
+    When sagemaker_session is provided, re-fetches the training job via
+    TrainingJob.get() with the user's boto_session, which avoids the
+    NoCredentialsError that occurs when refresh() uses the global default client.
+
+    Args:
+        training_job (TrainingJob): The training job to refresh.
+        sagemaker_session: SageMaker session with the correct credentials.
+            If None, falls back to the default _refresh_training_job(training_job, sagemaker_session).
+    """
+    if sagemaker_session is not None:
+        refreshed = TrainingJob.get(
+            training_job_name=training_job.training_job_name,
+            session=sagemaker_session.boto_session,
+        )
+        # Copy refreshed attributes back to the original object
+        for attr in ("training_job_status", "secondary_status", "failure_reason"):
+            if hasattr(refreshed, attr):
+                try:
+                    setattr(training_job, attr, getattr(refreshed, attr))
+                except (AttributeError, TypeError):
+                    pass
+    else:
+        _refresh_training_job(training_job, sagemaker_session)
+
+
 @contextmanager
 def _suppress_info_logging():
     """Context manager to temporarily suppress INFO level logging."""
@@ -218,14 +246,18 @@ def get_mlflow_url(training_job) -> str:
 def wait(
         training_job: TrainingJob,
         poll: int = 5,
-        timeout: Optional[int] = 3000
+        timeout: Optional[int] = 3000,
+        sagemaker_session=None,
 ) -> None:
     """Wait for training job to complete with progress tracking.
 
     Args:
         training_job (TrainingJob): The SageMaker training job to monitor.
-        poll (int): Polling interval in seconds. Defaults to 3.
-        timeout (Optional[int]): Maximum wait time in seconds. Defaults to None.
+        poll (int): Polling interval in seconds. Defaults to 5.
+        timeout (Optional[int]): Maximum wait time in seconds. Defaults to 3000.
+        sagemaker_session: SageMaker session to use for describe calls.
+            If provided, uses the session's sagemaker_client instead of the
+            global default client, fixing NoCredentialsError with assumed-role sessions.
 
     Raises:
         FailedStatusError: If the training job fails.
@@ -277,7 +309,7 @@ def get_cached_mlflow_url():
                     iteration += 1
                     time.sleep(0.5)
                     if iteration >= poll * 2:
-                        training_job.refresh()
+                        _refresh_training_job(training_job, sagemaker_session)
                         iteration = 0
 
                     status = training_job.training_job_status
@@ -360,7 +392,7 @@ def get_cached_mlflow_url():
                         if not progress_started:
                             progress_started = True
                             time.sleep(poll)
-                            training_job.refresh()
+                            _refresh_training_job(training_job, sagemaker_session)
 
                         training_progress_pct, training_progress_text = _calculate_training_progress(
                             training_job.progress_info, metrics_util, mlflow_run_name, training_job
@@ -442,7 +474,7 @@ def get_cached_mlflow_url():
             while True:
                 iteration += 1
                 time.sleep(poll)
-                training_job.refresh()
+                _refresh_training_job(training_job, sagemaker_session)
 
                 status = training_job.training_job_status
                 secondary_status = training_job.secondary_status
@@ -462,7 +494,7 @@ def get_cached_mlflow_url():
                             if not progress_started:
                                 progress_started = True
                                 time.sleep(20)
-                                training_job.refresh()
+                                _refresh_training_job(training_job, sagemaker_session)
 
                             progress_pct, progress_text = _calculate_training_progress(
                                 training_job.progress_info, metrics_util, mlflow_run_name, training_job

sagemaker-train/src/sagemaker/train/model_trainer.py

@@ -790,7 +790,11 @@ def train(
             self._latest_training_job = training_job
 
             if wait:
-                training_job.wait(logs=logs)
+                from sagemaker.train.common_utils.trainer_wait import wait as trainer_wait
+                trainer_wait(
+                    training_job=training_job,
+                    sagemaker_session=self.sagemaker_session,
+                )
             if logs and not wait:
                 logger.warning(
                     "Not displaing the training container logs as 'wait' is set to False."