workbench 0.8.162__py3-none-any.whl → 0.8.220__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,69 +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"
+
+
+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", "xgb_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", "xgb_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-sklearn-xgb-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", "xgb_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", "xgb_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-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
-        ("us-east-1", "xgb_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", "xgb_inference", "0.1", "arm64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1-arm64"
-        ),
-        # 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):
@@ -105,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
         """
 
@@ -143,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
@@ -272,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":
@@ -308,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)
         """
@@ -324,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")
 
@@ -546,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
 
@@ -570,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"]
@@ -580,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()
@@ -597,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"""
@@ -860,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]
@@ -879,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)
@@ -958,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
@@ -1149,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()}")
@@ -1221,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")