From 1dda496874648301f27e3ba37c61ca72daf6d282 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Fri, 25 Sep 2026 14:19:19 -0700 Subject: [PATCH] fix: Default AsyncPredictor upload prefix to endpoint name AsyncPredictor accepts name=None by default, but predict() and predict_async() called name_from_base(self.name) when uploading input data without an input_path, which raised "TypeError: 'NoneType' object is not subscriptable". Fall back to the wrapped predictor's endpoint name when no name is given, and document the name argument. An explicitly provided name is still used as the S3 key prefix. Fixes #3210 Fixes #4774 --- .../src/sagemaker/serve/predictor_async.py | 8 +++- .../tests/unit/test_predictor_async.py | 45 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/sagemaker-serve/src/sagemaker/serve/predictor_async.py b/sagemaker-serve/src/sagemaker/serve/predictor_async.py index 31e94318d5..7bc164ac0e 100644 --- a/sagemaker-serve/src/sagemaker/serve/predictor_async.py +++ b/sagemaker-serve/src/sagemaker/serve/predictor_async.py @@ -40,6 +40,9 @@ def __init__( predictor (sagemaker.predictor.Predictor): General ``Predictor`` object has useful methods and variables. ``AsyncPredictor`` stands on top of it with capability for async inference. + name (str): Optional. Name used as the prefix of the Amazon S3 key when + input data is uploaded for async inference. If not specified, the + endpoint name is used. (Default: None) """ self.predictor = predictor self.endpoint_name = predictor.endpoint_name @@ -169,10 +172,13 @@ def _upload_data_to_s3( my_uuid = str(uuid.uuid4()) timestamp = sagemaker_timestamp() bucket = self.sagemaker_session.default_bucket() + # ``name`` is optional; fall back to the endpoint name so the default + # ``AsyncPredictor(predictor)`` construction can upload input data. + base_name = self.name or self.endpoint_name key = s3.s3_path_join( self.sagemaker_session.default_bucket_prefix, "async-endpoint-inputs", - name_from_base(self.name, short=True), + name_from_base(base_name, short=True), "{}-{}".format(timestamp, my_uuid), ) diff --git a/sagemaker-serve/tests/unit/test_predictor_async.py b/sagemaker-serve/tests/unit/test_predictor_async.py index ccd011a973..fbeeea0bbb 100644 --- a/sagemaker-serve/tests/unit/test_predictor_async.py +++ b/sagemaker-serve/tests/unit/test_predictor_async.py @@ -86,6 +86,51 @@ def test_upload_data_to_s3(self, mock_parse): self.assertEqual(result, "s3://bucket/key") async_predictor.s3_client.put_object.assert_called_once() + def test_upload_data_to_s3_without_name_uses_endpoint_name(self): + # Regression test for https://github.com/aws/sagemaker-python-sdk/issues/3210 + self.mock_predictor.serializer.serialize.return_value = b"serialized_data" + + async_predictor = AsyncPredictor(self.mock_predictor) + async_predictor.sagemaker_session.default_bucket.return_value = "default-bucket" + async_predictor.sagemaker_session.default_bucket_prefix = None + + result = async_predictor._upload_data_to_s3("test_data") + + self.assertIsNone(async_predictor.name) + key = async_predictor.s3_client.put_object.call_args.kwargs["Key"] + self.assertTrue(key.startswith("async-endpoint-inputs/test-endpoint-")) + self.assertEqual(result, "s3://default-bucket/{}".format(key)) + + def test_upload_data_to_s3_with_name_uses_name(self): + self.mock_predictor.serializer.serialize.return_value = b"serialized_data" + + async_predictor = AsyncPredictor(self.mock_predictor, name="my-async") + async_predictor.sagemaker_session.default_bucket.return_value = "default-bucket" + async_predictor.sagemaker_session.default_bucket_prefix = None + + async_predictor._upload_data_to_s3("test_data") + + key = async_predictor.s3_client.put_object.call_args.kwargs["Key"] + self.assertTrue(key.startswith("async-endpoint-inputs/my-async-")) + + @patch.object(AsyncPredictor, "_submit_async_request") + def test_predict_async_with_data_and_no_name(self, mock_submit): + # Regression test for https://github.com/aws/sagemaker-python-sdk/issues/3210 + mock_submit.return_value = { + "OutputLocation": "s3://bucket/output", + "FailureLocation": "s3://bucket/failure", + } + self.mock_predictor.serializer.serialize.return_value = b"serialized_data" + + async_predictor = AsyncPredictor(self.mock_predictor) + async_predictor.sagemaker_session.default_bucket.return_value = "default-bucket" + async_predictor.sagemaker_session.default_bucket_prefix = None + + response = async_predictor.predict_async(data="test_data") + + self.assertEqual(response.output_path, "s3://bucket/output") + async_predictor.s3_client.put_object.assert_called_once() + def test_delete_endpoint(self): async_predictor = AsyncPredictor(self.mock_predictor) async_predictor.delete_endpoint()