Merge branch 'aws:master' into fix/update-inference-recommender-fields

Author: rahven14

src/sagemaker/clarify.py

@@ -53,6 +53,7 @@
     OfflineStoreConfig,
     DataCatalogConfig,
     FeatureValue,
+    FeatureParameter,
 )
 
 logger = logging.getLogger(__name__)
@@ -537,6 +538,64 @@ def describe(self, next_token: str = None) -> Dict[str, Any]:
             feature_group_name=self.name, next_token=next_token
         )
 
+    def update(self, feature_additions: Sequence[FeatureDefinition]) -> Dict[str, Any]:
+        """Update a FeatureGroup and add new features from the given feature definitions.
+
+        Args:
+            feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+
+        Returns:
+            Response dict from service.
+        """
+
+        return self.sagemaker_session.update_feature_group(
+            feature_group_name=self.name,
+            feature_additions=[
+                feature_addition.to_dict() for feature_addition in feature_additions
+            ],
+        )
+
+    def update_feature_metadata(
+        self,
+        feature_name: str,
+        description: str = None,
+        parameter_additions: Sequence[FeatureParameter] = None,
+        parameter_removals: Sequence[str] = None,
+    ) -> Dict[str, Any]:
+        """Update a feature metadata and add/remove metadata.
+
+        Args:
+            feature_name (str): name of the feature to update.
+            description (str): description of the feature to update.
+            parameter_additions (Sequence[Dict[str, str]): list of feature parameter to be added.
+            parameter_removals (Sequence[str]): list of feature parameter key to be removed.
+
+        Returns:
+            Response dict from service.
+        """
+        return self.sagemaker_session.update_feature_metadata(
+            feature_group_name=self.name,
+            feature_name=feature_name,
+            description=description,
+            parameter_additions=[
+                parameter_addition.to_dict() for parameter_addition in (parameter_additions or [])
+            ],
+            parameter_removals=(parameter_removals or []),
+        )
+
+    def describe_feature_metadata(self, feature_name: str) -> Dict[str, Any]:
+        """Describe feature metadata by feature name.
+
+        Args:
+            feature_name (str): name of the feature.
+        Returns:
+            Response dict from service.
+        """
+
+        return self.sagemaker_session.describe_feature_metadata(
+            feature_group_name=self.name, feature_name=feature_name
+        )
+
     def load_feature_definitions(
         self,
         data_frame: DataFrame,

src/sagemaker/feature_store/feature_group.py

@@ -207,3 +207,27 @@ def to_dict(self) -> Dict[str, Any]:
             FeatureName=self.feature_name,
             ValueAsString=self.value_as_string,
         )
+
+
+@attr.s
+class FeatureParameter(Config):
+    """FeatureParameter for FeatureStore.
+
+    Attributes:
+        key (str): key of the parameter.
+        value (str): value of the parameter.
+    """
+
+    key: str = attr.ib(default=None)
+    value: str = attr.ib(default=None)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Construct a dictionary based on the attributes provided.
+
+        Returns:
+            dict represents the attributes.
+        """
+        return Config.construct_dict(
+            Key=self.key,
+            Value=self.value,
+        )

src/sagemaker/feature_store/inputs.py

@@ -17,7 +17,6 @@
 import json
 import logging
 import os
-import re
 import copy
 from typing import List, Dict
 
@@ -50,6 +49,8 @@
     ["mxnet", "tensorflow", "keras", "pytorch", "onnx", "xgboost", "tflite"]
 )
 
+NEO_IOC_TARGET_DEVICES = ["ml_c4", "ml_c5", "ml_m4", "ml_m5", "ml_p2", "ml_p3", "ml_g4dn"]
+
 
 class ModelBase(abc.ABC):
     """An object that encapsulates a trained model.
@@ -763,13 +764,22 @@ def _compilation_job_config(
             "Framework": framework.upper(),
         }
 
-        multiple_version_supported_framework_list = ["pytorch", "tensorflow"]
-        if (
-            framework.lower() in multiple_version_supported_framework_list
-            and target_instance_type is not None
-            and re.match("(?=^ml_)(?!ml_inf)", target_instance_type) is not None
-            and framework_version is not None
+        def multi_version_compilation_supported(
+            target_instance_type: str, framework: str, framework_version: str
         ):
+            if target_instance_type and framework and framework_version:
+                framework = framework.lower()
+                multi_version_frameworks_support_mapping = {
+                    "inferentia": ["pytorch", "tensorflow", "mxnet"],
+                    "neo_ioc_targets": ["pytorch", "tensorflow"],
+                }
+                if target_instance_type in NEO_IOC_TARGET_DEVICES:
+                    return framework in multi_version_frameworks_support_mapping["neo_ioc_targets"]
+                if target_instance_type == "ml_inf":
+                    return framework in multi_version_frameworks_support_mapping["inferentia"]
+            return False
+
+        if multi_version_compilation_supported(target_instance_type, framework, framework_version):
             input_model_config["FrameworkVersion"] = utils.get_short_version(framework_version)
 
         role = self.sagemaker_session.expand_role(role)

src/sagemaker/model.py

@@ -4085,7 +4085,7 @@ def describe_feature_group(
         """Describe a FeatureGroup by name in FeatureStore service.
 
         Args:
-            feature_group_name (str): name of the FeatureGroup to descibe.
+            feature_group_name (str): name of the FeatureGroup to describe.
             next_token (str): next_token to get next page of features.
         Returns:
             Response dict from service.
@@ -4095,6 +4095,72 @@ def describe_feature_group(
         update_args(kwargs, NextToken=next_token)
         return self.sagemaker_client.describe_feature_group(**kwargs)
 
+    def update_feature_group(
+        self, feature_group_name: str, feature_additions: Sequence[Dict[str, str]]
+    ) -> Dict[str, Any]:
+        """Update a FeatureGroup and add new features from the given feature definitions.
+
+        Args:
+            feature_group_name (str): name of the FeatureGroup to update.
+            feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+        Returns:
+            Response dict from service.
+        """
+
+        return self.sagemaker_client.update_feature_group(
+            FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
+        )
+
+    def update_feature_metadata(
+        self,
+        feature_group_name: str,
+        feature_name: str,
+        description: str = None,
+        parameter_additions: Sequence[Dict[str, str]] = None,
+        parameter_removals: Sequence[str] = None,
+    ) -> Dict[str, Any]:
+        """Update a feature metadata and add/remove metadata.
+
+        Args:
+            feature_group_name (str): name of the FeatureGroup to update.
+            feature_name (str): name of the feature to update.
+            description (str): description of the feature to update.
+            parameter_additions (Sequence[Dict[str, str]): list of feature parameter to be added.
+            parameter_removals (Sequence[Dict[str, str]): list of feature parameter to be removed.
+        Returns:
+            Response dict from service.
+        """
+
+        request = {
+            "FeatureGroupName": feature_group_name,
+            "FeatureName": feature_name,
+        }
+
+        if description is not None:
+            request["Description"] = description
+        if parameter_additions is not None:
+            request["ParameterAdditions"] = parameter_additions
+        if parameter_removals is not None:
+            request["ParameterRemovals"] = parameter_removals
+
+        return self.sagemaker_client.update_feature_metadata(**request)
+
+    def describe_feature_metadata(
+        self, feature_group_name: str, feature_name: str
+    ) -> Dict[str, Any]:
+        """Describe feature metadata by feature name in FeatureStore service.
+
+        Args:
+            feature_group_name (str): name of the FeatureGroup.
+            feature_name (str): name of the feature.
+        Returns:
+            Response dict from service.
+        """
+
+        return self.sagemaker_client.describe_feature_metadata(
+            FeatureGroupName=feature_group_name, FeatureName=feature_name
+        )
+
     def put_record(
         self,
         feature_group_name: str,

src/sagemaker/session.py

@@ -22,8 +22,9 @@
 import pytest
 from pandas import DataFrame
 
+from sagemaker.feature_store.feature_definition import FractionalFeatureDefinition
 from sagemaker.feature_store.feature_group import FeatureGroup
-from sagemaker.feature_store.inputs import FeatureValue
+from sagemaker.feature_store.inputs import FeatureValue, FeatureParameter
 from sagemaker.session import get_execution_role, Session
 from tests.integ.timeout import timeout
 
@@ -237,6 +238,83 @@ def test_create_feature_store(
     assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
 
 
+def test_update_feature_group(
+    feature_store_session,
+    role,
+    feature_group_name,
+    offline_store_s3_uri,
+    pandas_data_frame,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=offline_store_s3_uri,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        new_feature_name = "new_feature"
+        new_features = [FractionalFeatureDefinition(feature_name=new_feature_name)]
+        feature_group.update(new_features)
+        _wait_for_feature_group_update(feature_group)
+        feature_definitions = feature_group.describe().get("FeatureDefinitions")
+        assert any([True for elem in feature_definitions if new_feature_name in elem.values()])
+
+
+def test_feature_metadata(
+    feature_store_session,
+    role,
+    feature_group_name,
+    offline_store_s3_uri,
+    pandas_data_frame,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=offline_store_s3_uri,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        parameter_additions = [
+            FeatureParameter(key="key1", value="value1"),
+            FeatureParameter(key="key2", value="value2"),
+        ]
+        description = "test description"
+        feature_name = "feature1"
+        feature_group.update_feature_metadata(
+            feature_name=feature_name,
+            description=description,
+            parameter_additions=parameter_additions,
+        )
+        describe_feature_metadata = feature_group.describe_feature_metadata(
+            feature_name=feature_name
+        )
+        print(describe_feature_metadata)
+        assert description == describe_feature_metadata.get("Description")
+        assert 2 == len(describe_feature_metadata.get("Parameters"))
+
+        parameter_removals = ["key1"]
+        feature_group.update_feature_metadata(
+            feature_name=feature_name, parameter_removals=parameter_removals
+        )
+        describe_feature_metadata = feature_group.describe_feature_metadata(
+            feature_name=feature_name
+        )
+        assert description == describe_feature_metadata.get("Description")
+        assert 1 == len(describe_feature_metadata.get("Parameters"))
+
+
 def test_ingest_without_string_feature(
     feature_store_session,
     role,
@@ -304,6 +382,18 @@ def _wait_for_feature_group_create(feature_group: FeatureGroup):
     print(f"FeatureGroup {feature_group.name} successfully created.")
 
 
+def _wait_for_feature_group_update(feature_group: FeatureGroup):
+    status = feature_group.describe().get("LastUpdateStatus").get("Status")
+    while status == "InProgress":
+        print("Waiting for Feature Group Update")
+        time.sleep(5)
+        status = feature_group.describe().get("LastUpdateStatus").get("Status")
+    if status != "Successful":
+        print(feature_group.describe())
+        raise RuntimeError(f"Failed to update feature group {feature_group.name}")
+    print(f"FeatureGroup {feature_group.name} successfully updated.")
+
+
 @contextmanager
 def cleanup_feature_group(feature_group: FeatureGroup):
     try: