def test_repack_model_step_with_source_dir(estimator, source_dir):
model_data = f"s3://{BUCKET}/model.tar.gz"
entry_point = "inference.py"
step = _RepackModelStep(
name="MyRepackModelStep",
sagemaker_session=estimator.sagemaker_session,
role=estimator.role,
model_data=model_data,
entry_point=entry_point,
source_dir=source_dir,
)
request_dict = step.to_request()
assert os.path.isfile(f"{source_dir}/_repack_model.py")
hyperparameters = request_dict["Arguments"]["HyperParameters"]
assert hyperparameters["inference_script"] == '"inference.py"'
assert hyperparameters["model_archive"] == '"model.tar.gz"'
assert hyperparameters["sagemaker_program"] == '"_repack_model.py"'
del request_dict["Arguments"]["HyperParameters"]
del request_dict["Arguments"]["AlgorithmSpecification"]["TrainingImage"]
assert request_dict == {
"Name": "MyRepackModelStep",
"Type": "Training",
"Arguments": {
"AlgorithmSpecification": {
"TrainingInputMode": "File"
},
"DebugHookConfig": {
"CollectionConfigurations": [],
"S3OutputPath": "s3://my-bucket/"
},
"InputDataConfig": [{
"ChannelName": "training",
"DataSource": {
"S3DataSource": {
"S3DataDistributionType": "FullyReplicated",
"S3DataType": "S3Prefix",
"S3Uri": f"s3://{BUCKET}",
}
},
}],
"OutputDataConfig": {
"S3OutputPath": f"s3://{BUCKET}/"
},
"ResourceConfig": {
"InstanceCount": 1,
"InstanceType": "ml.m5.large",
"VolumeSizeInGB": 30,
},
"RoleArn":
ROLE,
"StoppingCondition": {
"MaxRuntimeInSeconds": 86400
},
},
}
assert step.properties.TrainingJobName.expr == {
"Get": "Steps.MyRepackModelStep.TrainingJobName"
}
def test_inject_repack_script_s3(estimator, tmp, fake_s3):
create_file_tree(
tmp,
[
"model-dir/aa",
"model-dir/foo/inference.py",
],
)
model_data = Properties(path="Steps.MyStep", shape_name="DescribeModelOutput")
entry_point = "inference.py"
source_dir_path = "s3://fake/location"
step = _RepackModelStep(
name="MyRepackModelStep",
sagemaker_session=fake_s3.sagemaker_session,
role=estimator.role,
image_uri="foo",
model_data=model_data,
entry_point=entry_point,
source_dir=source_dir_path,
)
fake_s3.tar_and_upload("model-dir", "s3://fake/location")
step._inject_repack_script()
assert list_tar_files(fake_s3.fake_upload_path, tmp) == {
"/aa",
"/foo/inference.py",
"/_repack_model.py",
}
def test_repack_model_step(estimator):
model_data = f"s3://{BUCKET}/model.tar.gz"
entry_point = f"{DATA_DIR}/dummy_script.py"
step = _RepackModelStep(
name="MyRepackModelStep",
sagemaker_session=estimator.sagemaker_session,
role=estimator.role,
model_data=model_data,
entry_point=entry_point,
depends_on=["TestStep"],
)
request_dict = step.to_request()
hyperparameters = request_dict["Arguments"]["HyperParameters"]
assert hyperparameters["inference_script"] == '"dummy_script.py"'
assert hyperparameters["model_archive"] == '"s3://my-bucket/model.tar.gz"'
assert hyperparameters["sagemaker_program"] == '"_repack_model.py"'
assert (
hyperparameters["sagemaker_submit_directory"]
== '"s3://my-bucket/MyRepackModelStep-1be10316814854973ed1b445db3ef84e/source/sourcedir.tar.gz"'
)
del request_dict["Arguments"]["HyperParameters"]
del request_dict["Arguments"]["AlgorithmSpecification"]["TrainingImage"]
assert request_dict == {
"Name": "MyRepackModelStep",
"Type": "Training",
"DependsOn": ["TestStep"],
"Arguments": {
"AlgorithmSpecification": {"TrainingInputMode": "File"},
"DebugHookConfig": {"CollectionConfigurations": [], "S3OutputPath": "s3://my-bucket/"},
"InputDataConfig": [
{
"ChannelName": "training",
"DataSource": {
"S3DataSource": {
"S3DataDistributionType": "FullyReplicated",
"S3DataType": "S3Prefix",
"S3Uri": f"s3://{BUCKET}/model.tar.gz",
}
},
}
],
"OutputDataConfig": {"S3OutputPath": f"s3://{BUCKET}/"},
"ResourceConfig": {
"InstanceCount": 1,
"InstanceType": "ml.m5.large",
"VolumeSizeInGB": 30,
},
"RoleArn": ROLE,
"StoppingCondition": {"MaxRuntimeInSeconds": 86400},
},
}
assert step.properties.TrainingJobName.expr == {
"Get": "Steps.MyRepackModelStep.TrainingJobName"
}
def _append_repack_model_step(self):
"""Create and append a `_RepackModelStep` for the runtime repack"""
if isinstance(self._model, PipelineModel):
model_list = self._model.models
elif isinstance(self._model, Model):
model_list = [self._model]
else:
logging.warning("No models to repack")
return
security_group_ids = None
subnets = None
if self._model.vpc_config:
security_group_ids = self._model.vpc_config.get("SecurityGroupIds", None)
subnets = self._model.vpc_config.get("Subnets", None)
for i, model in enumerate(model_list):
runtime_repack_flg = (
self._need_runtime_repack and id(model) in self._need_runtime_repack
)
if runtime_repack_flg:
name_base = model.name or i
repack_model_step = _RepackModelStep(
name="{}-{}-{}".format(self.name, _REPACK_MODEL_NAME_BASE, name_base),
sagemaker_session=self._model.sagemaker_session or model.sagemaker_session,
role=self._model.role or model.role,
model_data=model.model_data,
entry_point=model.entry_point,
source_dir=model.source_dir,
dependencies=model.dependencies,
subnets=subnets,
security_group_ids=security_group_ids,
description=(
"Used to repack a model with customer scripts for a "
"register/create model step"
),
depends_on=self.depends_on,
retry_policies=self._repack_model_retry_policies,
)
self.steps.append(repack_model_step)
repacked_model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
if self._create_model_args:
if isinstance(self._model, PipelineModel):
container = self.step_args.create_model_request["Containers"][i]
else:
container = self.step_args.create_model_request["PrimaryContainer"]
else:
container = self.step_args.create_model_package_request[
"InferenceSpecification"
]["Containers"][i]
container["ModelDataUrl"] = repacked_model_data
def __init__(
self,
name: str,
estimator: EstimatorBase,
model_data,
content_types,
response_types,
inference_instances,
transform_instances,
depends_on: List[str] = None,
model_package_group_name=None,
model_metrics=None,
approval_status=None,
image_uri=None,
compile_model_family=None,
description=None,
**kwargs,
):
"""Construct steps `_RepackModelStep` and `_RegisterModelStep` based on the estimator.
Args:
name (str): The name of the training step.
estimator: The estimator instance.
model_data: The S3 uri to the model data from training.
content_types (list): The supported MIME types for the input data (default: None).
response_types (list): The supported MIME types for the output data (default: None).
inference_instances (list): A list of the instance types that are used to
generate inferences in real-time (default: None).
transform_instances (list): A list of the instance types on which a transformation
job can be run or on which an endpoint can be deployed (default: None).
depends_on (List[str]): The list of step names the first step in the collection
depends on
model_package_group_name (str): The Model Package Group name, exclusive to
`model_package_name`, using `model_package_group_name` makes the Model Package
versioned (default: None).
model_metrics (ModelMetrics): ModelMetrics object (default: None).
approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
or "PendingManualApproval" (default: "PendingManualApproval").
image_uri (str): The container image uri for Model Package, if not specified,
Estimator's training container image is used (default: None).
compile_model_family (str): The instance family for the compiled model. If
specified, a compiled model is used (default: None).
description (str): Model Package description (default: None).
**kwargs: additional arguments to `create_model`.
"""
steps: List[Step] = []
repack_model = False
if "entry_point" in kwargs:
repack_model = True
entry_point = kwargs["entry_point"]
source_dir = kwargs.get("source_dir")
dependencies = kwargs.get("dependencies")
repack_model_step = _RepackModelStep(
name=f"{name}RepackModel",
depends_on=depends_on,
estimator=estimator,
model_data=model_data,
entry_point=entry_point,
source_dir=source_dir,
dependencies=dependencies,
)
steps.append(repack_model_step)
model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
# remove kwargs consumed by model repacking step
kwargs.pop("entry_point", None)
kwargs.pop("source_dir", None)
kwargs.pop("dependencies", None)
register_model_step = _RegisterModelStep(
name=name,
estimator=estimator,
model_data=model_data,
content_types=content_types,
response_types=response_types,
inference_instances=inference_instances,
transform_instances=transform_instances,
model_package_group_name=model_package_group_name,
model_metrics=model_metrics,
approval_status=approval_status,
image_uri=image_uri,
compile_model_family=compile_model_family,
description=description,
**kwargs,
)
if not repack_model:
register_model_step.add_depends_on(depends_on)
steps.append(register_model_step)
self.steps = steps
def __init__(
self,
name: str,
estimator: EstimatorBase,
model_data,
model_inputs,
instance_count,
instance_type,
transform_inputs,
# model arguments
image_uri=None,
predictor_cls=None,
env=None,
# transformer arguments
strategy=None,
assemble_with=None,
output_path=None,
output_kms_key=None,
accept=None,
max_concurrent_transforms=None,
max_payload=None,
tags=None,
volume_kms_key=None,
depends_on: List[str] = None,
**kwargs,
):
"""Construct steps required for a Transformer step collection:
An estimator-centric step collection. It models what happens in workflows
when invoking the `transform()` method on an estimator instance:
First, if custom
model artifacts are required, a `_RepackModelStep` is included.
Second, a
`CreateModelStep` with the model data passed in from a training step or other
training job output.
Finally, a `TransformerStep`.
If repacking
the model artifacts is not necessary, only the CreateModelStep and TransformerStep
are in the step collection.
Args:
name (str): The name of the Transform Step.
estimator: The estimator instance.
instance_count (int): The number of EC2 instances to use.
instance_type (str): The type of EC2 instance to use.
strategy (str): The strategy used to decide how to batch records in
a single request (default: None). Valid values: 'MultiRecord'
and 'SingleRecord'.
assemble_with (str): How the output is assembled (default: None).
Valid values: 'Line' or 'None'.
output_path (str): The S3 location for saving the transform result. If
not specified, results are stored to a default bucket.
output_kms_key (str): Optional. A KMS key ID for encrypting the
transform output (default: None).
accept (str): The accept header passed by the client to
the inference endpoint. If it is supported by the endpoint,
it will be the format of the batch transform output.
env (dict): The Environment variables to be set for use during the
transform job (default: None).
depends_on (List[str]): The list of step names the first step in
the collection depends on
"""
steps = []
if "entry_point" in kwargs:
entry_point = kwargs["entry_point"]
source_dir = kwargs.get("source_dir")
dependencies = kwargs.get("dependencies")
repack_model_step = _RepackModelStep(
name=f"{name}RepackModel",
depends_on=depends_on,
estimator=estimator,
model_data=model_data,
entry_point=entry_point,
source_dir=source_dir,
dependencies=dependencies,
)
steps.append(repack_model_step)
model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
def predict_wrapper(endpoint, session):
return Predictor(endpoint, session)
predictor_cls = predictor_cls or predict_wrapper
model = Model(
image_uri=image_uri or estimator.training_image_uri(),
model_data=model_data,
predictor_cls=predictor_cls,
vpc_config=None,
sagemaker_session=estimator.sagemaker_session,
role=estimator.role,
**kwargs,
)
model_step = CreateModelStep(
name=f"{name}CreateModelStep",
model=model,
inputs=model_inputs,
)
if "entry_point" not in kwargs and depends_on:
# if the CreateModelStep is the first step in the collection
model_step.add_depends_on(depends_on)
steps.append(model_step)
transformer = Transformer(
model_name=model_step.properties.ModelName,
instance_count=instance_count,
instance_type=instance_type,
strategy=strategy,
assemble_with=assemble_with,
output_path=output_path,
output_kms_key=output_kms_key,
accept=accept,
max_concurrent_transforms=max_concurrent_transforms,
max_payload=max_payload,
env=env,
tags=tags,
base_transform_job_name=name,
volume_kms_key=volume_kms_key,
sagemaker_session=estimator.sagemaker_session,
)
transform_step = TransformStep(
name=f"{name}TransformStep",
transformer=transformer,
inputs=transform_inputs,
)
steps.append(transform_step)
self.steps = steps
def __init__(
self,
name: str,
content_types,
response_types,
inference_instances,
transform_instances,
estimator: EstimatorBase = None,
model_data=None,
depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
repack_model_step_retry_policies: List[RetryPolicy] = None,
register_model_step_retry_policies: List[RetryPolicy] = None,
model_package_group_name=None,
model_metrics=None,
approval_status=None,
image_uri=None,
compile_model_family=None,
display_name=None,
description=None,
tags=None,
model: Union[Model, PipelineModel] = None,
drift_check_baselines=None,
customer_metadata_properties=None,
**kwargs,
):
"""Construct steps `_RepackModelStep` and `_RegisterModelStep` based on the estimator.
Args:
name (str): The name of the training step.
estimator: The estimator instance.
model_data: The S3 uri to the model data from training.
content_types (list): The supported MIME types for the input data (default: None).
response_types (list): The supported MIME types for the output data (default: None).
inference_instances (list): A list of the instance types that are used to
generate inferences in real-time (default: None).
transform_instances (list): A list of the instance types on which a transformation
job can be run or on which an endpoint can be deployed (default: None).
depends_on (List[Union[str, Step, StepCollection]]): The list of `Step`/`StepCollection`
names or `Step` instances or `StepCollection` instances that the first step
in the collection depends on (default: None).
repack_model_step_retry_policies (List[RetryPolicy]): The list of retry policies
for the repack model step
register_model_step_retry_policies (List[RetryPolicy]): The list of retry policies
for register model step
model_package_group_name (str): The Model Package Group name or Arn, exclusive to
`model_package_name`, using `model_package_group_name` makes the Model Package
versioned (default: None).
model_metrics (ModelMetrics): ModelMetrics object (default: None).
approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
or "PendingManualApproval" (default: "PendingManualApproval").
image_uri (str): The container image uri for Model Package, if not specified,
Estimator's training container image is used (default: None).
compile_model_family (str): The instance family for the compiled model. If
specified, a compiled model is used (default: None).
description (str): Model Package description (default: None).
tags (List[dict[str, str]]): The list of tags to attach to the model package group. Note
that tags will only be applied to newly created model package groups; if the
name of an existing group is passed to "model_package_group_name",
tags will not be applied.
model (object or Model): A PipelineModel object that comprises a list of models
which gets executed as a serial inference pipeline or a Model object.
drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
metadata properties (default: None).
**kwargs: additional arguments to `create_model`.
"""
self.name = name
steps: List[Step] = []
repack_model = False
self.model_list = None
self.container_def_list = None
subnets = None
security_group_ids = None
if estimator is not None:
subnets = estimator.subnets
security_group_ids = estimator.security_group_ids
elif model is not None and model.vpc_config is not None:
subnets = model.vpc_config["Subnets"]
security_group_ids = model.vpc_config["SecurityGroupIds"]
if "entry_point" in kwargs:
repack_model = True
entry_point = kwargs.pop("entry_point", None)
source_dir = kwargs.pop("source_dir", None)
dependencies = kwargs.pop("dependencies", None)
kwargs = dict(**kwargs, output_kms_key=kwargs.pop("model_kms_key", None))
repack_model_step = _RepackModelStep(
name=f"{name}RepackModel",
depends_on=depends_on,
retry_policies=repack_model_step_retry_policies,
sagemaker_session=estimator.sagemaker_session,
role=estimator.role,
model_data=model_data,
entry_point=entry_point,
source_dir=source_dir,
dependencies=dependencies,
tags=tags,
subnets=subnets,
security_group_ids=security_group_ids,
description=description,
display_name=display_name,
**kwargs,
)
steps.append(repack_model_step)
model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
# remove kwargs consumed by model repacking step
kwargs.pop("output_kms_key", None)
elif model is not None:
if isinstance(model, PipelineModel):
self.model_list = model.models
elif isinstance(model, Model):
self.model_list = [model]
for model_entity in self.model_list:
if estimator is not None:
sagemaker_session = estimator.sagemaker_session
role = estimator.role
else:
sagemaker_session = model_entity.sagemaker_session
role = model_entity.role
if hasattr(model_entity, "entry_point") and model_entity.entry_point is not None:
repack_model = True
entry_point = model_entity.entry_point
source_dir = model_entity.source_dir
dependencies = model_entity.dependencies
kwargs = dict(**kwargs, output_kms_key=model_entity.model_kms_key)
model_name = model_entity.name or model_entity._framework_name
repack_model_step = _RepackModelStep(
name=f"{model_name}RepackModel",
depends_on=depends_on,
retry_policies=repack_model_step_retry_policies,
sagemaker_session=sagemaker_session,
role=role,
model_data=model_entity.model_data,
entry_point=entry_point,
source_dir=source_dir,
dependencies=dependencies,
tags=tags,
subnets=subnets,
security_group_ids=security_group_ids,
description=description,
display_name=display_name,
**kwargs,
)
steps.append(repack_model_step)
model_entity.model_data = (
repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
)
# remove kwargs consumed by model repacking step
kwargs.pop("output_kms_key", None)
if isinstance(model, PipelineModel):
self.container_def_list = model.pipeline_container_def(inference_instances[0])
elif isinstance(model, Model):
self.container_def_list = [model.prepare_container_def(inference_instances[0])]
register_model_step = _RegisterModelStep(
name=name,
estimator=estimator,
model_data=model_data,
content_types=content_types,
response_types=response_types,
inference_instances=inference_instances,
transform_instances=transform_instances,
model_package_group_name=model_package_group_name,
model_metrics=model_metrics,
drift_check_baselines=drift_check_baselines,
approval_status=approval_status,
image_uri=image_uri,
compile_model_family=compile_model_family,
description=description,
display_name=display_name,
tags=tags,
container_def_list=self.container_def_list,
retry_policies=register_model_step_retry_policies,
customer_metadata_properties=customer_metadata_properties,
**kwargs,
)
if not repack_model:
register_model_step.add_depends_on(depends_on)
steps.append(register_model_step)
self.steps = steps
# TODO: add public document link here once ready
warnings.warn(
(
"We are deprecating the use of RegisterModel. "
"Instead, please use the ModelStep, which simply takes in the step arguments "
"generated by model.register()."
),
DeprecationWarning,
)