feature: add verbose logging from cloudwatch for IR job (aws#790)

* feature: Add create inf rec api to session

* save

* fix error handling in submit. update docstring

* add in job_name param

* fix: type hint of PySparkProcessor __init__ (aws#3297)

* fix: fix PySparkProcessor __init__ params type (aws#3354)

* feature: add describe inf rec api to session

* feature: add describe inf rec api to session

* feature: add describe inf rec api to session

* add verbose logging in wait_for_inference_recommendations_job

* improve unit tests

* switch to cloudwatch logs

* add unit tests

* fix unit test

* make default logging level to Verbose

* address comments in PR

* address some comments

* fix typo missing s

* add initial print because it took a while for cw logstream to appear

* print and conditional optimizations and update tests

* optimize log print

* fix polling for logs empty events

* cover empty event case

* add spacing

Co-authored-by: Gary Wang <[email protected]>
Co-authored-by: Kevin <[email protected]>
Co-authored-by: André Perez <[email protected]>
Co-authored-by: Raymond Liu <[email protected]>

Author: 4 people

src/sagemaker/session.py

@@ -53,10 +53,12 @@
 _STATUS_CODE_TABLE = {
     "COMPLETED": "Completed",
     "INPROGRESS": "InProgress",
+    "IN_PROGRESS": "InProgress",
     "FAILED": "Failed",
     "STOPPED": "Stopped",
     "STOPPING": "Stopping",
     "STARTING": "Starting",
+    "PENDING": "Pending",
 }
 
 
@@ -4844,6 +4846,41 @@ def submit(request):
         )
         return job_name
 
+    def wait_for_inference_recommendations_job(
+        self, job_name: str, poll: int = 120, log_level: str = "Verbose"
+    ) -> Dict[str, Any]:
+        """Wait for an Amazon SageMaker Inference Recommender job to complete.
+
+        Args:
+            job_name (str): Name of the Inference Recommender job to wait for.
+            poll (int): Polling interval in seconds (default: 120).
+            log_level (str): The level of verbosity for the logs.
+            Can be "Quiet" or "Verbose" (default: "Quiet").
+
+        Returns:
+            (dict): Return value from the ``DescribeInferenceRecommendationsJob`` API.
+
+        Raises:
+            exceptions.CapacityError: If the Inference Recommender job fails with CapacityError.
+            exceptions.UnexpectedStatusException: If the Inference Recommender job fails.
+        """
+        if log_level == "Quiet":
+            _wait_until(
+                lambda: _describe_inference_recommendations_job_status(
+                    self.sagemaker_client, job_name
+                ),
+                poll,
+            )
+        elif log_level == "Verbose":
+            _display_inference_recommendations_job_steps_status(
+                self, self.sagemaker_client, job_name
+            )
+        else:
+            raise ValueError("log_level must be either Quiet or Verbose")
+        desc = _describe_inference_recommendations_job_status(self.sagemaker_client, job_name)
+        self._check_job_status(job_name, desc, "Status")
+        return desc
+
 
 def get_model_package_args(
     content_types,
@@ -5465,6 +5502,118 @@ def _create_model_package_status(sagemaker_client, model_package_name):
     return desc
 
 
+def _describe_inference_recommendations_job_status(sagemaker_client, job_name: str):
+    """Describes the status of a job and returns the job description.
+
+    Args:
+        sagemaker_client (boto3.client.sagemaker): A SageMaker client.
+        job_name (str): The name of the job.
+
+    Returns:
+        dict: The job description, or None if the job is still in progress.
+    """
+    inference_recommendations_job_status_codes = {
+        "PENDING": ".",
+        "IN_PROGRESS": ".",
+        "COMPLETED": "!",
+        "FAILED": "*",
+        "STOPPING": "_",
+        "STOPPED": "s",
+    }
+    in_progress_statuses = {"PENDING", "IN_PROGRESS", "STOPPING"}
+
+    desc = sagemaker_client.describe_inference_recommendations_job(JobName=job_name)
+    status = desc["Status"]
+
+    print(inference_recommendations_job_status_codes.get(status, "?"), end="", flush=True)
+
+    if status in in_progress_statuses:
+        return None
+
+    print("")
+    return desc
+
+
+def _display_inference_recommendations_job_steps_status(
+    sagemaker_session, sagemaker_client, job_name: str, poll: int = 60
+):
+    """Placeholder docstring"""
+    cloudwatch_client = sagemaker_session.boto_session.client("logs")
+    in_progress_statuses = {"PENDING", "IN_PROGRESS", "STOPPING"}
+    log_group_name = "/aws/sagemaker/InferenceRecommendationsJobs"
+    log_stream_name = job_name + "/execution"
+
+    initial_logs_batch = get_log_events_for_inference_recommender(
+        cloudwatch_client, log_group_name, log_stream_name
+    )
+    print(f"Retrieved logStream: {log_stream_name} from logGroup: {log_group_name}", flush=True)
+    events = initial_logs_batch["events"]
+    print(*[event["message"] for event in events], sep="\n", flush=True)
+
+    next_forward_token = initial_logs_batch["nextForwardToken"] if events else None
+    flush_remaining = True
+    while True:
+        logs_batch = (
+            cloudwatch_client.get_log_events(
+                logGroupName=log_group_name,
+                logStreamName=log_stream_name,
+                nextToken=next_forward_token,
+            )
+            if next_forward_token
+            else cloudwatch_client.get_log_events(
+                logGroupName=log_group_name, logStreamName=log_stream_name
+            )
+        )
+
+        events = logs_batch["events"]
+
+        desc = sagemaker_client.describe_inference_recommendations_job(JobName=job_name)
+        status = desc["Status"]
+
+        if not events:
+            if status in in_progress_statuses:
+                time.sleep(poll)
+                continue
+            if flush_remaining:
+                flush_remaining = False
+                time.sleep(poll)
+                continue
+
+        next_forward_token = logs_batch["nextForwardToken"]
+        print(*[event["message"] for event in events], sep="\n", flush=True)
+
+        if status not in in_progress_statuses:
+            break
+
+        time.sleep(poll)
+
+
+def get_log_events_for_inference_recommender(cw_client, log_group_name, log_stream_name):
+    """Retrieves log events from the specified CloudWatch log group and log stream.
+
+    Args:
+        cw_client (boto3.client): A boto3 CloudWatch client.
+        log_group_name (str): The name of the CloudWatch log group.
+        log_stream_name (str): The name of the CloudWatch log stream.
+
+    Returns:
+        (dict): A dictionary containing log events from CloudWatch log group and log stream.
+    """
+    print("Fetching logs from CloudWatch...", flush=True)
+    for _ in retries(
+        max_retry_count=30,  # 30*10 = 5min
+        exception_message_prefix="Waiting for cloudwatch stream to appear. ",
+        seconds_to_sleep=10,
+    ):
+        try:
+            return cw_client.get_log_events(
+                logGroupName=log_group_name, logStreamName=log_stream_name
+            )
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "ResourceNotFoundException":
+                pass
+
+
 def _deploy_done(sagemaker_client, endpoint_name):
     """Placeholder docstring"""
     hosting_status_codes = {

tests/unit/test_session.py

@@ -23,7 +23,7 @@
 from mock import ANY, MagicMock, Mock, patch, call, mock_open
 
 import sagemaker
-from sagemaker import TrainingInput, Session, get_execution_role
+from sagemaker import TrainingInput, Session, get_execution_role, exceptions
 from sagemaker.async_inference import AsyncInferenceConfig
 from sagemaker.session import (
     _tuning_job_status,
@@ -2267,7 +2267,6 @@ def test_train_done_in_progress(sagemaker_session):
     "GenerateCandidateDefinitionsOnly": False,
 }
 
-
 COMPLETE_EXPECTED_AUTO_ML_JOB_ARGS = {
     "AutoMLJobName": JOB_NAME,
     "InputDataConfig": [
@@ -3112,3 +3111,160 @@ def test_create_inference_recommendations_job_propogate_other_exception(sagemake
         )
 
     assert "AccessDeniedException" in str(error)
+
+
+DEFAULT_LOG_EVENTS_INFERENCE_RECOMMENDER = [
+    MockBotoException("ResourceNotFoundException"),
+    {"nextForwardToken": None, "events": [{"timestamp": 1, "message": "hi there #1"}]},
+    {"nextForwardToken": None, "events": [{"timestamp": 2, "message": "hi there #2"}]},
+    {"nextForwardToken": None, "events": [{"timestamp": 3, "message": "hi there #3"}]},
+    {"nextForwardToken": None, "events": [{"timestamp": 4, "message": "hi there #4"}]},
+]
+
+FLUSH_LOG_EVENTS_INFERENCE_RECOMMENDER = [
+    MockBotoException("ResourceNotFoundException"),
+    {"nextForwardToken": None, "events": [{"timestamp": 1, "message": "hi there #1"}]},
+    {"nextForwardToken": None, "events": [{"timestamp": 2, "message": "hi there #2"}]},
+    {"nextForwardToken": None, "events": []},
+    {"nextForwardToken": None, "events": [{"timestamp": 3, "message": "hi there #3"}]},
+    {"nextForwardToken": None, "events": []},
+    {"nextForwardToken": None, "events": [{"timestamp": 4, "message": "hi there #4"}]},
+]
+
+INFERENCE_RECOMMENDATIONS_DESC_STATUS_PENDING = {"Status": "PENDING"}
+INFERENCE_RECOMMENDATIONS_DESC_STATUS_IN_PROGRESS = {"Status": "IN_PROGRESS"}
+INFERENCE_RECOMMENDATIONS_DESC_STATUS_COMPLETED = {"Status": "COMPLETED"}
+
+
+@pytest.fixture()
+def sm_session_inference_recommender():
+    boto_mock = MagicMock(name="boto_session")
+    boto_mock.client("logs").get_log_events.side_effect = DEFAULT_LOG_EVENTS_INFERENCE_RECOMMENDER
+
+    ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
+
+    ims.sagemaker_client.describe_inference_recommendations_job.side_effect = [
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_PENDING,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_IN_PROGRESS,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_COMPLETED,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_COMPLETED,
+    ]
+
+    return ims
+
+
+@pytest.fixture()
+def sm_session_inference_recommender_flush():
+    boto_mock = MagicMock(name="boto_session")
+    boto_mock.client("logs").get_log_events.side_effect = FLUSH_LOG_EVENTS_INFERENCE_RECOMMENDER
+
+    ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
+
+    ims.sagemaker_client.describe_inference_recommendations_job.side_effect = [
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_PENDING,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_IN_PROGRESS,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_IN_PROGRESS,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_COMPLETED,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_COMPLETED,
+        INFERENCE_RECOMMENDATIONS_DESC_STATUS_COMPLETED,
+    ]
+
+    return ims
+
+
+@patch("time.sleep")
+def test_wait_for_inference_recommendations_job_completed(sleep, sm_session_inference_recommender):
+    assert (
+        sm_session_inference_recommender.wait_for_inference_recommendations_job(
+            JOB_NAME, log_level="Quiet"
+        )["Status"]
+        == "COMPLETED"
+    )
+
+    assert (
+        4
+        == sm_session_inference_recommender.sagemaker_client.describe_inference_recommendations_job.call_count
+    )
+    assert 2 == sleep.call_count
+    sleep.assert_has_calls([call(120), call(120)])
+
+
+def test_wait_for_inference_recommendations_job_failed(sagemaker_session):
+    inference_recommendations_desc_status_failed = {
+        "Status": "FAILED",
+        "FailureReason": "Mock Failure Reason",
+    }
+
+    sagemaker_session.sagemaker_client.describe_inference_recommendations_job = Mock(
+        name="describe_inference_recommendations_job",
+        return_value=inference_recommendations_desc_status_failed,
+    )
+
+    with pytest.raises(exceptions.UnexpectedStatusException) as error:
+        sagemaker_session.wait_for_inference_recommendations_job(JOB_NAME)
+
+    assert "Mock Failure Reason" in str(error)
+
+
+@patch("builtins.print")
+@patch("time.sleep")
+def test_wait_for_inference_recommendations_job_completed_verbose(
+    sleep, mock_print, sm_session_inference_recommender
+):
+    assert (
+        sm_session_inference_recommender.wait_for_inference_recommendations_job(
+            JOB_NAME, log_level="Verbose"
+        )["Status"]
+        == "COMPLETED"
+    )
+    assert (
+        4
+        == sm_session_inference_recommender.sagemaker_client.describe_inference_recommendations_job.call_count
+    )
+
+    assert (
+        5 == sm_session_inference_recommender.boto_session.client("logs").get_log_events.call_count
+    )
+
+    assert 3 == sleep.call_count
+    sleep.assert_has_calls([call(10), call(60), call(60)])
+
+    assert 8 == mock_print.call_count
+
+
+@patch("builtins.print")
+@patch("time.sleep")
+def test_wait_for_inference_recommendations_job_flush_completed(
+    sleep, mock_print, sm_session_inference_recommender_flush
+):
+    assert (
+        sm_session_inference_recommender_flush.wait_for_inference_recommendations_job(
+            JOB_NAME, log_level="Verbose"
+        )["Status"]
+        == "COMPLETED"
+    )
+    assert (
+        6
+        == sm_session_inference_recommender_flush.sagemaker_client.describe_inference_recommendations_job.call_count
+    )
+
+    assert (
+        7
+        == sm_session_inference_recommender_flush.boto_session.client(
+            "logs"
+        ).get_log_events.call_count
+    )
+
+    assert 5 == sleep.call_count
+    sleep.assert_has_calls([call(10), call(60), call(60), call(60), call(60)])
+
+    assert 8 == mock_print.call_count
+
+
+def test_wait_for_inference_recommendations_job_invalid_log_level(sagemaker_session):
+    with pytest.raises(ValueError) as error:
+        sagemaker_session.wait_for_inference_recommendations_job(
+            JOB_NAME, log_level="invalid_log_level"
+        )
+
+    assert "log_level must be either Quiet or Verbose" in str(error)