Sagemaker runtime: InvokeEndpointAsync does not provide FailureLocation parameter

[codeneobee]

Moto version: 5.0.12

When using moto to mock calls to invoke_endpoint_async in the Sagemaker runtime client, the output only provides the OutputLocation parameter but not the FailureLocation parameter, as specified in https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpointAsync.html. There also does not seem to be a way (that I have found) to override the default response of the invoke_endpoint_async call, it is only possible to override the response that is returned when reading the object at OutputLocation from the S3 bucket.

Expected: The invoke_endpoint_async call should return a FailureLocation object key or there should be an option to override the default response of the call.

[bblommers]

Hi @codeneobee, it is possible to override the regular response already. You can use the same endpoint as when configuring the output for the invoke_endpoint-method. See the docs here: https://docs.getmoto.org/en/latest/docs/services/sagemaker-runtime.html

We're discussing the best way to return/configure a FailureLocation in #7905 - any suggestions are welcome!

[codeneobee]

Hey @bblommers , in the case of invoke_endpoint_async this overrides the response returned from the json inside the S3 bucket at the OutputLocation, but there is no option to override the response of the invoke_endpoint_async call directly

So for example:

@mock_aws
def test_invoke_endpoint_async():
    client = boto3.client("sagemaker-runtime", region_name="us-east-1")
    requests.post(
        "http://motoapi.amazonaws.com/moto-api/static/sagemaker/endpoint-results",
        json= {"results": [{"Body": "some body"}]
    }
 
    body = client.invoke_endpoint_async(        # <------ This is the response I want to be able to override
        EndpointName="foo", InputLocation="bar", InferenceId="baz"
    )
    output_location = body["OutputLocation"]

    bucket_name, obj = bucket_and_name_from_url(output_location)
    resp = s3.get_object(Bucket=bucket_name, Key=obj)
    resp = json.loads(resp["Body"].read().decode("utf-8"))
    assert resp["Body"] == "some body"         # <------ This is the response that gets overriden above

[bblommers]

@codeneobee Moto 5.0.14 (just released) contains a dedicated endpoint to override the results for the async endpoint, including an option to specify failure results. See the docs here: https://docs.getmoto.org/en/latest/docs/services/sagemaker-runtime.html

I think that answers your question, so I'll close this - let us know if you run into any issues though.