handle existing sagemaker deployments gracefully

[samhita-alla]

Tracking issue

Why are the changes needed?

This can be useful for generating deterministic names for sagemaker entities with an "idempotence token", and to avoid raising an error if the deployment already exists, instead returning the existing deployment.

An idempotence token is useful for generating different tokens given different configurations. It avoids name collisions when updating the configuration.

What changes were proposed in this pull request?

• Add try-catch to avoid raising exceptions when the deployment already exists.
• Update the endpoint agent's output.
• Add "idempotence token" as the mixin's output.
• Update the workflow's output.

UX:

• idempotence_token corresponds to the config hash
• inputs.idempotence_token corresponds to the idempotence token of the previous task. Since create_sagemaker_deployment is a workflow, it injects the previous task's idempotence token into the current task as an input.

from flytekit import kwtypes, workflow
from flytekit.types.file import FlyteFile
from flytekitplugins.awssagemaker_inference import create_sagemaker_deployment

sd_deployment = create_sagemaker_deployment(
    name="stable-diffusion",
    model_input_types=kwtypes(deployment_name=str, ...),
    model_config={
        "ModelName": "{inputs.deployment_name}-{idempotence_token}",
        ...
    },
    endpoint_config_input_types=kwtypes(deployment_name=str, ...),
    endpoint_config_config={
        "EndpointConfigName": "{inputs.deployment_name}-{idempotence_token}",
        "ProductionVariants": [
            {
                "ModelName": "{inputs.deployment_name}-{inputs.idempotence_token}",
                ...
            },
        ],
    },
    endpoint_input_types=kwtypes(deployment_name=str),
    endpoint_config={
        "EndpointName": "{inputs.deployment_name}-{idempotence_token}",
        "EndpointConfigName": "{inputs.deployment_name}-{inputs.idempotence_token}",
    },
)


@workflow
def wf(
    execution_role_arn: str = "...",
    deployment_name: str = "stable-diffusion-model",
    instance_type: str = "ml.g5.2xlarge",
    initial_instance_count: int = 1,
    region: str = "us-east-2",
) -> list[dict]:
    return sd_deployment(...)

If the deployment already exists, here's how the result might look: [{'result': 'Entity already exists: arn:aws:sagemaker:us-east-2:123456789:model/stable-diffusion-model-0b2634d258f999f6'}, {'result': 'Entity already exists: arn:aws:sagemaker:us-east-2:123456789:endpoint-config/stable-diffusion-model-bc3afab2046ae0df'}, {'EndpointArn': 'arn:aws:sagemaker:us-east-2:123456789:endpoint/stable-diffusion-model-3cfe9fbf5941f3c2'}]

If only some entities exist, such as the model and endpoint configuration, the plugin detects that these two entities exist and only creates an endpoint to complete the deployment.

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[cosmicBboy]

I think a complementary approach is to create a task SagemakerEndpointExists so that users can compose a workflow themselves with the conditional:

@workflow
def wf():
    (
        if_(SagemakerEndpointExists(name="my-model")
        .then_(create_endpoint(..., override=True))  # delete then create
        .else(create_endpoint(...))  # create
    )

I guess once conditionals in imperative workflows is implemented this pattern wouldn't be necessary

[samhita-alla]

I think a complementary approach is to create a task SagemakerEndpointExists so that users can compose a workflow themselves with the conditional:

@cosmicBboy i think we can create a workflow to handle this case as we need to use the list API to check if the endpoint's available. i'll work on it.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 76.26%. Comparing base (a6a8651) to head (376906a).
Report is 15 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2400      +/-   ##
==========================================
- Coverage   76.27%   76.26%   -0.02%     
==========================================
  Files         183      183              
  Lines       18716    18728      +12     
  Branches     3694     3695       +1     
==========================================
+ Hits        14275    14282       +7     
- Misses       3804     3811       +7     
+ Partials      637      635       -2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.