Merge branch 'master' into jobNamePrefix

Author: icywang86rui

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## v2.46.1 (2021-06-22)
+
+### Bug Fixes and Other Changes
+
+ * Register model step tags
+
+### Documentation Changes
+
+ * update to include new batch_get_record api call
+ * Correct type annotation for TrainingStep inputs
+ * introduce input mode FastFile
+ * update hf transformer version
+
 ## v2.46.0 (2021-06-15)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.46.1.dev0
+2.46.2.dev0

doc/amazon_sagemaker_featurestore.rst

@@ -291,6 +291,13 @@ example identifier to retrieve the record.
    record_identifier_value = str(2990130)
    featurestore_runtime.get_record(FeatureGroupName=transaction_feature_group_name, RecordIdentifierValueAsString=record_identifier_value)
 
+You can use the ``batch_get_record`` function to retrieve multiple records simultaneously from your feature store. The following example uses this API to retrieve a batch of records.
+
+.. code:: python
+
+   record_identifier_values = ["573291", "109382", "828400", "124013"]
+   featurestore_runtime.batch_get_record(Identifiers=[{"FeatureGroupName": transaction_feature_group_name, "RecordIdentifiersValueAsString": record_identifier_values}])
+
 An example response from the fraud detection example:
 
 .. code:: python

src/sagemaker/estimator.py

