wehub-resource-sync
1091 lines
40 KiB
Python
1091 lines
40 KiB
Python
import json
|
|
import time
|
|
from datetime import datetime
|
|
from typing import Any, NamedTuple
|
|
|
|
from moto.core import DEFAULT_ACCOUNT_ID, BackendDict, BaseBackend, BaseModel
|
|
from moto.core.models import base_decorator
|
|
from moto.core.responses import BaseResponse
|
|
|
|
|
|
class SageMakerResourceWithArn(NamedTuple):
|
|
resource: Any
|
|
arn: str
|
|
|
|
|
|
class SageMakerResponse(BaseResponse):
|
|
"""
|
|
A collection of handlers for SageMaker API calls that produce API-conforming
|
|
JSON responses.
|
|
"""
|
|
|
|
@property
|
|
def sagemaker_backend(self):
|
|
return sagemaker_backends[DEFAULT_ACCOUNT_ID][self.region]
|
|
|
|
@property
|
|
def request_params(self):
|
|
return json.loads(self.body)
|
|
|
|
def create_endpoint_config(self):
|
|
"""
|
|
Handler for the SageMaker "CreateEndpointConfig" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateEndpointConfig.html.
|
|
"""
|
|
config_name = self.request_params["EndpointConfigName"]
|
|
production_variants = self.request_params.get("ProductionVariants")
|
|
tags = self.request_params.get("Tags", [])
|
|
async_inference_config = self.request_params.get("AsyncInferenceConfig")
|
|
new_config = self.sagemaker_backend.create_endpoint_config(
|
|
config_name=config_name,
|
|
production_variants=production_variants,
|
|
tags=tags,
|
|
region_name=self.region,
|
|
async_inference_config=async_inference_config,
|
|
)
|
|
return json.dumps({"EndpointConfigArn": new_config.arn})
|
|
|
|
def describe_endpoint_config(self):
|
|
"""
|
|
Handler for the SageMaker "DescribeEndpoint" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeEndpoint.html.
|
|
"""
|
|
config_name = self.request_params["EndpointConfigName"]
|
|
config_description = self.sagemaker_backend.describe_endpoint_config(config_name)
|
|
return json.dumps(config_description.response_object)
|
|
|
|
def delete_endpoint_config(self):
|
|
"""
|
|
Handler for the SageMaker "DeleteEndpointConfig" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DeleteEndpointConfig.html.
|
|
"""
|
|
config_name = self.request_params["EndpointConfigName"]
|
|
self.sagemaker_backend.delete_endpoint_config(config_name)
|
|
return ""
|
|
|
|
def create_endpoint(self):
|
|
"""
|
|
Handler for the SageMaker "CreateEndpoint" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateEndpoint.html.
|
|
"""
|
|
endpoint_name = self.request_params["EndpointName"]
|
|
endpoint_config_name = self.request_params["EndpointConfigName"]
|
|
tags = self.request_params.get("Tags", [])
|
|
new_endpoint = self.sagemaker_backend.create_endpoint(
|
|
endpoint_name=endpoint_name,
|
|
endpoint_config_name=endpoint_config_name,
|
|
tags=tags,
|
|
region_name=self.region,
|
|
)
|
|
return json.dumps({"EndpointArn": new_endpoint.arn})
|
|
|
|
def describe_endpoint(self):
|
|
"""
|
|
Handler for the SageMaker "DescribeEndpoint" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeEndpoint.html.
|
|
"""
|
|
endpoint_name = self.request_params["EndpointName"]
|
|
endpoint_description = self.sagemaker_backend.describe_endpoint(endpoint_name)
|
|
return json.dumps(endpoint_description.response_object)
|
|
|
|
def update_endpoint(self):
|
|
"""
|
|
Handler for the SageMaker "UpdateEndpoint" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_UpdateEndpoint.html.
|
|
"""
|
|
endpoint_name = self.request_params["EndpointName"]
|
|
new_config_name = self.request_params["EndpointConfigName"]
|
|
updated_endpoint = self.sagemaker_backend.update_endpoint(
|
|
endpoint_name=endpoint_name, new_config_name=new_config_name
|
|
)
|
|
return json.dumps({"EndpointArn": updated_endpoint.arn})
|
|
|
|
def delete_endpoint(self):
|
|
"""
|
|
Handler for the SageMaker "DeleteEndpoint" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DeleteEndpoint.html.
|
|
"""
|
|
endpoint_name = self.request_params["EndpointName"]
|
|
self.sagemaker_backend.delete_endpoint(endpoint_name)
|
|
return ""
|
|
|
|
def list_endpoints(self):
|
|
"""
|
|
Handler for the SageMaker "ListEndpoints" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListEndpoints.html.
|
|
|
|
This function does not support pagination. All endpoint configs are returned in a
|
|
single response.
|
|
"""
|
|
endpoint_summaries = self.sagemaker_backend.list_endpoints()
|
|
return json.dumps({
|
|
"Endpoints": [summary.response_object for summary in endpoint_summaries]
|
|
})
|
|
|
|
def list_endpoint_configs(self):
|
|
"""
|
|
Handler for the SageMaker "ListEndpointConfigs" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListEndpointConfigs.html.
|
|
|
|
This function does not support pagination. All endpoint configs are returned in a
|
|
single response.
|
|
"""
|
|
# Note:
|
|
endpoint_config_summaries = self.sagemaker_backend.list_endpoint_configs()
|
|
return json.dumps({
|
|
"EndpointConfigs": [summary.response_object for summary in endpoint_config_summaries]
|
|
})
|
|
|
|
def list_models(self):
|
|
"""
|
|
Handler for the SageMaker "ListModels" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListModels.html.
|
|
|
|
This function does not support pagination. All endpoint configs are returned in a
|
|
single response.
|
|
"""
|
|
model_summaries = self.sagemaker_backend.list_models()
|
|
return json.dumps({"Models": [summary.response_object for summary in model_summaries]})
|
|
|
|
def create_model(self):
|
|
"""
|
|
Handler for the SageMaker "CreateModel" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateModel.html.
|
|
"""
|
|
model_name = self.request_params["ModelName"]
|
|
primary_container = self.request_params["PrimaryContainer"]
|
|
execution_role_arn = self.request_params["ExecutionRoleArn"]
|
|
tags = self.request_params.get("Tags", [])
|
|
vpc_config = self.request_params.get("VpcConfig", None)
|
|
new_model = self.sagemaker_backend.create_model(
|
|
model_name=model_name,
|
|
primary_container=primary_container,
|
|
execution_role_arn=execution_role_arn,
|
|
tags=tags,
|
|
vpc_config=vpc_config,
|
|
region_name=self.region,
|
|
)
|
|
return json.dumps({"ModelArn": new_model.arn})
|
|
|
|
def describe_model(self):
|
|
"""
|
|
Handler for the SageMaker "DescribeModel" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeModel.html.
|
|
"""
|
|
model_name = self.request_params["ModelName"]
|
|
model_description = self.sagemaker_backend.describe_model(model_name)
|
|
return json.dumps(model_description.response_object)
|
|
|
|
def delete_model(self):
|
|
"""
|
|
Handler for the SageMaker "DeleteModel" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DeleteModel.html.
|
|
"""
|
|
model_name = self.request_params["ModelName"]
|
|
self.sagemaker_backend.delete_model(model_name)
|
|
return ""
|
|
|
|
def list_tags(self):
|
|
"""
|
|
Handler for the SageMaker "ListTags" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ListTags.html
|
|
"""
|
|
arn = self.request_params["ResourceArn"]
|
|
sagemaker_resource = (
|
|
"models" if "model" in arn else "endpoints" if "endpoint" in arn else None
|
|
)
|
|
results = self.sagemaker_backend.list_tags(
|
|
resource_arn=arn, region_name=self.region, resource_type=sagemaker_resource
|
|
)
|
|
|
|
return json.dumps({"Tags": results, "NextToken": None})
|
|
|
|
def create_transform_job(self):
|
|
"""
|
|
Handler for the SageMaker "CreateTransformJob" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTransformJob.html.
|
|
"""
|
|
job_name = self.request_params["TransformJobName"]
|
|
model_name = self.request_params.get("ModelName")
|
|
transform_input = self.request_params.get("TransformInput")
|
|
transform_output = self.request_params.get("TransformOutput")
|
|
transform_resources = self.request_params.get("TransformResources")
|
|
data_processing = self.request_params.get("DataProcessing")
|
|
tags = self.request_params.get("Tags", [])
|
|
new_job = self.sagemaker_backend.create_transform_job(
|
|
job_name=job_name,
|
|
model_name=model_name,
|
|
transform_input=transform_input,
|
|
transform_output=transform_output,
|
|
transform_resources=transform_resources,
|
|
data_processing=data_processing,
|
|
tags=tags,
|
|
region_name=self.region,
|
|
)
|
|
return json.dumps({"TransformJobArn": new_job.arn})
|
|
|
|
def stop_transform_job(self):
|
|
"""
|
|
Handler for the SageMaker "StopTransformJob" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_StopTransformJob.html.
|
|
"""
|
|
job_name = self.request_params["TransformJobName"]
|
|
self.sagemaker_backend.stop_transform_job(job_name)
|
|
return ""
|
|
|
|
def describe_transform_job(self):
|
|
"""
|
|
Handler for the SageMaker "DescribeTransformJob" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DescribeTransformJob.html.
|
|
"""
|
|
job_name = self.request_params["TransformJobName"]
|
|
transform_job_description = self.sagemaker_backend.describe_transform_job(job_name)
|
|
return json.dumps(transform_job_description.response_object)
|
|
|
|
def list_transform_jobs(self):
|
|
"""
|
|
Handler for the SageMaker "ListTransformJobs" API call documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ListTransformJobs.html.
|
|
|
|
This function does not support pagination. All transform jobs are returned in a
|
|
single response.
|
|
"""
|
|
transform_job_summaries = self.sagemaker_backend.list_transform_jobs()
|
|
return json.dumps({
|
|
"TransformJobSummaries": [
|
|
summary.response_object for summary in transform_job_summaries
|
|
]
|
|
})
|
|
|
|
|
|
class SageMakerBackend(BaseBackend):
|
|
"""
|
|
A mock backend for managing and exposing SageMaker resource state.
|
|
"""
|
|
|
|
BASE_SAGEMAKER_ARN = "arn:aws:sagemaker:{region_name}:{account_id}:"
|
|
|
|
def __init__(self, region_name, account_id=None):
|
|
super().__init__(region_name, account_id)
|
|
self.models = {}
|
|
self.endpoints = {}
|
|
self.endpoint_configs = {}
|
|
self.transform_jobs = {}
|
|
self._endpoint_update_latency_seconds = 0
|
|
self._transform_job_update_latency_seconds = 0
|
|
|
|
def set_endpoint_update_latency(self, latency_seconds):
|
|
"""
|
|
Sets the latency for the following operations that update endpoint state:
|
|
- "create_endpoint"
|
|
- "update_endpoint"
|
|
"""
|
|
self._endpoint_update_latency_seconds = latency_seconds
|
|
|
|
def set_transform_job_update_latency(self, latency_seconds):
|
|
"""
|
|
Sets the latency for the following operations that update transform job state:
|
|
- "create_transform_job"
|
|
- "terminate_transform_job"
|
|
"""
|
|
self._transform_job_update_latency_seconds = latency_seconds
|
|
|
|
def set_endpoint_latest_operation(self, endpoint_name, operation):
|
|
if endpoint_name not in self.endpoints:
|
|
raise ValueError(
|
|
"Attempted to manually set the latest operation for an endpoint"
|
|
" that does not exist!"
|
|
)
|
|
self.endpoints[endpoint_name].resource.latest_operation = operation
|
|
|
|
def set_transform_job_latest_operation(self, transform_job_name, operation):
|
|
if transform_job_name not in self.transform_jobs:
|
|
raise ValueError(
|
|
"Attempted to manually set the latest operation for a transform job"
|
|
" that does not exist!"
|
|
)
|
|
self.transform_jobs[transform_job_name].resource.latest_operation = operation
|
|
|
|
@property
|
|
def _url_module(self):
|
|
"""
|
|
Required override from the Moto "BaseBackend" object that reroutes requests from the
|
|
specified SageMaker URLs to the mocked SageMaker backend.
|
|
"""
|
|
urls_module_name = "tests.sagemaker.mock.mock_sagemaker_urls"
|
|
return __import__(urls_module_name, fromlist=["url_bases", "url_paths"])
|
|
|
|
def _get_base_arn(self, region_name):
|
|
"""
|
|
Returns:
|
|
A SageMaker ARN prefix that can be prepended to a resource name.
|
|
"""
|
|
return SageMakerBackend.BASE_SAGEMAKER_ARN.format(
|
|
region_name=region_name, account_id=DEFAULT_ACCOUNT_ID
|
|
)
|
|
|
|
def create_endpoint_config(
|
|
self, config_name, production_variants, tags, region_name, async_inference_config
|
|
):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "CreateEndpointConfig" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateEndpointConfig.html.
|
|
"""
|
|
if config_name in self.endpoint_configs:
|
|
raise ValueError(
|
|
"Attempted to create an endpoint configuration with name:"
|
|
f" {config_name}, but an endpoint configuration with this"
|
|
" name already exists."
|
|
)
|
|
for production_variant in production_variants:
|
|
if "ModelName" not in production_variant:
|
|
raise ValueError("Production variant must specify a model name.")
|
|
elif production_variant["ModelName"] not in self.models:
|
|
raise ValueError(
|
|
"Production variant specifies a model name that does not exist"
|
|
" Model name: '{model_name}'".format(model_name=production_variant["ModelName"])
|
|
)
|
|
|
|
new_config = EndpointConfig(
|
|
config_name=config_name,
|
|
production_variants=production_variants,
|
|
tags=tags,
|
|
async_inference_config=async_inference_config,
|
|
)
|
|
new_config_arn = self._get_base_arn(region_name=region_name) + new_config.arn_descriptor
|
|
new_resource = SageMakerResourceWithArn(resource=new_config, arn=new_config_arn)
|
|
self.endpoint_configs[config_name] = new_resource
|
|
return new_resource
|
|
|
|
def describe_endpoint_config(self, config_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "DescribeEndpointConfig" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeEndpointConfig.html.
|
|
"""
|
|
if config_name not in self.endpoint_configs:
|
|
raise ValueError(
|
|
f"Attempted to describe an endpoint config with name: `{config_name}`"
|
|
" that does not exist."
|
|
)
|
|
|
|
config = self.endpoint_configs[config_name]
|
|
return EndpointConfigDescription(config=config.resource, arn=config.arn)
|
|
|
|
def delete_endpoint_config(self, config_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "DeleteEndpointConfig" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DeleteEndpointConfig.html.
|
|
"""
|
|
if config_name not in self.endpoint_configs:
|
|
raise ValueError(
|
|
f"Attempted to delete an endpoint config with name: `{config_name}`"
|
|
" that does not exist."
|
|
)
|
|
|
|
del self.endpoint_configs[config_name]
|
|
|
|
def create_endpoint(self, endpoint_name, endpoint_config_name, tags, region_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "CreateEndpoint" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateEndpoint.html.
|
|
"""
|
|
if endpoint_name in self.endpoints:
|
|
raise ValueError(
|
|
f"Attempted to create an endpoint with name: `{endpoint_name}`"
|
|
" but an endpoint with this name already exists."
|
|
)
|
|
|
|
if endpoint_config_name not in self.endpoint_configs:
|
|
raise ValueError(
|
|
"Attempted to create an endpoint with a configuration named:"
|
|
f" `{endpoint_config_name}` However, this configuration does not exist."
|
|
)
|
|
|
|
new_endpoint = Endpoint(
|
|
endpoint_name=endpoint_name,
|
|
config_name=endpoint_config_name,
|
|
tags=tags,
|
|
latest_operation=EndpointOperation.create_successful(
|
|
latency_seconds=self._endpoint_update_latency_seconds
|
|
),
|
|
)
|
|
new_endpoint_arn = self._get_base_arn(region_name=region_name) + new_endpoint.arn_descriptor
|
|
new_resource = SageMakerResourceWithArn(resource=new_endpoint, arn=new_endpoint_arn)
|
|
self.endpoints[endpoint_name] = new_resource
|
|
return new_resource
|
|
|
|
def describe_endpoint(self, endpoint_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "DescribeEndpoint" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeEndpoint.html.
|
|
"""
|
|
if endpoint_name not in self.endpoints:
|
|
raise ValueError(
|
|
f"Attempted to describe an endpoint with name: `{endpoint_name}`"
|
|
" that does not exist."
|
|
)
|
|
|
|
endpoint = self.endpoints[endpoint_name]
|
|
config = self.endpoint_configs[endpoint.resource.config_name]
|
|
return EndpointDescription(
|
|
endpoint=endpoint.resource, config=config.resource, arn=endpoint.arn
|
|
)
|
|
|
|
def update_endpoint(self, endpoint_name, new_config_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "UpdateEndpoint" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_UpdateEndpoint.html.
|
|
"""
|
|
if endpoint_name not in self.endpoints:
|
|
raise ValueError(
|
|
f"Attempted to update an endpoint with name: `{endpoint_name}` that does not exist."
|
|
)
|
|
|
|
if new_config_name not in self.endpoint_configs:
|
|
raise ValueError(
|
|
f"Attempted to update an endpoint named `{endpoint_name}` with a new"
|
|
f" configuration named: `{new_config_name}`. However, this configuration"
|
|
" does not exist."
|
|
)
|
|
|
|
endpoint = self.endpoints[endpoint_name]
|
|
endpoint.resource.latest_operation = EndpointOperation.update_successful(
|
|
latency_seconds=self._endpoint_update_latency_seconds
|
|
)
|
|
endpoint.resource.config_name = new_config_name
|
|
return endpoint
|
|
|
|
def delete_endpoint(self, endpoint_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "DeleteEndpoint" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DeleteEndpoint.html.
|
|
"""
|
|
if endpoint_name not in self.endpoints:
|
|
raise ValueError(
|
|
f"Attempted to delete an endpoint with name: `{endpoint_name}` that does not exist."
|
|
)
|
|
|
|
del self.endpoints[endpoint_name]
|
|
|
|
def list_endpoints(self):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "ListEndpoints" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListEndpoints.html.
|
|
"""
|
|
return [
|
|
EndpointSummary(endpoint=endpoint.resource, arn=endpoint.arn)
|
|
for endpoint in self.endpoints.values()
|
|
]
|
|
|
|
def list_endpoint_configs(self):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "ListEndpointConfigs" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListEndpointConfigs.html.
|
|
"""
|
|
return [
|
|
EndpointConfigSummary(config=endpoint_config.resource, arn=endpoint_config.arn)
|
|
for endpoint_config in self.endpoint_configs.values()
|
|
]
|
|
|
|
def list_models(self):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "ListModels" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListModels.html.
|
|
"""
|
|
return [ModelSummary(model=model.resource, arn=model.arn) for model in self.models.values()]
|
|
|
|
def list_tags(self, resource_arn, region_name, resource_type):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "ListTags" API
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ListTags.html
|
|
"""
|
|
resource_values = getattr(self, resource_type).values()
|
|
for sagemaker_resource in resource_values:
|
|
if sagemaker_resource.arn == resource_arn:
|
|
return sagemaker_resource.resource.tags
|
|
|
|
def create_model(
|
|
self, model_name, primary_container, execution_role_arn, tags, region_name, vpc_config=None
|
|
):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "CreateModel" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateModel.html.
|
|
"""
|
|
if model_name in self.models:
|
|
raise ValueError(
|
|
f"Attempted to create a model with name: `{model_name}`"
|
|
" but a model with this name already exists."
|
|
)
|
|
|
|
new_model = Model(
|
|
model_name=model_name,
|
|
primary_container=primary_container,
|
|
execution_role_arn=execution_role_arn,
|
|
tags=tags,
|
|
vpc_config=vpc_config,
|
|
)
|
|
new_model_arn = self._get_base_arn(region_name=region_name) + new_model.arn_descriptor
|
|
new_resource = SageMakerResourceWithArn(resource=new_model, arn=new_model_arn)
|
|
self.models[model_name] = new_resource
|
|
return new_resource
|
|
|
|
def describe_model(self, model_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "DescribeModel" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeModel.html.
|
|
"""
|
|
if model_name not in self.models:
|
|
raise ValueError(
|
|
f"Attempted to describe a model with name: `{model_name}` that does not exist."
|
|
)
|
|
|
|
model = self.models[model_name]
|
|
return ModelDescription(model=model.resource, arn=model.arn)
|
|
|
|
def delete_model(self, model_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "DeleteModel" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DeleteModel.html.
|
|
"""
|
|
if model_name not in self.models:
|
|
raise ValueError(
|
|
f"Attempted to delete an model with name: `{model_name}` that does not exist."
|
|
)
|
|
|
|
del self.models[model_name]
|
|
|
|
def create_transform_job(
|
|
self,
|
|
job_name,
|
|
model_name,
|
|
transform_input,
|
|
transform_output,
|
|
transform_resources,
|
|
data_processing,
|
|
tags,
|
|
region_name,
|
|
):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "CreateTransformJob" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTransformJob.html.
|
|
"""
|
|
if job_name in self.transform_jobs:
|
|
raise ValueError(
|
|
"Attempted to create a transform job with name:"
|
|
f" {job_name}, but a transform job with this"
|
|
" name already exists."
|
|
)
|
|
|
|
if model_name not in self.models:
|
|
raise ValueError(
|
|
"Attempted to create a transform job with a model named:"
|
|
f" `{model_name}` However, this model does not exist."
|
|
)
|
|
|
|
new_job = TransformJob(
|
|
job_name=job_name,
|
|
model_name=model_name,
|
|
transform_input=transform_input,
|
|
transform_output=transform_output,
|
|
transform_resources=transform_resources,
|
|
data_processing=data_processing,
|
|
tags=tags,
|
|
latest_operation=TransformJobOperation.create_successful(
|
|
latency_seconds=self._transform_job_update_latency_seconds
|
|
),
|
|
)
|
|
new_job_arn = self._get_base_arn(region_name=region_name) + new_job.arn_descriptor
|
|
new_resource = SageMakerResourceWithArn(resource=new_job, arn=new_job_arn)
|
|
self.transform_jobs[job_name] = new_resource
|
|
return new_resource
|
|
|
|
def describe_transform_job(self, job_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "DescribeTransformJob" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DescribeTransformJob.html.
|
|
"""
|
|
if job_name not in self.transform_jobs:
|
|
raise ValueError(
|
|
f"Attempted to describe a transform job with name: `{job_name}`"
|
|
" that does not exist."
|
|
)
|
|
|
|
transform_job = self.transform_jobs[job_name]
|
|
return TransformJobDescription(transform_job=transform_job.resource, arn=transform_job.arn)
|
|
|
|
def stop_transform_job(self, job_name):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "StopTransformJob" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_StopTransformJob.html.
|
|
"""
|
|
if job_name not in self.transform_jobs:
|
|
raise ValueError(
|
|
f"Attempted to stop a transform job with name: `{job_name}` that does not exist."
|
|
)
|
|
|
|
self.transform_jobs[
|
|
job_name
|
|
].resource.latest_operation = TransformJobOperation.stop_successful(
|
|
latency_seconds=self._transform_job_update_latency_seconds
|
|
)
|
|
|
|
def list_transform_jobs(self):
|
|
"""
|
|
Modifies backend state during calls to the SageMaker "ListTransformJobs" API
|
|
documented here:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ListTransformJobs.html.
|
|
"""
|
|
return [
|
|
TransformJobSummary(transform_job=transform_job.resource, arn=transform_job.arn)
|
|
for transform_job in self.transform_jobs.values()
|
|
]
|
|
|
|
|
|
class TimestampedResource(BaseModel):
|
|
TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
|
|
|
|
def __init__(self):
|
|
curr_time = datetime.now().strftime(TimestampedResource.TIMESTAMP_FORMAT)
|
|
self.creation_time = curr_time
|
|
self.last_modified_time = curr_time
|
|
|
|
|
|
class Endpoint(TimestampedResource):
|
|
"""
|
|
Object representing a SageMaker endpoint. The SageMakerBackend will create
|
|
and manage Endpoints.
|
|
"""
|
|
|
|
STATUS_IN_SERVICE = "InService"
|
|
STATUS_FAILED = "Failed"
|
|
STATUS_CREATING = "Creating"
|
|
STATUS_UPDATING = "Updating"
|
|
|
|
def __init__(self, endpoint_name, config_name, tags, latest_operation):
|
|
"""
|
|
Args:
|
|
endpoint_name: The name of the Endpoint.
|
|
config_name: The name of the EndpointConfiguration to associate with the Endpoint.
|
|
tags: Arbitrary tags to associate with the endpoint.
|
|
latest_operation: The most recent operation that was invoked on the endpoint,
|
|
represented as an EndpointOperation object.
|
|
"""
|
|
super().__init__()
|
|
self.endpoint_name = endpoint_name
|
|
self.config_name = config_name
|
|
self.tags = tags
|
|
self.latest_operation = latest_operation
|
|
|
|
@property
|
|
def arn_descriptor(self):
|
|
return f":endpoint/{self.endpoint_name}"
|
|
|
|
@property
|
|
def status(self):
|
|
return self.latest_operation.status()
|
|
|
|
|
|
class TransformJob(TimestampedResource):
|
|
"""
|
|
Object representing a SageMaker transform job. The SageMakerBackend will create
|
|
and manage transform jobs.
|
|
"""
|
|
|
|
STATUS_IN_PROGRESS = "InProgress"
|
|
STATUS_FAILED = "Failed"
|
|
STATUS_COMPLETED = "Completed"
|
|
STATUS_STOPPING = "Stopping"
|
|
STATUS_STOPPED = "Stopped"
|
|
|
|
def __init__(
|
|
self,
|
|
job_name,
|
|
model_name,
|
|
transform_input,
|
|
transform_output,
|
|
transform_resources,
|
|
data_processing,
|
|
tags,
|
|
latest_operation,
|
|
):
|
|
"""
|
|
Args:
|
|
job_name: The name of the TransformJob.
|
|
model_name: The name of the model to associate with the TransformJob.
|
|
transform_input: The input data source and the way transform job consumes it.
|
|
transform_output: The output results of the transform job.
|
|
transform_resources: The ML instance types and instance count to use for the
|
|
transform job.
|
|
data_processing: The data structure to specify the inference data and associate data
|
|
to the prediction results.
|
|
tags: Arbitrary tags to associate with the transform job.
|
|
latest_operation: The most recent operation that was invoked on the transform job,
|
|
represented as an TransformJobOperation object.
|
|
"""
|
|
super().__init__()
|
|
self.job_name = job_name
|
|
self.model_name = model_name
|
|
self.transform_input = transform_input
|
|
self.transform_output = transform_output
|
|
self.transform_resources = transform_resources
|
|
self.data_processing = data_processing
|
|
self.tags = tags
|
|
self.latest_operation = latest_operation
|
|
|
|
@property
|
|
def arn_descriptor(self):
|
|
return f":transform-job/{self.job_name}"
|
|
|
|
@property
|
|
def status(self):
|
|
return self.latest_operation.status()
|
|
|
|
|
|
class EndpointOperation:
|
|
"""
|
|
Object representing a SageMaker endpoint operation ("create" or "update"). Every
|
|
Endpoint is associated with the operation that was most recently invoked on it.
|
|
"""
|
|
|
|
def __init__(self, latency_seconds, pending_status, completed_status):
|
|
"""
|
|
Args:
|
|
latency_seconds: The latency of the operation, in seconds. Before the time window
|
|
specified by this latency elapses, the operation will have the status specified by
|
|
``pending_status``. After the time window elapses, the operation will
|
|
have the status specified by ``completed_status``.
|
|
pending_status: The status that the operation should reflect *before* the latency
|
|
window has elapsed.
|
|
completed_status: The status that the operation should reflect *after* the latency
|
|
window has elapsed.
|
|
"""
|
|
self.latency_seconds = latency_seconds
|
|
self.pending_status = pending_status
|
|
self.completed_status = completed_status
|
|
self.start_time = time.time()
|
|
|
|
def status(self):
|
|
if time.time() - self.start_time < self.latency_seconds:
|
|
return self.pending_status
|
|
else:
|
|
return self.completed_status
|
|
|
|
@classmethod
|
|
def create_successful(cls, latency_seconds):
|
|
return cls(
|
|
latency_seconds=latency_seconds,
|
|
pending_status=Endpoint.STATUS_CREATING,
|
|
completed_status=Endpoint.STATUS_IN_SERVICE,
|
|
)
|
|
|
|
@classmethod
|
|
def create_unsuccessful(cls, latency_seconds):
|
|
return cls(
|
|
latency_seconds=latency_seconds,
|
|
pending_status=Endpoint.STATUS_CREATING,
|
|
completed_status=Endpoint.STATUS_FAILED,
|
|
)
|
|
|
|
@classmethod
|
|
def update_successful(cls, latency_seconds):
|
|
return cls(
|
|
latency_seconds=latency_seconds,
|
|
pending_status=Endpoint.STATUS_UPDATING,
|
|
completed_status=Endpoint.STATUS_IN_SERVICE,
|
|
)
|
|
|
|
@classmethod
|
|
def update_unsuccessful(cls, latency_seconds):
|
|
return cls(
|
|
latency_seconds=latency_seconds,
|
|
pending_status=Endpoint.STATUS_UPDATING,
|
|
completed_status=Endpoint.STATUS_FAILED,
|
|
)
|
|
|
|
|
|
class TransformJobOperation:
|
|
"""
|
|
Object representing a SageMaker transform job operation ("create" or "stop"). Every
|
|
transform job is associated with the operation that was most recently invoked on it.
|
|
"""
|
|
|
|
def __init__(self, latency_seconds, pending_status, completed_status):
|
|
"""
|
|
Args:
|
|
latency_seconds: The latency of the operation, in seconds. Before the time window
|
|
specified by this latency elapses, the operation will have the status
|
|
specified by ``pending_status``. After the time window elapses, the
|
|
operation will have the status specified by ``completed_status``.
|
|
pending_status: The status that the operation should reflect *before* the latency
|
|
window has elapsed.
|
|
completed_status: The status that the operation should reflect *after* the latency
|
|
window has elapsed.
|
|
"""
|
|
self.latency_seconds = latency_seconds
|
|
self.pending_status = pending_status
|
|
self.completed_status = completed_status
|
|
self.start_time = time.time()
|
|
|
|
def status(self):
|
|
if time.time() - self.start_time < self.latency_seconds:
|
|
return self.pending_status
|
|
else:
|
|
return self.completed_status
|
|
|
|
@classmethod
|
|
def create_successful(cls, latency_seconds):
|
|
return cls(
|
|
latency_seconds=latency_seconds,
|
|
pending_status=TransformJob.STATUS_IN_PROGRESS,
|
|
completed_status=TransformJob.STATUS_COMPLETED,
|
|
)
|
|
|
|
@classmethod
|
|
def create_unsuccessful(cls, latency_seconds):
|
|
return cls(
|
|
latency_seconds=latency_seconds,
|
|
pending_status=TransformJob.STATUS_IN_PROGRESS,
|
|
completed_status=TransformJob.STATUS_FAILED,
|
|
)
|
|
|
|
@classmethod
|
|
def stop_successful(cls, latency_seconds):
|
|
return cls(
|
|
latency_seconds=latency_seconds,
|
|
pending_status=TransformJob.STATUS_STOPPING,
|
|
completed_status=TransformJob.STATUS_STOPPED,
|
|
)
|
|
|
|
|
|
class EndpointSummary:
|
|
"""
|
|
Object representing an endpoint entry in the endpoints list returned by
|
|
SageMaker's "ListEndpoints" API:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListEndpoints.html.
|
|
"""
|
|
|
|
def __init__(self, endpoint, arn):
|
|
self.endpoint = endpoint
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"EndpointName": self.endpoint.endpoint_name,
|
|
"CreationTime": self.endpoint.creation_time,
|
|
"LastModifiedTime": self.endpoint.last_modified_time,
|
|
"EndpointStatus": self.endpoint.status,
|
|
"EndpointArn": self.arn,
|
|
}
|
|
|
|
|
|
class EndpointDescription:
|
|
"""
|
|
Object representing an endpoint description returned by SageMaker's
|
|
"DescribeEndpoint" API:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeEndpoint.html.
|
|
"""
|
|
|
|
def __init__(self, endpoint, config, arn):
|
|
self.endpoint = endpoint
|
|
self.config = config
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"EndpointName": self.endpoint.endpoint_name,
|
|
"EndpointArn": self.arn,
|
|
"EndpointConfigName": self.endpoint.config_name,
|
|
"ProductionVariants": self.config.production_variants,
|
|
"EndpointStatus": self.endpoint.status,
|
|
"CreationTime": self.endpoint.creation_time,
|
|
"LastModifiedTime": self.endpoint.last_modified_time,
|
|
}
|
|
|
|
|
|
class EndpointConfig(TimestampedResource):
|
|
"""
|
|
Object representing a SageMaker endpoint configuration. The SageMakerBackend will create
|
|
and manage EndpointConfigs.
|
|
"""
|
|
|
|
def __init__(self, config_name, production_variants, tags, async_inference_config=None):
|
|
super().__init__()
|
|
self.config_name = config_name
|
|
self.production_variants = production_variants
|
|
self.tags = tags
|
|
self.async_inference_config = async_inference_config
|
|
|
|
@property
|
|
def arn_descriptor(self):
|
|
return f":endpoint-config/{self.config_name}"
|
|
|
|
|
|
class EndpointConfigSummary:
|
|
"""
|
|
Object representing an endpoint configuration entry in the configurations list returned by
|
|
SageMaker's "ListEndpointConfigs" API:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListEndpointConfigs.html.
|
|
"""
|
|
|
|
def __init__(self, config, arn):
|
|
self.config = config
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"EndpointConfigName": self.config.config_name,
|
|
"EndpointArn": self.arn,
|
|
"CreationTime": self.config.creation_time,
|
|
}
|
|
|
|
|
|
class EndpointConfigDescription:
|
|
"""
|
|
Object representing an endpoint configuration description returned by SageMaker's
|
|
"DescribeEndpointConfig" API:
|
|
https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeEndpointConfig.html.
|
|
"""
|
|
|
|
def __init__(self, config, arn):
|
|
self.config = config
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"EndpointConfigName": self.config.config_name,
|
|
"EndpointConfigArn": self.arn,
|
|
"ProductionVariants": self.config.production_variants,
|
|
"CreationTime": self.config.creation_time,
|
|
"AsyncInferenceConfig": self.config.async_inference_config,
|
|
}
|
|
|
|
|
|
class Model(TimestampedResource):
|
|
"""
|
|
Object representing a SageMaker model. The SageMakerBackend will create and manage Models.
|
|
"""
|
|
|
|
def __init__(self, model_name, primary_container, execution_role_arn, tags, vpc_config):
|
|
super().__init__()
|
|
self.model_name = model_name
|
|
self.primary_container = primary_container
|
|
self.execution_role_arn = execution_role_arn
|
|
self.tags = tags
|
|
self.vpc_config = vpc_config
|
|
|
|
@property
|
|
def arn_descriptor(self):
|
|
return f":model/{self.model_name}"
|
|
|
|
|
|
class ModelSummary:
|
|
"""
|
|
Object representing a model entry in the models list returned by SageMaker's
|
|
"ListModels" API: https://docs.aws.amazon.com/sagemaker/latest/dg/API_ListModels.html.
|
|
"""
|
|
|
|
def __init__(self, model, arn):
|
|
self.model = model
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"ModelArn": self.arn,
|
|
"ModelName": self.model.model_name,
|
|
"CreationTime": self.model.creation_time,
|
|
}
|
|
|
|
|
|
class ModelDescription:
|
|
"""
|
|
Object representing a model description returned by SageMaker's
|
|
"DescribeModel" API: https://docs.aws.amazon.com/sagemaker/latest/dg/API_DescribeModel.html.
|
|
"""
|
|
|
|
def __init__(self, model, arn):
|
|
self.model = model
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"ModelArn": self.arn,
|
|
"ModelName": self.model.model_name,
|
|
"PrimaryContainer": self.model.primary_container,
|
|
"ExecutionRoleArn": self.model.execution_role_arn,
|
|
"VpcConfig": self.model.vpc_config or {},
|
|
"CreationTime": self.model.creation_time,
|
|
}
|
|
|
|
|
|
class TransformJobSummary:
|
|
"""
|
|
Object representing a TransformJobSummary entry in the TransformJobSummaries list returned by
|
|
SageMaker's "ListTransformJobs" API:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ListTransformJobs.html.
|
|
"""
|
|
|
|
def __init__(self, transform_job, arn):
|
|
self.transform_job = transform_job
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"TransformJobName": self.transform_job.job_name,
|
|
"TransformJobArn": self.arn,
|
|
"CreationTime": self.transform_job.creation_time,
|
|
"LastModifiedTime": self.transform_job.last_modified_time,
|
|
"TransformJobStatus": self.transform_job.status,
|
|
}
|
|
|
|
|
|
class TransformJobDescription:
|
|
"""
|
|
Object representing a transform job description returned by SageMaker's
|
|
"DescribeTransformJob" API:
|
|
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DescribeTransformJob.html.
|
|
"""
|
|
|
|
def __init__(self, transform_job, arn):
|
|
self.transform_job = transform_job
|
|
self.arn = arn
|
|
|
|
@property
|
|
def response_object(self):
|
|
return {
|
|
"TransformJobName": self.transform_job.job_name,
|
|
"TransformJobArn": self.arn,
|
|
"CreationTime": self.transform_job.creation_time,
|
|
"LastModifiedTime": self.transform_job.last_modified_time,
|
|
"TransformJobStatus": self.transform_job.status,
|
|
"ModelName": self.transform_job.model_name,
|
|
}
|
|
|
|
|
|
# Create a SageMaker backend for EC2 region: "us-west-2"
|
|
sagemaker_backends = BackendDict(SageMakerBackend, "sagemaker")
|
|
|
|
mock_sagemaker = base_decorator(sagemaker_backends)
|