One client for each sagemaker endpoint vs multiple client

[mohitk1995]

The below snippet make use single client to make request to multiple sagemaker endpoints with different payloads. Does the below implementation makes sense? Or should we use different client for each endpoint?

import json
from time import perf_counter

from aiobotocore.client import AioBaseClient
from aiobotocore.session import get_session
from asyncer import asyncify


async def call_sagemaker(client, payload):
    invocation_details = [
        {
            "endpoint_name": "endpoint-1",
            "payload": payload,
            "content_type": "application/json",
        },
        {"endpoint_name": "endpoint-2", "payload": payload, "content_type": "application/json"},
        {
            "endpoint_name": "endpoint-3",
            "payload": payload,
            "content_type": "application/json",
        },
    ]
    tasks = [
        asyncio.create_task(
            _invoke_endpoint_helper(
                client=client,
                endpoint_name=invocation_detail["endpoint_name"],
                payload=invocation_detail["payload"],
                content_type=invocation_detail["content_type"],
            )
        )
        for invocation_detail in invocation_details
    ]
    res = await asyncio.gather(*tasks, return_exceptions=True)
    return res


async def _invoke_endpoint_helper(client, endpoint_name, payload, content_type):
    invoke_endpoint_cr = await asyncify(client.invoke_endpoint)(
        EndpointName=endpoint_name,
        ContentType=content_type,
        Body=json.dumps(payload),
        CustomAttributes="accept_eula=true",
    )
    response = await invoke_endpoint_cr
    response = await response["Body"].read()
    return json.loads(response.decode("utf8"))


async def get_new_client() -> AioBaseClient:
    session = get_session()
    client_ctx = session.create_client(
        "sagemaker-runtime",
        region_name="us-west-2",
    )
    client = await client_ctx.__aenter__()
    return client


async def main():
    client = await get_new_client()
    payload = json.dumps({"inputs": "How to create policy?"})
    for i in range(5):
        start_time = perf_counter()
        responses = await call_sagemaker(client, payload)
        print(responses)
        print(f"Time taken: {perf_counter() - start_time}")
    await client.__aexit__(None, None, None)


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

[geraldstanje1]

cc @tim-finnigan @RyanFitzSimmonsAK @indrora @nateprewitt could you answer above?

[tim-finnigan]

I think using a single client makes sense, but since aiobotocore is being used here you can also consider asking for help in their repository: https://github.com/aio-libs/aiobotocore. Or for general use questions like this in the re:Post forums: https://repost.aws/.

[geraldstanje1]

@tim-finnigan the main question is: can you share the same client/session across multiple endpoint calls with async - is that safe?

but since aiobotocore is being used here you can also consider asking for help in their repository: https://github.com/aio-libs/aiobotocore.

this question was asked in aio-libs/aiobotocore#1119 - than they said reach out here - we go in circles :)

[tim-finnigan]

Please see: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/clients.html#multithreading-or-multiprocessing-with-clients

[geraldstanje1]

yes i saw that, that why we asked it here to clarify - does that mean we cannot share the same client/session across multiple endpoint calls with async?

[tim-finnigan]

Regarding the async behavior with aiobotocore it looks like someone already commented:

"that code won't work, after you return the client it will be closed...I highly recommend re-structuring your code, you're instructing aiobotocore to close the client and session after get_client returns."

If you're asking about Boto3/Botocore/SageMaker behavior outside of aiobotocore, please be more specific. If you can, provide a Boto3 code snippet illustrating what you're trying to do. In addition to the invoke_endpoint command there is also the invoke_endpoint_async command, in which "Inference requests sent to this API are enqueued for asynchronous processing.". But note both of these commands involve calls to the upstream SageMaker APIs InvokeEndpoint and InvokeEndpointAsync.

[geraldstanje1]

hi @tim-finnigan i refer to the following - you see that 4 endpoints (client.invoke_endpoint) are called using asyncio.gather - meaning those should run in parallel and they share the same client/session:

tasks = [
        asyncio.create_task(
            _invoke_endpoint_helper(
                client=client,
                endpoint_name=invocation_detail["endpoint_name"],
                payload=invocation_detail["payload"],
                content_type=invocation_detail["content_type"],
            )
        )
        for invocation_detail in invocation_details
    ]
    res = await asyncio.gather(*tasks, return_exceptions=True)

async def _invoke_endpoint_helper(client, endpoint_name, payload, content_type):
    invoke_endpoint_cr = await asyncify(client.invoke_endpoint)(
        EndpointName=endpoint_name,
        ContentType=content_type,
        Body=json.dumps(payload),
        CustomAttributes="accept_eula=true",
    )
    response = await invoke_endpoint_cr
    response = await response["Body"].read()
    return json.loads(response.decode("utf8"))