@@ -1007,6 +1007,8 @@ def register(
         if compile_model_family is not None:
             model = self._compiled_models[compile_model_family]
         else:
+            if "model_kms_key" not in kwargs:
+                kwargs["model_kms_key"] = self.output_kms_key
             model = self.create_model(image_uri=image_uri, **kwargs)
         model.name = model_name
         return model.register(

src/sagemaker/inputs.py

@@ -70,6 +70,8 @@ def __init__(
                         a local directory.
                     * 'Pipe' - Amazon SageMaker streams data directly from S3 to the container via
                         a Unix-named pipe.
+                    * 'FastFile' - Amazon SageMaker streams data from S3 on demand instead of
+                        downloading the entire dataset before training begins.
 
             attribute_names (list[str]): A list of one or more attribute names to use that are
                 found in a specified AugmentedManifestFile.

src/sagemaker/model.py

@@ -195,6 +195,7 @@ def _get_model_package_args(
         marketplace_cert=False,
         approval_status=None,
         description=None,
+        tags=None,
     ):
         """Get arguments for session.create_model_package method.
 
@@ -250,6 +251,8 @@ def _get_model_package_args(
             model_package_args["approval_status"] = approval_status
         if description is not None:
             model_package_args["description"] = description
+        if tags is not None:
+            model_package_args["tags"] = tags
         return model_package_args
 
     def _init_sagemaker_session_if_does_not_exist(self, instance_type):

src/sagemaker/session.py

@@ -467,6 +467,8 @@ def train(  # noqa: C901
                 a directory in the Docker container.
                 * 'Pipe' - Amazon SageMaker streams data directly from S3 to the container via a
                 Unix-named pipe.
+                * 'FastFile' - Amazon SageMaker streams data from S3 on demand instead of
+                downloading the entire dataset before training begins.
             input_config (list): A list of Channel objects. Each channel is a named input source.
                 Please refer to the format details described:
                 https://botocore.readthedocs.io/en/latest/reference/services/sagemaker.html#SageMaker.Client.create_training_job
@@ -609,6 +611,8 @@ def _get_train_request(  # noqa: C901
                 a directory in the Docker container.
                 * 'Pipe' - Amazon SageMaker streams data directly from S3 to the container via a
                 Unix-named pipe.
+                * 'FastFile' - Amazon SageMaker streams data from S3 on demand instead of
+                downloading the entire dataset before training begins.
             input_config (list): A list of Channel objects. Each channel is a named input source.
                 Please refer to the format details described:
                 https://botocore.readthedocs.io/en/latest/reference/services/sagemaker.html#SageMaker.Client.create_training_job
@@ -1897,6 +1901,8 @@ def tune(  # noqa: C901
                 a directory in the Docker container.
                 * 'Pipe' - Amazon SageMaker streams data directly from S3 to the container via a
                 Unix-named pipe.
+                * 'FastFile' - Amazon SageMaker streams data from S3 on demand instead of
+                downloading the entire dataset before training begins.
             metric_definitions (list[dict]): A list of dictionaries that defines the metric(s)
                 used to evaluate the training jobs. Each dictionary contains two keys: 'Name' for
                 the name of the metric, and 'Regex' for the regular expression used to extract the
@@ -2180,6 +2186,8 @@ def _map_training_config(
                     a directory in the Docker container.
                 * 'Pipe' - Amazon SageMaker streams data directly from S3 to the container via a
                     Unix-named pipe.
+                * 'FastFile' - Amazon SageMaker streams data from S3 on demand instead of
+                    downloading the entire dataset before training begins.
             role (str): An AWS IAM role (either name or full ARN). The Amazon SageMaker training
                 jobs and APIs that create Amazon SageMaker endpoints use this role to access
                 training data and model artifacts. You must grant sufficient permissions to
@@ -2716,6 +2724,7 @@ def _get_create_model_package_request(
         marketplace_cert=False,
         approval_status="PendingManualApproval",
         description=None,
+        tags=None,
     ):
         """Get request dictionary for CreateModelPackage API.
 
@@ -2753,6 +2762,8 @@ def _get_create_model_package_request(
             request_dict["ModelPackageGroupName"] = model_package_group_name
         if description is not None:
             request_dict["ModelPackageDescription"] = description
+        if tags is not None:
+            request_dict["Tags"] = tags
         if model_metrics:
             request_dict["ModelMetrics"] = model_metrics
         if metadata_properties:

src/sagemaker/workflow/_utils.py

@@ -60,6 +60,7 @@ def __init__(
         source_dir: str = None,
         dependencies: List = None,
         depends_on: List[str] = None,
+        **kwargs,
     ):
         """Constructs a TrainingStep, given an `EstimatorBase` instance.
 
@@ -98,6 +99,7 @@ def __init__(
                 "inference_script": self._entry_point_basename,
                 "model_archive": self._model_archive,
             },
+            **kwargs,
         )
         repacker.disable_profiler = True
         inputs = TrainingInput(self._model_prefix)
@@ -225,6 +227,7 @@ def __init__(
         compile_model_family=None,
         description=None,
         depends_on: List[str] = None,
+        tags=None,
         **kwargs,
     ):
         """Constructor of a register model step.
@@ -264,6 +267,7 @@ def __init__(
         self.inference_instances = inference_instances
         self.transform_instances = transform_instances
         self.model_package_group_name = model_package_group_name
+        self.tags = tags
         self.model_metrics = model_metrics
         self.metadata_properties = metadata_properties
         self.approval_status = approval_status
@@ -324,10 +328,12 @@ def arguments(self) -> RequestType:
             metadata_properties=self.metadata_properties,
             approval_status=self.approval_status,
             description=self.description,
+            tags=self.tags,
         )
         request_dict = model.sagemaker_session._get_create_model_package_request(
             **model_package_args
         )
+
         # these are not available in the workflow service and will cause rejection
         if "CertifyForMarketplace" in request_dict:
             request_dict.pop("CertifyForMarketplace")

src/sagemaker/workflow/callback_step.py

@@ -58,19 +58,20 @@ def to_request(self) -> RequestType:
             "OutputType": self.output_type.value,
         }
 
-    @property
-    def expr(self) -> Dict[str, str]:
-        """The 'Get' expression dict for a `Parameter`."""
-        return CallbackOutput._expr(self.output_name)
+    def expr(self, step_name) -> Dict[str, str]:
+        """The 'Get' expression dict for a `CallbackOutput`."""
+        return CallbackOutput._expr(self.output_name, step_name)
 
     @classmethod
-    def _expr(cls, name):
+    def _expr(cls, name, step_name):
         """An internal classmethod for the 'Get' expression dict for a `CallbackOutput`.
 
         Args:
             name (str): The name of the callback output.
+            step_name (str): The name of the step the callback step associated
+                with this output belongs to.
         """
-        return {"Get": f"Steps.{name}.OutputParameters['{name}']"}
+        return {"Get": f"Steps.{step_name}.OutputParameters['{name}']"}
 
 
 class CallbackStep(Step):

src/sagemaker/workflow/pipeline.py

@@ -24,7 +24,7 @@
 
 from sagemaker._studio import _append_project_tags
 from sagemaker.session import Session
-from sagemaker.workflow.callback_step import CallbackOutput
+from sagemaker.workflow.callback_step import CallbackOutput, CallbackStep
 from sagemaker.workflow.entities import (
     Entity,
     Expression,
@@ -240,9 +240,12 @@ def definition(self) -> str:
         """Converts a request structure to string representation for workflow service calls."""
         request_dict = self.to_request()
         request_dict["PipelineExperimentConfig"] = interpolate(
-            request_dict["PipelineExperimentConfig"]
+            request_dict["PipelineExperimentConfig"], {}
+        )
+        callback_output_to_step_map = _map_callback_outputs(self.steps)
+        request_dict["Steps"] = interpolate(
+            request_dict["Steps"], callback_output_to_step_map=callback_output_to_step_map
         )
-        request_dict["Steps"] = interpolate(request_dict["Steps"])
 
         return json.dumps(request_dict)
 
@@ -263,38 +266,62 @@ def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]:
     return [{"Name": name, "Value": str(value)} for name, value in parameters.items()]
 
 
-def interpolate(request_obj: RequestType) -> RequestType:
+def interpolate(
+    request_obj: RequestType, callback_output_to_step_map: Dict[str, str]
+) -> RequestType:
     """Replaces Parameter values in a list of nested Dict[str, Any] with their workflow expression.
 
     Args:
         request_obj (RequestType): The request dict.
+        callback_output_to_step_map (Dict[str, str]): A dict of output name -> step name.
 
     Returns:
         RequestType: The request dict with Parameter values replaced by their expression.
     """
     request_obj_copy = deepcopy(request_obj)
-    return _interpolate(request_obj_copy)
+    return _interpolate(request_obj_copy, callback_output_to_step_map=callback_output_to_step_map)
 
 
-def _interpolate(obj: Union[RequestType, Any]):
+def _interpolate(obj: Union[RequestType, Any], callback_output_to_step_map: Dict[str, str]):
     """Walks the nested request dict, replacing Parameter type values with workflow expressions.
 
     Args:
         obj (Union[RequestType, Any]): The request dict.
+        callback_output_to_step_map (Dict[str, str]): A dict of output name -> step name.
     """
-    if isinstance(obj, (Expression, Parameter, Properties, CallbackOutput)):
+    if isinstance(obj, (Expression, Parameter, Properties)):
         return obj.expr
+    if isinstance(obj, CallbackOutput):
+        step_name = callback_output_to_step_map[obj.output_name]
+        return obj.expr(step_name)
     if isinstance(obj, dict):
         new = obj.__class__()
         for key, value in obj.items():
-            new[key] = interpolate(value)
+            new[key] = interpolate(value, callback_output_to_step_map)
     elif isinstance(obj, (list, set, tuple)):
-        new = obj.__class__(interpolate(value) for value in obj)
+        new = obj.__class__(interpolate(value, callback_output_to_step_map) for value in obj)
     else:
         return obj
     return new
 
 
+def _map_callback_outputs(steps: List[Step]):
+    """Iterate over the provided steps, building a map of callback output parameters to step names.
+
+    Args:
+        step (List[Step]): The steps list.
+    """
+
+    callback_output_map = {}
+    for step in steps:
+        if isinstance(step, CallbackStep):
+            if step.outputs:
+                for output in step.outputs:
+                    callback_output_map[output.output_name] = step.name
+
+    return callback_output_map
+
+
 def update_args(args: Dict[str, Any], **kwargs):
     """Updates the request arguments dict with a value, if populated.
 

src/sagemaker/workflow/step_collections.py

@@ -67,6 +67,7 @@ def __init__(
         image_uri=None,
         compile_model_family=None,
         description=None,
+        tags=None,
         **kwargs,
     ):
         """Construct steps `_RepackModelStep` and `_RegisterModelStep` based on the estimator.
@@ -94,15 +95,21 @@ def __init__(
             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.
             **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"]
+            entry_point = kwargs.pop("entry_point", None)
             source_dir = kwargs.get("source_dir")
             dependencies = kwargs.get("dependencies")
+            kwargs = dict(**kwargs, output_kms_key=kwargs.pop("model_kms_key", None))
+
             repack_model_step = _RepackModelStep(
                 name=f"{name}RepackModel",
                 depends_on=depends_on,
@@ -111,6 +118,7 @@ def __init__(
                 entry_point=entry_point,
                 source_dir=source_dir,
                 dependencies=dependencies,
+                **kwargs,
             )
             steps.append(repack_model_step)
             model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
@@ -119,6 +127,7 @@ def __init__(
         kwargs.pop("entry_point", None)
         kwargs.pop("source_dir", None)
         kwargs.pop("dependencies", None)
+        kwargs.pop("output_kms_key", None)
 
         register_model_step = _RegisterModelStep(
             name=name,
@@ -134,6 +143,7 @@ def __init__(
             image_uri=image_uri,
             compile_model_family=compile_model_family,
             description=description,
+            tags=tags,
             **kwargs,
         )
         if not repack_model:

src/sagemaker/workflow/steps.py

@@ -16,16 +16,12 @@
 import abc
 
 from enum import Enum
-from typing import Dict, List
+from typing import Dict, List, Union
 
 import attr
 
 from sagemaker.estimator import EstimatorBase, _TrainingJob
-from sagemaker.inputs import (
-    CreateModelInput,
-    TrainingInput,
-    TransformInput,
-)
+from sagemaker.inputs import CreateModelInput, TrainingInput, TransformInput, FileSystemInput
 from sagemaker.model import Model
 from sagemaker.processing import (
     ProcessingInput,
@@ -145,7 +141,7 @@ def __init__(
         self,
         name: str,
         estimator: EstimatorBase,
-        inputs: TrainingInput = None,
+        inputs: Union[TrainingInput, dict, str, FileSystemInput] = None,
         cache_config: CacheConfig = None,
         depends_on: List[str] = None,
     ):
@@ -157,7 +153,23 @@ def __init__(
         Args:
             name (str): The name of the training step.
             estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance.
-            inputs (TrainingInput): A `sagemaker.inputs.TrainingInput` instance. Defaults to `None`.
+            inputs (str or dict or sagemaker.inputs.TrainingInput
+                or sagemaker.inputs.FileSystemInput): Information
+                about the training data. This can be one of three types:
+
+                * (str) the S3 location where training data is saved, or a file:// path in
+                    local mode.
+                * (dict[str, str] or dict[str, sagemaker.inputs.TrainingInput]) If using multiple
+                    channels for training data, you can specify a dict mapping channel names to
+                    strings or :func:`~sagemaker.inputs.TrainingInput` objects.
+                * (sagemaker.inputs.TrainingInput) - channel configuration for S3 data sources
+                    that can provide additional information as well as the path to the training
+                    dataset.
+                    See :func:`sagemaker.inputs.TrainingInput` for full details.
+                * (sagemaker.inputs.FileSystemInput) - channel configuration for
+                    a file system data source that can provide additional information as well as
+                    the path to the training dataset.
+
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
             depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.TrainingStep`
                 depends on