wehub-resource-sync
527 lines
19 KiB
Python
527 lines
19 KiB
Python
import os
|
|
import time
|
|
from functools import wraps
|
|
from typing import NamedTuple
|
|
from unittest import mock
|
|
|
|
import boto3
|
|
import botocore
|
|
import numpy as np
|
|
import pytest
|
|
from click.testing import CliRunner
|
|
from moto.core import DEFAULT_ACCOUNT_ID
|
|
from sklearn.linear_model import LogisticRegression
|
|
|
|
import mlflow
|
|
import mlflow.pyfunc
|
|
import mlflow.sagemaker as mfs
|
|
import mlflow.sagemaker.cli as mfscli
|
|
import mlflow.sklearn
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import Model
|
|
from mlflow.protos.databricks_pb2 import (
|
|
INTERNAL_ERROR,
|
|
INVALID_PARAMETER_VALUE,
|
|
RESOURCE_DOES_NOT_EXIST,
|
|
ErrorCode,
|
|
)
|
|
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
|
|
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
|
|
|
|
from tests.helper_functions import set_boto_credentials # noqa: F401
|
|
from tests.sagemaker.mock import TransformJob, TransformJobOperation, mock_sagemaker
|
|
|
|
|
|
class TrainedModel(NamedTuple):
|
|
model_path: str
|
|
run_id: str
|
|
model_uri: str
|
|
|
|
|
|
@pytest.fixture
|
|
def pretrained_model():
|
|
model_path = "model"
|
|
with mlflow.start_run():
|
|
X = np.array([-2, -1, 0, 1, 2, 1]).reshape(-1, 1)
|
|
y = np.array([0, 0, 1, 1, 1, 0])
|
|
lr = LogisticRegression(solver="lbfgs")
|
|
lr.fit(X, y)
|
|
mlflow.sklearn.log_model(lr, name=model_path)
|
|
run_id = mlflow.active_run().info.run_id
|
|
model_uri = "runs:/" + run_id + "/" + model_path
|
|
return TrainedModel(model_path, run_id, model_uri)
|
|
|
|
|
|
@pytest.fixture
|
|
def sagemaker_client():
|
|
return boto3.client("sagemaker", region_name="us-west-2")
|
|
|
|
|
|
def get_sagemaker_backend(region_name):
|
|
return mock_sagemaker.backends[DEFAULT_ACCOUNT_ID][region_name]
|
|
|
|
|
|
def mock_sagemaker_aws_services(fn):
|
|
from moto import mock_ecr, mock_iam, mock_s3, mock_sts
|
|
|
|
@mock_ecr
|
|
@mock_iam
|
|
@mock_s3
|
|
@mock_sagemaker
|
|
@mock_sts
|
|
@wraps(fn)
|
|
def mock_wrapper(*args, **kwargs):
|
|
# Create an ECR repository for the `mlflow-pyfunc` SageMaker docker image
|
|
ecr_client = boto3.client("ecr", region_name="us-west-2")
|
|
ecr_client.create_repository(repositoryName=mfs.DEFAULT_IMAGE_NAME)
|
|
|
|
# Create the moto IAM role
|
|
role_policy = """
|
|
{
|
|
"Version": "2012-10-17",
|
|
"Statement": [
|
|
{
|
|
"Effect": "Allow",
|
|
"Action": "*",
|
|
"Resource": "*"
|
|
}
|
|
]
|
|
}
|
|
"""
|
|
iam_client = boto3.client("iam", region_name="us-west-2")
|
|
iam_client.create_role(RoleName="moto", AssumeRolePolicyDocument=role_policy)
|
|
|
|
return fn(*args, **kwargs)
|
|
|
|
return mock_wrapper
|
|
|
|
|
|
def test_batch_deployment_with_unsupported_flavor_raises_exception(pretrained_model):
|
|
unsupported_flavor = "this is not a valid flavor"
|
|
match = "The specified flavor: `this is not a valid flavor` is not supported for deployment"
|
|
with pytest.raises(MlflowException, match=match) as exc:
|
|
mfs.deploy_transform_job(
|
|
job_name="bad_flavor",
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
flavor=unsupported_flavor,
|
|
)
|
|
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
def test_batch_deployment_of_model_with_no_supported_flavors_raises_exception(pretrained_model):
|
|
logged_model_path = _download_artifact_from_uri(pretrained_model.model_uri)
|
|
model_config_path = os.path.join(logged_model_path, "MLmodel")
|
|
model_config = Model.load(model_config_path)
|
|
del model_config.flavors[mlflow.pyfunc.FLAVOR_NAME]
|
|
model_config.save(path=model_config_path)
|
|
|
|
match = "The specified model does not contain any of the supported flavors for deployment"
|
|
with pytest.raises(MlflowException, match=match) as exc:
|
|
mfs.deploy_transform_job(
|
|
job_name="missing-flavor",
|
|
model_uri=logged_model_path,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
flavor=None,
|
|
)
|
|
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_deploy_sagemaker_transform_job_in_asynchronous_mode_without_archiving_throws_exception(
|
|
pretrained_model,
|
|
):
|
|
with pytest.raises(MlflowException, match="Resources must be archived") as exc:
|
|
mfs.deploy_transform_job(
|
|
job_name="test-job",
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
archive=False,
|
|
synchronous=False,
|
|
)
|
|
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploy_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_local(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
job_name = "test-job"
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
archive=True,
|
|
)
|
|
|
|
region_name = sagemaker_client.meta.region_name
|
|
s3_client = boto3.client("s3", region_name=region_name)
|
|
default_bucket = mfs._get_default_s3_bucket(region_name)
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
model_name = transform_job_description["ModelName"]
|
|
assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]
|
|
object_names = [
|
|
entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
|
|
]
|
|
assert any(model_name in object_name for object_name in object_names)
|
|
assert job_name in [
|
|
transform_job["TransformJobName"]
|
|
for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
|
|
]
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploy_cli_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_local(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
job_name = "test-job"
|
|
result = CliRunner(env={"LC_ALL": "en_US.UTF-8", "LANG": "en_US.UTF-8"}).invoke(
|
|
mfscli.commands,
|
|
[
|
|
"deploy-transform-job",
|
|
"--job-name",
|
|
job_name,
|
|
"--model-uri",
|
|
pretrained_model.model_uri,
|
|
"--input-data-type",
|
|
"Some Data Type",
|
|
"--input-uri",
|
|
"Some Input Uri",
|
|
"--content-type",
|
|
"Some Content Type",
|
|
"--output-path",
|
|
"Some Output Path",
|
|
"--archive",
|
|
],
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
region_name = sagemaker_client.meta.region_name
|
|
s3_client = boto3.client("s3", region_name=region_name)
|
|
default_bucket = mfs._get_default_s3_bucket(region_name)
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
model_name = transform_job_description["ModelName"]
|
|
assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]
|
|
object_names = [
|
|
entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
|
|
]
|
|
assert any(model_name in object_name for object_name in object_names)
|
|
assert job_name in [
|
|
transform_job["TransformJobName"]
|
|
for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
|
|
]
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploy_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_s3(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
local_model_path = _download_artifact_from_uri(pretrained_model.model_uri)
|
|
artifact_path = "model"
|
|
region_name = sagemaker_client.meta.region_name
|
|
default_bucket = mfs._get_default_s3_bucket(region_name)
|
|
s3_artifact_repo = S3ArtifactRepository(f"s3://{default_bucket}")
|
|
s3_artifact_repo.log_artifacts(local_model_path, artifact_path=artifact_path)
|
|
model_s3_uri = f"s3://{default_bucket}/{pretrained_model.model_path}"
|
|
|
|
job_name = "test-job"
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=model_s3_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
archive=True,
|
|
)
|
|
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
model_name = transform_job_description["ModelName"]
|
|
assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]
|
|
|
|
s3_client = boto3.client("s3", region_name=region_name)
|
|
object_names = [
|
|
entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
|
|
]
|
|
assert any(model_name in object_name for object_name in object_names)
|
|
assert job_name in [
|
|
transform_job["TransformJobName"]
|
|
for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
|
|
]
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploy_cli_creates_sagemaker_transform_job_and_s3_resources_with_expected_names_from_s3(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
local_model_path = _download_artifact_from_uri(pretrained_model.model_uri)
|
|
artifact_path = "model"
|
|
region_name = sagemaker_client.meta.region_name
|
|
default_bucket = mfs._get_default_s3_bucket(region_name)
|
|
s3_artifact_repo = S3ArtifactRepository(f"s3://{default_bucket}")
|
|
s3_artifact_repo.log_artifacts(local_model_path, artifact_path=artifact_path)
|
|
model_s3_uri = f"s3://{default_bucket}/{pretrained_model.model_path}"
|
|
|
|
job_name = "test-job"
|
|
result = CliRunner(env={"LC_ALL": "en_US.UTF-8", "LANG": "en_US.UTF-8"}).invoke(
|
|
mfscli.commands,
|
|
[
|
|
"deploy-transform-job",
|
|
"--job-name",
|
|
job_name,
|
|
"--model-uri",
|
|
model_s3_uri,
|
|
"--input-data-type",
|
|
"Some Data Type",
|
|
"--input-uri",
|
|
"Some Input Uri",
|
|
"--content-type",
|
|
"Some Content Type",
|
|
"--output-path",
|
|
"Some Output Path",
|
|
"--archive",
|
|
],
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
region_name = sagemaker_client.meta.region_name
|
|
s3_client = boto3.client("s3", region_name=region_name)
|
|
default_bucket = mfs._get_default_s3_bucket(region_name)
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
model_name = transform_job_description["ModelName"]
|
|
assert model_name in [model["ModelName"] for model in sagemaker_client.list_models()["Models"]]
|
|
object_names = [
|
|
entry["Key"] for entry in s3_client.list_objects(Bucket=default_bucket)["Contents"]
|
|
]
|
|
assert any(model_name in object_name for object_name in object_names)
|
|
assert job_name in [
|
|
transform_job["TransformJobName"]
|
|
for transform_job in sagemaker_client.list_transform_jobs()["TransformJobSummaries"]
|
|
]
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploying_sagemaker_transform_job_with_preexisting_name_in_create_mode_throws_exception(
|
|
pretrained_model,
|
|
):
|
|
job_name = "test-job"
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match="a batch transform job with the same name already exists"
|
|
) as exc:
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
)
|
|
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploy_in_synchronous_mode_waits_for_transform_job_creation_to_complete_before_returning(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
transform_job_creation_latency = 10
|
|
get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
|
|
transform_job_creation_latency
|
|
)
|
|
|
|
job_name = "test-job"
|
|
deployment_start_time = time.time()
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
synchronous=True,
|
|
)
|
|
deployment_end_time = time.time()
|
|
|
|
assert (deployment_end_time - deployment_start_time) >= transform_job_creation_latency
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_COMPLETED
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploy_create_in_asynchronous_mode_returns_before_transform_job_creation_completes(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
transform_job_creation_latency = 10
|
|
get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
|
|
transform_job_creation_latency
|
|
)
|
|
|
|
job_name = "test-job"
|
|
deployment_start_time = time.time()
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
archive=True,
|
|
synchronous=False,
|
|
)
|
|
deployment_end_time = time.time()
|
|
|
|
assert (deployment_end_time - deployment_start_time) < transform_job_creation_latency
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_IN_PROGRESS
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_deploy_in_throw_exception_after_transform_job_creation_fails(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
transform_job_creation_latency = 10
|
|
sagemaker_backend = get_sagemaker_backend(sagemaker_client.meta.region_name)
|
|
sagemaker_backend.set_transform_job_update_latency(transform_job_creation_latency)
|
|
|
|
boto_caller = botocore.client.BaseClient._make_api_call
|
|
|
|
def fail_transform_job_creations(self, operation_name, operation_kwargs):
|
|
"""
|
|
Processes all boto3 client operations according to the following rules:
|
|
- If the operation is a transform job creation, create the transform job and
|
|
set its status to ``TransformJob.STATUS_FAILED``.
|
|
- Else, execute the client operation as normal
|
|
"""
|
|
result = boto_caller(self, operation_name, operation_kwargs)
|
|
if operation_name == "CreateTransformJob":
|
|
transform_job_name = operation_kwargs["TransformJobName"]
|
|
sagemaker_backend.set_transform_job_latest_operation(
|
|
transform_job_name=transform_job_name,
|
|
operation=TransformJobOperation.create_unsuccessful(
|
|
latency_seconds=transform_job_creation_latency
|
|
),
|
|
)
|
|
return result
|
|
|
|
with (
|
|
mock.patch("botocore.client.BaseClient._make_api_call", new=fail_transform_job_creations),
|
|
pytest.raises(MlflowException, match="batch transform job failed") as exc,
|
|
):
|
|
mfs.deploy_transform_job(
|
|
job_name="test-job",
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
)
|
|
|
|
assert exc.value.error_code == ErrorCode.Name(INTERNAL_ERROR)
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_attempting_to_terminate_in_asynchronous_mode_without_archiving_throws_exception(
|
|
pretrained_model,
|
|
):
|
|
job_name = "test-job"
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="Resources must be archived") as exc:
|
|
mfs.terminate_transform_job(
|
|
job_name=job_name,
|
|
archive=False,
|
|
synchronous=False,
|
|
)
|
|
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_terminate_in_sync_mode_waits_for_transform_job_termination_to_complete_before_returning(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
transform_job_termination_latency = 10
|
|
get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
|
|
transform_job_termination_latency
|
|
)
|
|
|
|
job_name = "test-job"
|
|
termination_start_time = time.time()
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
archive=True,
|
|
synchronous=True,
|
|
)
|
|
|
|
mfs.terminate_transform_job(job_name=job_name, synchronous=True)
|
|
termination_end_time = time.time()
|
|
|
|
assert (termination_end_time - termination_start_time) >= transform_job_termination_latency
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_STOPPED
|
|
|
|
|
|
@mock_sagemaker_aws_services
|
|
def test_terminate_in_asynchronous_mode_returns_before_transform_job_termination_completes(
|
|
pretrained_model, sagemaker_client
|
|
):
|
|
transform_job_termination_latency = 10
|
|
get_sagemaker_backend(sagemaker_client.meta.region_name).set_transform_job_update_latency(
|
|
transform_job_termination_latency
|
|
)
|
|
|
|
job_name = "test-job"
|
|
termination_start_time = time.time()
|
|
mfs.deploy_transform_job(
|
|
job_name=job_name,
|
|
model_uri=pretrained_model.model_uri,
|
|
s3_input_data_type="Some Data Type",
|
|
s3_input_uri="Some Input Uri",
|
|
content_type="Some Content Type",
|
|
s3_output_path="Some Output Path",
|
|
archive=True,
|
|
synchronous=False,
|
|
)
|
|
|
|
mfs.terminate_transform_job(job_name=job_name, archive=True, synchronous=False)
|
|
termination_end_time = time.time()
|
|
|
|
assert (termination_end_time - termination_start_time) < transform_job_termination_latency
|
|
transform_job_description = sagemaker_client.describe_transform_job(TransformJobName=job_name)
|
|
assert transform_job_description["TransformJobStatus"] == TransformJob.STATUS_STOPPING
|