Merge branch 'master' into feature/new_fg_utils

Author: JoseJuan98

CHANGELOG.md

@@ -1,5 +1,21 @@
 # Changelog
 
+## v2.130.0 (2023-01-26)
+
+### Features
+
+ * Add PyTorch 1.13.1 to SDK
+ * Adding image_uri config for DJL containers
+ * Support specifying env-vars when creating model from model package
+ * local download dir for Model and Estimator classes
+
+### Bug Fixes and Other Changes
+
+ * increase creation time slack minutes
+ * Enable load_run auto pass in experiment config
+ * Add us-isob-east-1 accounts and configs
+ * Clean up Pipeline unit tests
+
 ## v2.129.0 (2023-01-19)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.129.1.dev0
+2.130.1.dev0

src/sagemaker/estimator.py

@@ -155,6 +155,8 @@ def __init__(
         entry_point: Optional[Union[str, PipelineVariable]] = None,
         dependencies: Optional[List[Union[str]]] = None,
         instance_groups: Optional[List[InstanceGroup]] = None,
+        training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None,
+        training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -489,6 +491,18 @@ def __init__(
                 `Train Using a Heterogeneous Cluster
                 <https://docs.aws.amazon.com/sagemaker/latest/dg/train-heterogeneous-cluster.html>`_
                 in the *Amazon SageMaker developer guide*.
+            training_repository_access_mode (str): Optional. Specifies how SageMaker accesses the
+                Docker image that contains the training algorithm (default: None).
+                Set this to one of the following values:
+                * 'Platform' - The training image is hosted in Amazon ECR.
+                * 'Vpc' - The training image is hosted in a private Docker registry in your VPC.
+                When it's default to None, its behavior will be same as 'Platform' - image is hosted
+                in ECR.
+            training_repository_credentials_provider_arn (str): Optional. The Amazon Resource Name
+                (ARN) of an AWS Lambda function that provides credentials to authenticate to the
+                private Docker registry where your training image is hosted (default: None).
+                When it's set to None, SageMaker will not do authentication before pulling the image
+                in the private Docker registry.
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -536,7 +550,9 @@ def __init__(
         self.dependencies = dependencies or []
         self.uploaded_code = None
         self.tags = add_jumpstart_tags(
-            tags=tags, training_model_uri=self.model_uri, training_script_uri=self.source_dir
+            tags=tags,
+            training_model_uri=self.model_uri,
+            training_script_uri=self.source_dir,
         )
         if self.instance_type in ("local", "local_gpu"):
             if self.instance_type == "local_gpu" and self.instance_count > 1:
@@ -571,6 +587,12 @@ def __init__(
         self.subnets = subnets
         self.security_group_ids = security_group_ids
 
+        # training image configs
+        self.training_repository_access_mode = training_repository_access_mode
+        self.training_repository_credentials_provider_arn = (
+            training_repository_credentials_provider_arn
+        )
+
         self.encrypt_inter_container_traffic = encrypt_inter_container_traffic
         self.use_spot_instances = use_spot_instances
         self.max_wait = max_wait
@@ -651,7 +673,8 @@ def _ensure_base_job_name(self):
             self.base_job_name
             or get_jumpstart_base_name_if_jumpstart_model(self.source_dir, self.model_uri)
             or base_name_from_image(
-                self.training_image_uri(), default_base_name=EstimatorBase.JOB_CLASS_NAME
+                self.training_image_uri(),
+                default_base_name=EstimatorBase.JOB_CLASS_NAME,
             )
         )
 
@@ -1405,7 +1428,10 @@ def deploy(
         self._ensure_base_job_name()
 
         jumpstart_base_name = get_jumpstart_base_name_if_jumpstart_model(
-            kwargs.get("source_dir"), self.source_dir, kwargs.get("model_data"), self.model_uri
+            kwargs.get("source_dir"),
+            self.source_dir,
+            kwargs.get("model_data"),
+            self.model_uri,
         )
         default_name = (
             name_from_base(jumpstart_base_name)
@@ -1638,6 +1664,15 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
             init_params["algorithm_arn"] = job_details["AlgorithmSpecification"]["AlgorithmName"]
         elif "TrainingImage" in job_details["AlgorithmSpecification"]:
             init_params["image_uri"] = job_details["AlgorithmSpecification"]["TrainingImage"]
+            if "TrainingImageConfig" in job_details["AlgorithmSpecification"]:
+                init_params["training_repository_access_mode"] = job_details[
+                    "AlgorithmSpecification"
+                ]["TrainingImageConfig"].get("TrainingRepositoryAccessMode")
+                init_params["training_repository_credentials_provider_arn"] = (
+                    job_details["AlgorithmSpecification"]["TrainingImageConfig"]
+                    .get("TrainingRepositoryAuthConfig", {})
+                    .get("TrainingRepositoryCredentialsProviderArn")
+                )
         else:
             raise RuntimeError(
                 "Invalid AlgorithmSpecification. Either TrainingImage or "
@@ -2118,6 +2153,17 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         else:
             train_args["retry_strategy"] = None
 
+        if estimator.training_repository_access_mode is not None:
+            training_image_config = {
+                "TrainingRepositoryAccessMode": estimator.training_repository_access_mode
+            }
+            if estimator.training_repository_credentials_provider_arn is not None:
+                training_image_config["TrainingRepositoryAuthConfig"] = {}
+                training_image_config["TrainingRepositoryAuthConfig"][
+                    "TrainingRepositoryCredentialsProviderArn"
+                ] = estimator.training_repository_credentials_provider_arn
+            train_args["training_image_config"] = training_image_config
+
         # encrypt_inter_container_traffic may be a pipeline variable place holder object
         # which is parsed in execution time
         if estimator.encrypt_inter_container_traffic:
@@ -2182,7 +2228,11 @@ def _is_local_channel(cls, input_uri):
 
     @classmethod
     def update(
-        cls, estimator, profiler_rule_configs=None, profiler_config=None, resource_config=None
+        cls,
+        estimator,
+        profiler_rule_configs=None,
+        profiler_config=None,
+        resource_config=None,
     ):
         """Update a running Amazon SageMaker training job.
 
@@ -2321,6 +2371,8 @@ def __init__(
         entry_point: Optional[Union[str, PipelineVariable]] = None,
         dependencies: Optional[List[str]] = None,
         instance_groups: Optional[List[InstanceGroup]] = None,
+        training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None,
+        training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -2654,6 +2706,18 @@ def __init__(
                 `Train Using a Heterogeneous Cluster
                 <https://docs.aws.amazon.com/sagemaker/latest/dg/train-heterogeneous-cluster.html>`_
                 in the *Amazon SageMaker developer guide*.
+            training_repository_access_mode (str): Optional. Specifies how SageMaker accesses the
+                Docker image that contains the training algorithm (default: None).
+                Set this to one of the following values:
+                * 'Platform' - The training image is hosted in Amazon ECR.
+                * 'Vpc' - The training image is hosted in a private Docker registry in your VPC.
+                When it's default to None, its behavior will be same as 'Platform' - image is hosted
+                in ECR.
+            training_repository_credentials_provider_arn (str): Optional. The Amazon Resource Name
+                (ARN) of an AWS Lambda function that provides credentials to authenticate to the
+                private Docker registry where your training image is hosted (default: None).
+                When it's set to None, SageMaker will not do authentication before pulling the image
+                in the private Docker registry.
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
@@ -2698,6 +2762,8 @@ def __init__(
             dependencies=dependencies,
             hyperparameters=hyperparameters,
             instance_groups=instance_groups,
+            training_repository_access_mode=training_repository_access_mode,
+            training_repository_credentials_provider_arn=training_repository_credentials_provider_arn,  # noqa: E501 # pylint: disable=line-too-long
             **kwargs,
         )
 

src/sagemaker/feature_store/feature_group.py

@@ -710,20 +710,25 @@ def ingest(
     ) -> IngestionManagerPandas:
         """Ingest the content of a pandas DataFrame to feature store.
 
-        ``max_worker`` number of thread will be created to work on different partitions of
-        the ``data_frame`` in parallel.
+        ``max_worker`` the number of threads created to work on different partitions of the
+        ``data_frame`` in parallel.
 
-        ``max_processes`` number of processes will be created to work on different partitions
-        of the ``data_frame`` in parallel, each with ``max_worker`` threads.
+        ``max_processes`` the number of processes will be created to work on different
+        partitions of the ``data_frame`` in parallel, each with ``max_worker`` threads.
 
-        The ingest function will attempt to ingest all records in the data frame. If ``wait``
-        is True, then an exception is thrown after all records have been processed. If ``wait``
-        is False, then a later call to the returned instance IngestionManagerPandas' ``wait()``
-        function will throw an exception.
+        The ingest function attempts to ingest all records in the data frame. SageMaker
+        Feature Store throws an exception if it fails to ingest any records.
 
-        Zero based indices of rows that failed to be ingested can be found in the exception.
-        They can also be found from the IngestionManagerPandas' ``failed_rows`` function after
-        the exception is thrown.
+        If ``wait`` is ``True``, Feature Store runs the ``ingest`` function synchronously.
+        You receive an ``IngestionError`` if there are any records that can't be ingested.
+        If ``wait`` is ``False``, Feature Store runs the ``ingest`` function asynchronously.
+
+        Instead of setting ``wait`` to ``True`` in the ``ingest`` function, you can invoke
+        the ``wait`` function on the returned instance of ``IngestionManagerPandas`` to run
+        the ``ingest`` function synchronously.
+
+        To access the rows that failed to ingest, set ``wait`` to ``False``. The
+        ``IngestionError.failed_rows`` object saves all of the rows that failed to ingest.
 
         `profile_name` argument is an optional one. It will use the default credential if None is
         passed. This `profile_name` is used in the sagemaker_featurestore_runtime client only. See

src/sagemaker/feature_store/feature_store.py

@@ -26,6 +26,13 @@
 from sagemaker import Session
 from sagemaker.feature_store.dataset_builder import DatasetBuilder
 from sagemaker.feature_store.feature_group import FeatureGroup
+from sagemaker.feature_store.inputs import (
+    Filter,
+    ResourceEnum,
+    SearchOperatorEnum,
+    SortOrderEnum,
+    Identifier,
+)
 
 
 @attr.s
@@ -114,6 +121,7 @@ def list_feature_groups(
             sort_by (str): The value on which the FeatureGroup list is sorted.
             max_results (int): The maximum number of results returned by ListFeatureGroups.
             next_token (str): A token to resume pagination of ListFeatureGroups results.
+
         Returns:
             Response dict from service.
         """
@@ -128,3 +136,61 @@ def list_feature_groups(
             max_results=max_results,
             next_token=next_token,
         )
+
+    def batch_get_record(self, identifiers: Sequence[Identifier]) -> Dict[str, Any]:
+        """Get record in batch from FeatureStore
+
+        Args:
+            identifiers (Sequence[Identifier]): A list of identifiers to uniquely identify records
+                in FeatureStore.
+
+        Returns:
+            Response dict from service.
+        """
+        batch_get_record_identifiers = [identifier.to_dict() for identifier in identifiers]
+        return self.sagemaker_session.batch_get_record(identifiers=batch_get_record_identifiers)
+
+    def search(
+        self,
+        resource: ResourceEnum,
+        filters: Sequence[Filter] = None,
+        operator: SearchOperatorEnum = None,
+        sort_by: str = None,
+        sort_order: SortOrderEnum = None,
+        next_token: str = None,
+        max_results: int = None,
+    ) -> Dict[str, Any]:
+        """Search for FeatureGroups or FeatureMetadata satisfying given filters.
+
+        Args:
+            resource (ResourceEnum): The name of the Amazon SageMaker resource to search for.
+                Valid values are ``FeatureGroup`` or ``FeatureMetadata``.
+            filters (Sequence[Filter]): A list of filter objects (Default: None).
+            operator (SearchOperatorEnum): A Boolean operator used to evaluate the filters.
+                Valid values are ``And`` or ``Or``. The default is ``And`` (Default: None).
+            sort_by (str): The name of the resource property used to sort the ``SearchResults``.
+                The default is ``LastModifiedTime``.
+            sort_order (SortOrderEnum): How ``SearchResults`` are ordered.
+                Valid values are ``Ascending`` or ``Descending``. The default is ``Descending``.
+            next_token (str): If more than ``MaxResults`` resources match the specified
+                filters, the response includes a ``NextToken``. The ``NextToken`` can be passed to
+                the next ``SearchRequest`` to continue retrieving results (Default: None).
+            max_results (int): The maximum number of results to return (Default: None).
+
+        Returns:
+            Response dict from service.
+        """
+        search_expression = {}
+        if filters:
+            search_expression["Filters"] = [filter.to_dict() for filter in filters]
+        if operator:
+            search_expression["Operator"] = str(operator)
+
+        return self.sagemaker_session.search(
+            resource=str(resource),
+            search_expression=search_expression,
+            sort_by=sort_by,
+            sort_order=None if not sort_order else str(sort_order),
+            next_token=next_token,
+            max_results=max_results,
+        )