fix: Default AsyncPredictor upload prefix to endpoint name (#6336)

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

Author: lucasjia-aws

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
@@ -195,10 +198,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),
             )
 

sagemaker-serve/tests/unit/test_predictor_async.py

@@ -90,6 +90,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()