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, )