workbench 0.8.174__py3-none-any.whl → 0.8.227__py3-none-any.whl

workbench/core/artifacts/model_core.py

@@ -21,6 +21,7 @@ from workbench.utils.aws_utils import newest_path, pull_s3_data
 from workbench.utils.s3_utils import compute_s3_object_hash
 from workbench.utils.shap_utils import shap_values_data, shap_feature_importance
 from workbench.utils.deprecated_utils import deprecated
+from workbench.utils.model_utils import published_proximity_model, get_model_hyperparameters
 
 
 class ModelType(Enum):
@@ -29,92 +30,64 @@ class ModelType(Enum):
     CLASSIFIER = "classifier"
     REGRESSOR = "regressor"
     CLUSTERER = "clusterer"
-    TRANSFORMER = "transformer"
     PROXIMITY = "proximity"
     PROJECTION = "projection"
     UQ_REGRESSOR = "uq_regressor"
     ENSEMBLE_REGRESSOR = "ensemble_regressor"
+    TRANSFORMER = "transformer"
     UNKNOWN = "unknown"
 
 
-# Deprecated Images
-"""
-        # US East 1 images
-        "py312-general-ml-training"
-        ("us-east-1", "training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-sklearn-xgb-training:0.1"
-        ),
-        ("us-east-1", "inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1"
-        ),
-
-        # US West 2 images
-        ("us-west-2", "training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-training:0.1"
-        ),
-        ("us-west-2", "inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1"
-        ),
-
-        # ARM64 images
-        ("us-east-1", "inference", "0.1", "arm64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1-arm64"
-        ),
-        ("us-west-2", "inference", "0.1", "arm64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1-arm64"
-        ),
-"""
+class ModelFramework(Enum):
+    """Enumerated Types for Workbench Model Frameworks"""
+
+    SKLEARN = "sklearn"
+    XGBOOST = "xgboost"
+    LIGHTGBM = "lightgbm"
+    PYTORCH = "pytorch"
+    CHEMPROP = "chemprop"
+    TRANSFORMER = "transformer"
+    META = "meta"
+    UNKNOWN = "unknown"
 
 
 class ModelImages:
     """Class for retrieving workbench inference images"""
 
