From 2322ccefc39b8c042687c264e69af98006f5d08a Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Mon, 12 Oct 2026 12:46:25 -0500 Subject: [PATCH] fix(core): truncate the Run failure status message to the API limit When an exception escapes a 'with Run(...)' block, Run.__exit__ records it as the trial component's failure status with message=str(exc_value). The UpdateTrialComponent API limits Status.Message to 1024 characters, so a long exception message (a stack dump, an OOM message) made close() fail with a ValidationException from the service. That error replaced the user's original exception, and the run was never marked Failed or given an end time. v2 truncates the message (#5276); the v3 port lost it. Truncate the message to MAX_STATUS_MESSAGE_LEN (1024), as v2 does. --- .../src/sagemaker/core/experiments/run.py | 4 +- .../tests/unit/experiments/test_run_exit.py | 49 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) create mode 100644 sagemaker-core/tests/unit/experiments/test_run_exit.py diff --git a/sagemaker-core/src/sagemaker/core/experiments/run.py b/sagemaker-core/src/sagemaker/core/experiments/run.py index 2082fc9f5b..79e78188a1 100644 --- a/sagemaker-core/src/sagemaker/core/experiments/run.py +++ b/sagemaker-core/src/sagemaker/core/experiments/run.py @@ -69,6 +69,8 @@ TRIAL_NAME_TEMPLATE = "Default-Run-Group-{}" MAX_RUN_TC_ARTIFACTS_LEN = 30 MAX_NAME_LEN_IN_BACKEND = 120 +# UpdateTrialComponent rejects a Status.Message longer than this. +MAX_STATUS_MESSAGE_LEN = 1024 EXPERIMENT_NAME = "ExperimentName" TRIAL_NAME = "TrialName" RUN_NAME = "RunName" @@ -760,7 +762,7 @@ def __exit__(self, exc_type, exc_value, exc_traceback): if exc_value: self._trial_component.status = _api_types.TrialComponentStatus( primary_status=_TrialComponentStatusType.Failed.value, - message=str(exc_value), + message=(str(exc_value) or "")[:MAX_STATUS_MESSAGE_LEN], ) else: self._trial_component.status = _api_types.TrialComponentStatus( diff --git a/sagemaker-core/tests/unit/experiments/test_run_exit.py b/sagemaker-core/tests/unit/experiments/test_run_exit.py new file mode 100644 index 0000000000..294b45c974 --- /dev/null +++ b/sagemaker-core/tests/unit/experiments/test_run_exit.py @@ -0,0 +1,49 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Status recorded by Run.__exit__.""" + +from __future__ import absolute_import + +from unittest.mock import MagicMock, patch + +import pytest + +from sagemaker.core.experiments.run import Run +from sagemaker.core.experiments.trial_component import _TrialComponent + + +def _run_with_trial_component(): + # Skip the network-bound __init__; __exit__ only needs these attributes. + run = Run.__new__(Run) + run._trial_component = _TrialComponent( + sagemaker_session=MagicMock(), trial_component_name="trial-component" + ) + run._in_load = False + run._inside_init_context = True + run._inside_load_context = False + return run + + +@pytest.mark.parametrize("message_length, expected_length", [(2000, 1024), (10, 10)]) +def test_exit_with_exception_limits_status_message_to_api_maximum(message_length, expected_length): + """UpdateTrialComponent rejects Status.Message longer than 1024 characters.""" + run = _run_with_trial_component() + error = RuntimeError("x" * message_length) + + with patch.object(Run, "close") as mock_close: + run.__exit__(type(error), error, None) + + status = run._trial_component.status + assert status.primary_status == "Failed" + assert len(status.message) == expected_length + mock_close.assert_called_once()