truncate sagemaker agent outputs and automate idempotence token handling

[samhita-alla]

Tracking issue

Why are the changes needed?

• To return the relevant outputs and discard the irrelevant entries from the output dictionary.
• To make specifying the idempotence token optional in the config and delegate the heavy lifting to the plugin.

What changes were proposed in this pull request?

This PR truncates the SageMaker agent outputs to only return the ARNs. It also adds a new idempotence_token parameter to the create_sagemaker_deployment function, which, when set to True, appends an idempotence token to the relevant fields.

How was this patch tested?

I've run the example locally. New output format: [{'ModelArn': 'arn:aws:sagemaker:us-east-2:123456789:model/stable-diffusion-model-0b2634d258f999f6'}, {'EndpointConfigArn': 'arn:aws:sagemaker:us-east-2:123456789:endpoint-config/stable-diffusion-model-5825e66eb1014642'}, {'EndpointArn': 'arn:aws:sagemaker:us-east-2:123456789:endpoint/stable-diffusion-model-1-dc4144fef7e49561'}]

New deployment workflow:

sd_deployment = create_sagemaker_deployment(
    name="stable-diffusion",
    model_input_types=kwtypes(
        deployment_name=str, model_path=FlyteFile, execution_role_arn=str
    ),
    model_config={
        "ModelName": "{inputs.deployment_name}",
        "PrimaryContainer": {
            "Image": "{images.sd_deployment_image}",
            "ModelDataUrl": "{inputs.model_path}",
            "Environment": {
                "SAGEMAKER_TRITON_DEFAULT_MODEL_NAME": "pipeline",
                "SAGEMAKER_TRITON_LOG_INFO": "false --load-model=text_encoder --load-model=vae",
            },
        },
        "ExecutionRoleArn": "{inputs.execution_role_arn}",
    },
    endpoint_config_input_types=kwtypes(
        deployment_name=str,
        initial_instance_count=int,
        instance_type=str,
    ),
    endpoint_config_config={
        "ProductionVariants": [
            {
                "VariantName": "AllTraffic-2",
                "ModelName": "{inputs.deployment_name}",
                "InitialInstanceCount": "{inputs.initial_instance_count}",
                "InstanceType": "{inputs.instance_type}",
            },
        ],
    },
    endpoint_input_types=kwtypes(deployment_name=str),
    endpoint_config={
        "EndpointName": "{inputs.deployment_name}",
    },
    images={"sd_deployment_image": triton_image_uri(version="23.12")},
    region_at_runtime=True,
)

The idempotence_token is set to True by default, so the agent appends an idempotence token to model name, endpoint config name, and endpoint.

If the field value isn't provided, such as ModelName, the agent appends the idempotence token to the workflow name and uses that as the ModelName.

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

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 72.22%. Comparing base (097e9e8) to head (67e0594).
Report is 6 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2588      +/-   ##
==========================================
- Coverage   76.22%   72.22%   -4.00%     
==========================================
  Files         187      185       -2     
  Lines       18938    18832     -106     
  Branches     3706     3705       -1     
==========================================
- Hits        14435    13601     -834     
- Misses       3870     4568     +698     
- Partials      633      663      +30     

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