-    image_uris = {
-        # US East 1 images
-        ("us-east-1", "training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-general-ml-training:0.1"
-        ),
-        ("us-east-1", "inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-general-ml-inference:0.1"
-        ),
-        ("us-east-1", "pytorch_training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-pytorch-training:0.1"
-        ),
-        ("us-east-1", "pytorch_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-pytorch-inference:0.1"
-        ),
-        # US West 2 images
-        ("us-west-2", "training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-general-ml-training:0.1"
-        ),
-        ("us-west-2", "inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-general-ml-inference:0.1"
-        ),
-        ("us-west-2", "pytorch_training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-pytorch-training:0.1"
-        ),
-        ("us-west-2", "pytorch_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-pytorch-inference:0.1"
-        ),
-        # ARM64 images
-        # Meta Endpoint inference images
-        ("us-east-1", "meta-endpoint", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-meta-endpoint:0.1"
-        ),
-        ("us-west-2", "meta-endpoint", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-meta-endpoint:0.1"
-        ),
+    # Account ID
+    ACCOUNT_ID = "507740646243"
+
+    # Image name mappings
+    IMAGE_NAMES = {
+        "training": "py312-general-ml-training",
+        "inference": "py312-general-ml-inference",
+        "pytorch_training": "py312-pytorch-training",
+        "pytorch_inference": "py312-pytorch-inference",
+        "meta_training": "py312-meta-training",
+        "meta_inference": "py312-meta-inference",
     }
 
     @classmethod
-    def get_image_uri(cls, region, image_type, version="0.1", architecture="x86_64"):
-        key = (region, image_type, version, architecture)
-        if key in cls.image_uris:
-            return cls.image_uris[key]
-        else:
-            raise ValueError(
-                f"No matching image found for region: {region}, image_type: {image_type}, version: {version}"
-            )
+    def get_image_uri(cls, region, image_type, version="latest", architecture="x86_64"):
+        """
+        Dynamically construct ECR image URI.
+
+        Args:
+            region: AWS region (e.g., 'us-east-1', 'us-west-2')
+            image_type: Type of image (e.g., 'training', 'inference', 'pytorch_training')
+            version: Image version (e.g., '0.1', '0.2' defaults to 'latest')
+            architecture: CPU architecture (default: 'x86_64', currently unused but kept for compatibility)
+
+        Returns:
+            ECR image URI string
+        """
+        if image_type not in cls.IMAGE_NAMES:
+            raise ValueError(f"Unknown image_type: {image_type}. Valid types: {list(cls.IMAGE_NAMES.keys())}")
+
+        image_name = cls.IMAGE_NAMES[image_type]
+        uri = f"{cls.ACCOUNT_ID}.dkr.ecr.{region}.amazonaws.com/aws-ml-images/{image_name}:{version}"
+
+        return uri
 
 
 class ModelCore(Artifact):
@@ -128,11 +101,10 @@ class ModelCore(Artifact):
         ```
     """
 
-    def __init__(self, model_name: str, model_type: ModelType = None, **kwargs):
+    def __init__(self, model_name: str, **kwargs):
         """ModelCore Initialization
         Args:
             model_name (str): Name of Model in Workbench.
-            model_type (ModelType, optional): Set this for newly created Models. Defaults to None.
             **kwargs: Additional keyword arguments
         """
 
@@ -166,10 +138,8 @@ class ModelCore(Artifact):
                 self.latest_model = self.model_meta["ModelPackageList"][0]
                 self.description = self.latest_model.get("ModelPackageDescription", "-")
                 self.training_job_name = self._extract_training_job_name()
-                if model_type:
-                    self._set_model_type(model_type)
-                else:
-                    self.model_type = self._get_model_type()
+                self.model_type = self._get_model_type()
+                self.model_framework = self._get_model_framework()
             except (IndexError, KeyError):
                 self.log.critical(f"Model {self.model_name} appears to be malformed. Delete and recreate it!")
                 return
@@ -295,21 +265,25 @@ class ModelCore(Artifact):
         else:
             self.log.important(f"No inference data found for {self.model_name}!")
 
-    def get_inference_metrics(self, capture_name: str = "latest") -> Union[pd.DataFrame, None]:
+    def get_inference_metrics(self, capture_name: str = "auto") -> Union[pd.DataFrame, None]:
         """Retrieve the inference performance metrics for this model
 
         Args:
-            capture_name (str, optional): Specific capture_name or "training" (default: "latest")
+            capture_name (str, optional): Specific capture_name (default: "auto")
         Returns:
             pd.DataFrame: DataFrame of the Model Metrics
 
         Note:
-            If a capture_name isn't specified this will try to return something reasonable
+            If a capture_name isn't specified this will try to the 'first' available metrics
         """
         # Try to get the auto_capture 'training_holdout' or the training
-        if capture_name == "latest":
-            metrics_df = self.get_inference_metrics("auto_inference")
-            return metrics_df if metrics_df is not None else self.get_inference_metrics("model_training")
+        if capture_name == "auto":
+            metric_list = self.list_inference_runs()
+            if metric_list:
+                return self.get_inference_metrics(metric_list[0])
+            else:
+                self.log.warning(f"No performance metrics found for {self.model_name}!")
+                return None
 
         # Grab the metrics captured during model training (could return None)
         if capture_name == "model_training":
@@ -331,11 +305,11 @@ class ModelCore(Artifact):
                 self.log.warning(f"Performance metrics {capture_name} not found for {self.model_name}!")
                 return None
 
-    def confusion_matrix(self, capture_name: str = "latest") -> Union[pd.DataFrame, None]:
+    def confusion_matrix(self, capture_name: str = "auto") -> Union[pd.DataFrame, None]:
         """Retrieve the confusion_matrix for this model
 
         Args:
-            capture_name (str, optional): Specific capture_name or "training" (default: "latest")
+            capture_name (str, optional): Specific capture_name or "training" (default: "auto")
         Returns:
             pd.DataFrame: DataFrame of the Confusion Matrix (might be None)
         """
@@ -347,7 +321,7 @@ class ModelCore(Artifact):
             raise ValueError(error_msg)
 
         # Grab the metrics from the Workbench Metadata (try inference first, then training)
-        if capture_name == "latest":
+        if capture_name == "auto":
             cm = self.confusion_matrix("auto_inference")
             return cm if cm is not None else self.confusion_matrix("model_training")
 
@@ -569,6 +543,17 @@ class ModelCore(Artifact):
         else:
             self.log.error(f"Model {self.model_name} is not a classifier!")
 
+    def summary(self) -> dict:
+        """Summary information about this Model
+
+        Returns:
+            dict: Dictionary of summary information about this Model
+        """
+        self.log.info("Computing Model Summary...")
+        summary = super().summary()
+        summary["hyperparameters"] = get_model_hyperparameters(self)
+        return summary
+
     def details(self) -> dict:
         """Additional Details about this Model
 
@@ -593,6 +578,7 @@ class ModelCore(Artifact):
         details["status"] = self.latest_model["ModelPackageStatus"]
         details["approval_status"] = self.latest_model.get("ModelApprovalStatus", "unknown")
         details["image"] = self.container_image().split("/")[-1]  # Shorten the image uri
+        details["hyperparameters"] = get_model_hyperparameters(self)
 
         # Grab the inference and container info
         inference_spec = self.latest_model["InferenceSpecification"]
@@ -603,16 +589,6 @@ class ModelCore(Artifact):
         details["transform_types"] = inference_spec["SupportedTransformInstanceTypes"]
         details["content_types"] = inference_spec["SupportedContentTypes"]
         details["response_types"] = inference_spec["SupportedResponseMIMETypes"]
-        details["model_metrics"] = self.get_inference_metrics()
-        if self.model_type == ModelType.CLASSIFIER:
-            details["confusion_matrix"] = self.confusion_matrix()
-            details["predictions"] = None
-        elif self.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
-            details["confusion_matrix"] = None
-            details["predictions"] = self.get_inference_predictions()
-        else:
-            details["confusion_matrix"] = None
-            details["predictions"] = None
 
         # Grab the inference metadata
         details["inference_meta"] = self.get_inference_metadata()
@@ -620,6 +596,24 @@ class ModelCore(Artifact):
         # Return the details
         return details
 
+    # Training View for this model
+    def training_view(self):
+        """Get the training view for this model"""
+        from workbench.core.artifacts.feature_set_core import FeatureSetCore
+        from workbench.core.views import View
+
+        # Grab our FeatureSet
+        fs = FeatureSetCore(self.get_input())
+
+        # See if we have a training view for this model
+        my_model_training_view = f"{self.name.replace('-', '_')}_training".lower()
+        view = View(fs, my_model_training_view, auto_create_view=False)
+        if view.exists():
+            return view
+        else:
+            self.log.important(f"No specific training view {my_model_training_view}, returning default training view")
+            return fs.view("training")
+
     # Pipeline for this model
     def get_pipeline(self) -> str:
         """Get the pipeline for this model"""
@@ -883,7 +877,7 @@ class ModelCore(Artifact):
             return self.df_store.get(f"/workbench/models/{self.name}/shap_data")
         else:
             # Loop over the SHAP data and return a dict of DataFrames
-            shap_dfs = self.df_store.list_subfiles(f"/workbench/models/{self.name}/shap_data")
+            shap_dfs = self.df_store.list(f"/workbench/models/{self.name}/shap_data")
             shap_data = {}
             for df_location in shap_dfs:
                 key = df_location.split("/")[-1]
@@ -902,10 +896,24 @@ class ModelCore(Artifact):
         except (KeyError, IndexError, TypeError):
             return None
 
+    def publish_prox_model(self, prox_model_name: str = None, include_all_columns: bool = False):
+        """Create and publish a Proximity Model for this Model
+
+        Args:
+            prox_model_name (str, optional): Name of the Proximity Model (if not specified, a name will be generated)
+            include_all_columns (bool): Include all DataFrame columns in results (default: False)
+
+        Returns:
+            Model: The published Proximity Model
+        """
+        if prox_model_name is None:
+            prox_model_name = self.model_name + "-prox"
+        return published_proximity_model(self, prox_model_name, include_all_columns=include_all_columns)
+
     def delete(self):
         """Delete the Model Packages and the Model Group"""
         if not self.exists():
-            self.log.warning(f"Trying to delete an Model that doesn't exist: {self.name}")
+            self.log.warning(f"Trying to delete a Model that doesn't exist: {self.name}")
 
         # Call the Class Method to delete the Model Group
         ModelCore.managed_delete(model_group_name=self.name)
@@ -981,6 +989,27 @@ class ModelCore(Artifact):
             self.log.warning(f"Could not determine model type for {self.model_name}!")
             return ModelType.UNKNOWN
 
+    def _set_model_framework(self, model_framework: ModelFramework):
+        """Internal: Set the Model Framework for this Model"""
+        self.model_framework = model_framework
+        self.upsert_workbench_meta({"workbench_model_framework": self.model_framework.value})
+        self.remove_health_tag("model_framework_unknown")
+
+    def _get_model_framework(self) -> ModelFramework:
+        """Internal: Query the Workbench Metadata to get the model framework
+        Returns:
+            ModelFramework: The ModelFramework of this Model
+        Notes:
+            This is an internal method that should not be called directly
+            Use the model_framework attribute instead
+        """
+        model_framework = self.workbench_meta().get("workbench_model_framework")
+        try:
+            return ModelFramework(model_framework)
+        except ValueError:
+            self.log.warning(f"Could not determine model framework for {self.model_name}!")
+            return ModelFramework.UNKNOWN
+
     def _load_training_metrics(self):
         """Internal: Retrieve the training metrics and Confusion Matrix for this model
                      and load the data into the Workbench Metadata
@@ -1172,13 +1201,11 @@ if __name__ == "__main__":
     # Grab a ModelCore object and pull some information from it
     my_model = ModelCore("abalone-regression")
 
-    # Call the various methods
-
     # Let's do a check/validation of the Model
     print(f"Model Check: {my_model.exists()}")
 
     # Make sure the model is 'ready'
-    # my_model.onboard()
+    my_model.onboard()
 
     # Get the ARN of the Model Group
     print(f"Model Group ARN: {my_model.group_arn()}")
@@ -1244,5 +1271,10 @@ if __name__ == "__main__":
     # Delete the Model
     # ModelCore.managed_delete("wine-classification")
 
+    # Check the training view logic
+    model = ModelCore("wine-class-test-251112-BW")
+    training_view = model.training_view()
+    print(f"Training View Name: {training_view.name}")
+
     # Check for a model that doesn't exist
     my_model = ModelCore("empty-model-group")

workbench/core/artifacts/parameter_store_core.py

@@ -0,0 +1,98 @@
+"""ParameterStoreCore: Manages Workbench parameters in a Cloud Based Parameter Store."""
+
+import logging
+
+# Workbench Imports
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+# Workbench Bridges Import
+from workbench_bridges.api import ParameterStore as BridgesParameterStore
+
+
+class ParameterStoreCore(BridgesParameterStore):
+    """ParameterStoreCore: Manages Workbench parameters in a Cloud Based Parameter Store.
+
+    Common Usage:
+        ```python
+        params = ParameterStoreCore()
+
+        # List Parameters
+        params.list()
+
+        ['/workbench/abalone_info',
+         '/workbench/my_data',
+         '/workbench/test',
+         '/workbench/pipelines/my_pipeline']
+
+        # Add Key
+        params.upsert("key", "value")
+        value = params.get("key")
+
+        # Add any data (lists, dictionaries, etc..)
+        my_data = {"key": "value", "number": 4.2, "list": [1,2,3]}
+        params.upsert("my_data", my_data)
+
+        # Retrieve data
+        return_value = params.get("my_data")
+        pprint(return_value)
+
+        {'key': 'value', 'list': [1, 2, 3], 'number': 4.2}
+
+        # Delete parameters
+        param_store.delete("my_data")
+        ```
+    """
+
+    def __init__(self):
+        """ParameterStoreCore Init Method"""
+        session = AWSAccountClamp().boto3_session
+
+        # Initialize parent with workbench config
+        super().__init__(boto3_session=session)
+        self.log = logging.getLogger("workbench")
+
+
+if __name__ == "__main__":
+    """Exercise the ParameterStoreCore Class"""
+
+    # Create a ParameterStoreCore manager
+    param_store = ParameterStoreCore()
+
+    # List the parameters
+    print("Listing Parameters...")
+    print(param_store.list())
+
+    # Add a new parameter
+    param_store.upsert("/workbench/test", "value")
+
+    # Get the parameter
+    print(f"Getting parameter 'test': {param_store.get('/workbench/test')}")
+
+    # Add a dictionary as a parameter
+    sample_dict = {"key": "str_value", "awesome_value": 4.2}
+    param_store.upsert("/workbench/my_data", sample_dict)
+
+    # Retrieve the parameter as a dictionary
+    retrieved_value = param_store.get("/workbench/my_data")
+    print("Retrieved value:", retrieved_value)
+
+    # List the parameters
+    print("Listing Parameters...")
+    print(param_store.list())
+
+    # List the parameters with a prefix
+    print("Listing Parameters with prefix '/workbench':")
+    print(param_store.list("/workbench"))
+
+    # Delete the parameters
+    param_store.delete("/workbench/test")
+    param_store.delete("/workbench/my_data")
+
+    # Out of scope tests
+    param_store.upsert("test", "value")
+    param_store.delete("test")
+
+    # Recursive delete test
+    param_store.upsert("/workbench/test/test1", "value1")
+    param_store.upsert("/workbench/test/test2", "value2")
+    param_store.delete_recursive("workbench/test/")

workbench/core/cloud_platform/aws/aws_account_clamp.py

@@ -55,9 +55,10 @@ class AWSAccountClamp:
         # Check our Assume Role
         self.log.info("Checking Workbench Assumed Role...")
         role_info = self.aws_session.assumed_role_info()
+        self.log.info(f"Assumed Role: {role_info}")
 
-        # Check if the Role is a 'ReadOnly' role
-        self.read_only_role = "readonly" in role_info["AssumedRoleArn"].lower()
+        # Check if we have tag write permissions (if we don't, we are read-only)
+        self.read_only = not self.check_tag_permissions()
 
         # Check our Workbench API Key and Load the License
         self.log.info("Checking Workbench API License...")
@@ -141,6 +142,45 @@ class AWSAccountClamp:
         """
         return self.boto3_session.client("sagemaker")
 
+    def check_tag_permissions(self):
+        """Check if current role has permission to add tags to SageMaker endpoints.
+
+        Returns:
+            bool: True if AddTags is allowed, False otherwise
+        """
+        try:
+            sagemaker = self.boto3_session.client("sagemaker")
+
+            # Use a non-existent endpoint name
+            fake_endpoint = "workbench-permission-check-dummy-endpoint"
+
+            # Try to add tags to the non-existent endpoint
+            sagemaker.add_tags(
+                ResourceArn=f"arn:aws:sagemaker:{self.region}:{self.account_id}:endpoint/{fake_endpoint}",
+                Tags=[{"Key": "PermissionCheck", "Value": "Test"}],
+            )
+
+            # If we get here, we have permission (but endpoint doesn't exist)
+            return True
+
+        except ClientError as e:
+            error_code = e.response["Error"]["Code"]
+
+            # AccessDeniedException = no permission
+            if error_code == "AccessDeniedException":
+                self.log.debug("No AddTags permission (AccessDeniedException)")
+                return False
+
+            # ResourceNotFound = we have permission, but endpoint doesn't exist
+            elif error_code in ["ResourceNotFound", "ValidationException"]:
+                self.log.debug("AddTags permission verified (resource not found)")
+                return True
+
+            # Unexpected error, assume no permission for safety
+            else:
+                self.log.debug(f"Unexpected error checking permissions: {error_code}")
+                return False
+
 
 if __name__ == "__main__":
     """Exercise the AWS Account Clamp Class"""
@@ -165,3 +205,9 @@ if __name__ == "__main__":
     print("\n\n*** AWS Sagemaker Session/Client Check ***")
     sm_client = aws_account_clamp.sagemaker_client()
     print(sm_client.list_feature_groups()["FeatureGroupSummaries"])
+
+    print("\n\n*** AWS Tag Permission Check ***")
+    if aws_account_clamp.check_tag_permissions():
+        print("Tag Permission Check Success...")
+    else:
+        print("Tag Permission Check Failed...")

workbench/core/cloud_platform/cloud_meta.py

@@ -7,7 +7,6 @@ import logging
 from typing import Union
 import pandas as pd
 
-
 # Workbench Imports
 from workbench.core.cloud_platform.aws.aws_meta import AWSMeta
 

workbench/core/pipelines/pipeline_executor.py

@@ -123,7 +123,7 @@ class PipelineExecutor:
                 if "model" in workbench_objects and (not subset or "endpoint" in subset):
                     workbench_objects["model"].to_endpoint(**kwargs)
                     endpoint = Endpoint(kwargs["name"])
-                    endpoint.auto_inference(capture=True)
+                    endpoint.auto_inference()
 
             # Found something weird
             else: