def sdk_model_package(name_suffix):
model_package_group_name = name_suffix + "-group"
(
model_package_group_input,
model_package_group_response,
) = sdk_make_model_package_group(model_package_group_name)
model_package_input, model_package_response = sdk_make_model_package(
model_package_group_name)
yield (
model_package_group_input,
model_package_group_response,
model_package_input,
model_package_response,
)
model_package_arn = model_package_response.get("ModelPackageArn")
if get_sagemaker_model_package(model_package_arn) is not None:
wait_sagemaker_model_package_status(model_package_arn,
cfg.JOB_STATUS_COMPLETED)
sagemaker_client().delete_model_package(
ModelPackageName=model_package_arn)
if get_sagemaker_model_package_group(model_package_group_name) is not None:
sagemaker_client().delete_model_package_group(
ModelPackageGroupName=model_package_group_name)
def sdk_make_model_package(model_package_group_name):
data_bucket = REPLACEMENT_VALUES["SAGEMAKER_DATA_BUCKET"]
model_package_input = {
"ModelPackageGroupName": model_package_group_name,
"InferenceSpecification": {
"Containers": [{
"Image":
REPLACEMENT_VALUES["XGBOOST_IMAGE_URI"],
"ModelDataUrl":
f"s3://{data_bucket}/sagemaker/model/xgboost-mnist-model.tar.gz",
}],
"SupportedContentTypes": [
"text/csv",
],
"SupportedResponseMIMETypes": [
"text/csv",
],
},
}
model_package_response = sagemaker_client().create_model_package(
**model_package_input)
assert model_package_response.get("ModelPackageArn", None) is not None
return model_package_input, model_package_response
def test_versioned_model_package_completed(
self, xgboost_versioned_model_package):
(reference, spec, resource) = xgboost_versioned_model_package
assert k8s.get_resource_exists(reference)
model_package_group_name = resource["spec"].get(
"modelPackageGroupName")
# Model package name for Versioned Model packages is the ARN of the resource
model_package_name = sagemaker_client().list_model_packages(
ModelPackageGroupName=model_package_group_name
)["ModelPackageSummaryList"][0]["ModelPackageArn"]
model_package_desc = get_sagemaker_model_package(model_package_name)
if k8s.get_resource_arn(resource) is None:
logging.error(
f"ARN for this resource is None, resource status is: {resource['status']}"
)
assert k8s.get_resource_arn(resource) == model_package_name
self._assert_model_package_status_in_sync(model_package_name,
reference,
cfg.JOB_STATUS_INPROGRESS)
assert k8s.wait_on_condition(reference, "ACK.ResourceSynced", "False")
self._assert_model_package_status_in_sync(model_package_name,
reference,
cfg.JOB_STATUS_COMPLETED)
assert k8s.wait_on_condition(reference, "ACK.ResourceSynced", "True")
# Update the resource
new_model_approval_status = "Approved"
approval_description = "Approved modelpackage"
spec["spec"]["modelApprovalStatus"] = new_model_approval_status
spec["spec"]["approvalDescription"] = approval_description
resource = k8s.patch_custom_resource(reference, spec)
resource = k8s.wait_resource_consumed_by_controller(reference)
assert resource is not None
self._assert_model_package_status_in_sync(model_package_name,
reference,
cfg.JOB_STATUS_COMPLETED)
assert k8s.wait_on_condition(reference, "ACK.ResourceSynced", "True")
model_package_desc = get_sagemaker_model_package(model_package_name)
assert model_package_desc[
"ModelApprovalStatus"] == new_model_approval_status
assert model_package_desc[
"ApprovalDescription"] == approval_description
assert (resource["spec"].get("modelApprovalStatus",
None) == new_model_approval_status)
assert resource["spec"].get("approvalDescription",
None) == approval_description
# Check that you can delete a completed resource from k8s
_, deleted = k8s.delete_custom_resource(reference, DELETE_WAIT_PERIOD,
DELETE_WAIT_LENGTH)
assert deleted is True
assert get_sagemaker_model_package(model_package_name) is None
def sdk_make_endpoint(endpoint_name, endpoint_config_name):
endpoint_input = {
"EndpointName": endpoint_name,
"EndpointConfigName": endpoint_config_name,
}
endpoint_response = sagemaker_client().create_endpoint(**endpoint_input)
assert endpoint_response.get("EndpointArn", None) is not None
return endpoint_input, endpoint_response
def get_sagemaker_hpo_job(hpo_job_name: str):
try:
hpo_desc = sagemaker_client().describe_hyper_parameter_tuning_job(
HyperParameterTuningJobName=hpo_job_name)
return hpo_desc
except botocore.exceptions.ClientError as error:
logging.error(
f"SageMaker could not find an hpo job with the name {hpo_job_name}. Error {error}"
)
return None
def get_sagemaker_processing_job(processing_job_name: str):
try:
processing_job = sagemaker_client().describe_processing_job(
ProcessingJobName=processing_job_name)
return processing_job
except botocore.exceptions.ClientError as error:
logging.error(
f"SageMaker could not find a processing job with the name {processing_job_name}. Error {error}"
)
return None
def get_notebook_instance_lifecycle_config(notebook_instance_lfc_name: str):
try:
resp = sagemaker_client().describe_notebook_instance_lifecycle_config(
NotebookInstanceLifecycleConfigName=notebook_instance_lfc_name)
return resp
except botocore.exceptions.ClientError as error:
logging.error(
f"SageMaker could not find a Notebook Instance Lifecycle Configuration with the name {notebook_instance_lfc_name}. Error {error}"
)
return None
def get_sagemaker_feature_group(feature_group_name: str):
"""Used to check if there is an existing feature group with a given feature_group_name."""
try:
return sagemaker_client().describe_feature_group(
FeatureGroupName=feature_group_name)
except botocore.exceptions.ClientError as error:
logging.error(
f"SageMaker could not find a feature group with the name {feature_group_name}. Error {error}"
)
return None
def sdk_make_model_package_group(model_package_group_name):
model_package_group_input = {
"ModelPackageGroupName": model_package_group_name,
"ModelPackageGroupDescription": "ModelPackageGroup for adoption",
}
model_package_group_response = sagemaker_client(
).create_model_package_group(**model_package_group_input)
assert model_package_group_response.get("ModelPackageGroupArn",
None) is not None
return model_package_group_input, model_package_group_response
def get_sagemaker_transform_job(transform_job_name: str):
try:
transform_desc = sagemaker_client().describe_transform_job(
TransformJobName=transform_job_name
)
return transform_desc
except botocore.exceptions.ClientError as error:
logging.error(
f"SageMaker could not find a transform job with the name {transform_job_name}. Error {error}"
)
return None
def sdk_endpoint(name_suffix):
model_name = name_suffix + "-model"
endpoint_config_name = name_suffix + "-config"
endpoint_name = name_suffix
model_input, model_response = sdk_make_model(model_name)
endpoint_config_input, endpoint_config_response = sdk_make_endpoint_config(
model_name, endpoint_config_name
)
endpoint_input, endpoint_response = sdk_make_endpoint(
endpoint_name, endpoint_config_name
)
yield (
model_input,
model_response,
endpoint_config_input,
endpoint_config_response,
endpoint_input,
endpoint_response,
)
if get_sagemaker_endpoint(endpoint_name) is not None:
wait_sagemaker_endpoint_status(endpoint_name, cfg.ENDPOINT_STATUS_INSERVICE)
sagemaker_client().delete_endpoint(EndpointName=endpoint_name)
if get_sagemaker_endpoint_config(endpoint_config_name) is not None:
sagemaker_client().delete_endpoint_config(
EndpointConfigName=endpoint_config_name
)
if get_sagemaker_model(model_name) is not None:
sagemaker_client().delete_model(ModelName=model_name)
def test_create_cross_region_model(self, cross_region_model):
(reference, resource) = cross_region_model
assert k8s.get_resource_exists(reference)
sm_client = sagemaker_client(get_cross_region())
model_name = resource["spec"].get("modelName", None)
model_desc = get_sagemaker_model(model_name, sm_client)
cross_region_model_arn = model_desc["ModelArn"]
assert k8s.get_resource_arn(resource) == cross_region_model_arn
# Delete the k8s resource.
_, deleted = k8s.delete_custom_resource(reference, 3, 10)
assert deleted
assert get_sagemaker_model(model_name, sm_client) is None
def sdk_make_model(model_name):
data_bucket = REPLACEMENT_VALUES["SAGEMAKER_DATA_BUCKET"]
model_input = {
"ModelName": model_name,
"Containers": [
{
"Image": REPLACEMENT_VALUES["XGBOOST_IMAGE_URI"],
"ModelDataUrl": f"s3://{data_bucket}/sagemaker/model/xgboost-mnist-model.tar.gz",
}
],
"ExecutionRoleArn": REPLACEMENT_VALUES["SAGEMAKER_EXECUTION_ROLE_ARN"],
}
model_response = sagemaker_client().create_model(**model_input)
assert model_response.get("ModelArn", None) is not None
return model_input, model_response
def sdk_make_endpoint_config(model_name, endpoint_config_name):
endpoint_config_input = {
"EndpointConfigName": endpoint_config_name,
"ProductionVariants": [
{
"VariantName": "variant-1",
"ModelName": model_name,
"InitialInstanceCount": 1,
"InstanceType": REPLACEMENT_VALUES["ENDPOINT_INSTANCE_TYPE"],
}
],
}
endpoint_config_response = sagemaker_client().create_endpoint_config(
**endpoint_config_input
)
assert endpoint_config_response.get("EndpointConfigArn", None) is not None
return endpoint_config_input, endpoint_config_response