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()