workbench 0.8.169__tar.gz → 0.8.170__tar.gz

{workbench-0.8.169 → workbench-0.8.170}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.169
+Version: 0.8.170
 Summary: Workbench: A Dashboard and Python API for creating and deploying AWS SageMaker Model Pipelines
 Author-email: SuperCowPowers LLC <support@supercowpowers.com>
 License-Expression: MIT

{workbench-0.8.169 → workbench-0.8.170}/aws_setup/workbench_core/workbench_core/workbench_core_stack.py

@@ -79,7 +79,7 @@ class WorkbenchCoreStack(Stack):
         self.model_policy = self.workbench_model_policy()
         self.endpoint_read_policy = self.workbench_endpoint_read_policy()
         self.endpoint_policy = self.workbench_endpoint_policy()
-        self.pipeline_policy = self.workbench_pipeline_policy()
+        # self.sagemaker_pipeline_full_policy = self.sagemaker_pipeline_full_policy()  # Not currently used
         self.dataframe_store_read_policy = self.workbench_dataframe_store_read_policy()
         self.dataframe_store_full_policy = self.workbench_dataframe_store_full_policy()
         self.parameter_store_read_policy = self.workbench_parameter_store_read_policy()
@@ -235,6 +235,16 @@ class WorkbenchCoreStack(Stack):
     #####################
     #     Glue Jobs     #
     #####################
+    def glue_jobs_s3_read(self) -> iam.PolicyStatement:
+        """S3 Read for Glue Jobs default script location"""
+        return iam.PolicyStatement(
+            actions=["s3:GetObject", "s3:GetBucketLocation", "s3:ListBucket"],
+            resources=[
+                f"arn:aws:s3:::aws-glue-assets-{self.account}-{self.region}",
+                f"arn:aws:s3:::aws-glue-assets-{self.account}-{self.region}/*",
+            ],
+        )
+
     def glue_job_logs(self) -> iam.PolicyStatement:
         """Create a policy statement for Glue job CloudWatch logs.
 
@@ -410,7 +420,6 @@ class WorkbenchCoreStack(Stack):
             compute_environment_name="workbench-compute-env",
             vpc=vpc,
             vpc_subnets=vpc_subnets,
-            # replace_compute_environment=True,
         )
 
     def create_batch_job_queue(self) -> batch.JobQueue:
@@ -484,11 +493,13 @@ class WorkbenchCoreStack(Stack):
     #####################
     #  Inference Store  #
     #####################
-    def inference_store_read(self) -> iam.PolicyStatement:
+    def glue_database_read_just_inference_store(self) -> iam.PolicyStatement:
         """Create a policy statement for reading from the Parameter Store.
 
+        Note: This is basically glue_databases_read but scoped to just the inference_store database.
+
         Returns:
-            iam.PolicyStatement: The policy statement for reading from the Parameter Store.
+            iam.PolicyStatement: The policy statement for reading from the Inference Store.
         """
         return iam.PolicyStatement(
             actions=[
@@ -501,13 +512,15 @@ class WorkbenchCoreStack(Stack):
             resources=self._inference_database_arns(),
         )
 
-    def inference_store_full(self) -> iam.PolicyStatement:
+    def glue_database_full_just_inference_store(self) -> iam.PolicyStatement:
         """Create a policy statement for full access to the Inference Store.
 
+        Note: This is basically glue_databases_full but scoped to just the inference_store database.
+
         Returns:
             iam.PolicyStatement: The policy statement for full access to the Inference Store.
         """
-        read_statement = self.inference_store_read()
+        read_statement = self.glue_database_read_just_inference_store()
         return iam.PolicyStatement(
             actions=read_statement.actions
             + [
@@ -515,7 +528,19 @@ class WorkbenchCoreStack(Stack):
                 "glue:UpdateTable",
                 "glue:DeleteTable",
             ],
-            resources=self._inference_database_arns(),
+            resources=read_statement.resources,
+        )
+
+    def s3_full_just_inference_store(self) -> iam.PolicyStatement:
+        """S3 write permissions for inference store data only."""
+        return iam.PolicyStatement(
+            actions=[
+                "s3:PutObject",  # Write inference store data
+                "s3:DeleteObject",  # Delete inference store data
+            ],
+            resources=[
+                f"arn:aws:s3:::{self.workbench_bucket}/athena/inference_store/inference_store/*",
+            ],
         )
 
     ####################
@@ -579,6 +604,20 @@ class WorkbenchCoreStack(Stack):
             resources=["*"],
         )
 
+    def athena_query_results_s3(self) -> iam.PolicyStatement:
+        """S3 permissions for Athena query results bucket."""
+        return iam.PolicyStatement(
+            actions=[
+                "s3:PutObject",  # Required: Athena writes query results even for SELECT queries
+                "s3:GetObject",  # Needed to retrieve results later
+                "s3:ListBucket",  # Needed for bucket operations
+            ],
+            resources=[
+                f"arn:aws:s3:::aws-athena-query-results-{self.account}-{self.region}",
+                f"arn:aws:s3:::aws-athena-query-results-{self.account}-{self.region}/*",
+            ],
+        )
+
     ######################
     #    FeatureStore    #
     ######################
@@ -852,7 +891,7 @@ class WorkbenchCoreStack(Stack):
         )
 
     @staticmethod
-    def pipeline_list_policy_statement() -> iam.PolicyStatement:
+    def pipeline_list() -> iam.PolicyStatement:
         """Create a policy statement for listing SageMaker pipelines.
 
         Returns:
@@ -865,7 +904,7 @@ class WorkbenchCoreStack(Stack):
             resources=["*"],  # Broad permission necessary for listing operations
         )
 
-    def pipeline_policy_statement(self) -> iam.PolicyStatement:
+    def pipeline_full(self) -> iam.PolicyStatement:
         """Create a policy statement for inspecting and running SageMaker Pipelines.
 
         Returns:
@@ -899,7 +938,8 @@ class WorkbenchCoreStack(Stack):
             ],
         )
 
-    def ecr_policy_statement(self) -> iam.PolicyStatement:
+    @staticmethod
+    def ecr_policy_statement() -> iam.PolicyStatement:
         """Create a policy statement for pulling ECR Images.
 
         Returns:
@@ -989,7 +1029,8 @@ class WorkbenchCoreStack(Stack):
     ##########################
     #   Workbench Dashboard  #
     ##########################
-    def dashboard_policy_statement(self) -> iam.PolicyStatement:
+    @staticmethod
+    def dashboard_policy_statement() -> iam.PolicyStatement:
         """Create a policy statement for additional permissions needed by Workbench Dashboard.
 
         Returns:
@@ -1230,22 +1271,27 @@ class WorkbenchCoreStack(Stack):
             managed_policy_name="WorkbenchEndpointPolicy",
         )
 
-    def workbench_pipeline_policy(self) -> iam.ManagedPolicy:
-        """Create a managed policy for the Workbench Pipelines"""
+    def sagemaker_full_policy(self) -> iam.ManagedPolicy:
+        """Create a managed policy for the Workbench Pipelines (FULL)"""
         policy_statements = [
-            self.pipeline_list_policy_statement(),
-            self.pipeline_policy_statement(),
+            self.pipeline_list(),
+            self.pipeline_full(),
         ]
         return iam.ManagedPolicy(
             self,
-            id="WorkbenchPipelinePolicy",
+            id="WorkbenchSagemakerPipelineFullPolicy",
             statements=policy_statements,
-            managed_policy_name="WorkbenchPipelinePolicy",
+            managed_policy_name="WorkbenchSagemakerPipelineFullPolicy",
         )
 
     def workbench_dataframe_store_read_policy(self) -> iam.ManagedPolicy:
         """Create a managed policy for the Workbench DataFrame Store (READ-ONLY)"""
         policy_statements = [
+            self.s3_read(),
+            self.glue_jobs_s3_read(),
+            self.glue_job_logs(),
+            self.cloudwatch_logs(),
+            self.parameter_store_read(),  # Get Workbench Bucket from Parameter Store
             self.dataframe_store_read(),
         ]
 
@@ -1259,6 +1305,11 @@ class WorkbenchCoreStack(Stack):
     def workbench_dataframe_store_full_policy(self) -> iam.ManagedPolicy:
         """Create a managed policy for the Workbench DataFrame Store (FULL)"""
         policy_statements = [
+            self.s3_read(),
+            self.glue_jobs_s3_read(),
+            self.glue_job_logs(),
+            self.cloudwatch_logs(),
+            self.parameter_store_read(),  # Get Workbench Bucket from Parameter Store
             self.dataframe_store_full(),
         ]
 
@@ -1266,12 +1317,16 @@ class WorkbenchCoreStack(Stack):
             self,
             id="WorkbenchDFStoreFullPolicy",
             statements=policy_statements,
-            managed_policy_name="WorkbenchDFStorePolicy",
+            managed_policy_name="WorkbenchDFStoreFullPolicy",
         )
 
     def workbench_parameter_store_read_policy(self) -> iam.ManagedPolicy:
         """Create a managed policy for the Workbench Parameter Store (READ-ONLY)"""
         policy_statements = [
+            self.s3_read(),
+            self.glue_jobs_s3_read(),
+            self.glue_job_logs(),
+            self.cloudwatch_logs(),
             self.parameter_store_discover(),
             self.parameter_store_read(),
         ]
@@ -1286,6 +1341,10 @@ class WorkbenchCoreStack(Stack):
     def workbench_parameter_store_full_policy(self) -> iam.ManagedPolicy:
         """Create a managed policy for the Workbench Parameter Store (FULL)"""
         policy_statements = [
+            self.s3_read(),
+            self.glue_jobs_s3_read(),
+            self.glue_job_logs(),
+            self.cloudwatch_logs(),
             self.parameter_store_discover(),
             self.parameter_store_full(),
         ]
@@ -1294,13 +1353,20 @@ class WorkbenchCoreStack(Stack):
             self,
             id="WorkbenchParameterStoreFullPolicy",
             statements=policy_statements,
-            managed_policy_name="WorkbenchParameterStorePolicy",
+            managed_policy_name="WorkbenchParameterStoreFullPolicy",
         )
 
     def workbench_inference_store_read_policy(self) -> iam.ManagedPolicy:
         """Create a managed policy for the Workbench Inference Store (READ-ONLY)"""
         policy_statements = [
-            self.inference_store_read(),
+            self.s3_read(),
+            self.glue_jobs_s3_read(),
+            self.athena_query_results_s3(),
+            self.glue_job_logs(),
+            self.cloudwatch_logs(),
+            self.glue_catalog_read(),
+            self.glue_database_read_just_inference_store(),
+            self.athena_read(),
         ]
 
         return iam.ManagedPolicy(
@@ -1313,14 +1379,23 @@ class WorkbenchCoreStack(Stack):
     def workbench_inference_store_full_policy(self) -> iam.ManagedPolicy:
         """Create a managed policy for the Workbench Inference Store (FULL)"""
         policy_statements = [
-            self.inference_store_full(),
+            self.s3_read(),
+            self.glue_jobs_s3_read(),
+            self.s3_full_just_inference_store(),
+            self.athena_query_results_s3(),
+            self.glue_job_logs(),
+            self.cloudwatch_logs(),
+            self.glue_catalog_read(),
+            self.glue_database_full_just_inference_store(),
+            self.athena_read(),
+            self.parameter_store_read(),  # Get Workbench Bucket from Parameter Store
         ]
 
         return iam.ManagedPolicy(
             self,
             id="WorkbenchInferenceStoreFullPolicy",
             statements=policy_statements,
-            managed_policy_name="WorkbenchInferenceStorePolicy",
+            managed_policy_name="WorkbenchInferenceStoreFullPolicy",
         )
 
     def create_execution_role(self) -> iam.Role:
@@ -1352,12 +1427,12 @@ class WorkbenchCoreStack(Stack):
         api_execution_role.add_to_policy(self.batch_pass_role())
         api_execution_role.add_to_policy(self.parameter_store_discover())
         api_execution_role.add_to_policy(self.parameter_store_full())
+        api_execution_role.add_to_policy(self.cloudwatch_logs())
         api_execution_role.add_to_policy(self.cloudwatch_monitor())
         api_execution_role.add_managed_policy(self.datasource_policy)
         api_execution_role.add_managed_policy(self.featureset_policy)
         api_execution_role.add_managed_policy(self.model_policy)
         api_execution_role.add_managed_policy(self.endpoint_policy)
-        api_execution_role.add_managed_policy(self.pipeline_policy)
         return api_execution_role
 
     def create_readonly_role(self) -> iam.Role:
@@ -1382,6 +1457,7 @@ class WorkbenchCoreStack(Stack):
         readonly_role.add_to_policy(self.glue_jobs_read())
         readonly_role.add_to_policy(self.parameter_store_discover())
         readonly_role.add_to_policy(self.parameter_store_read())
+        readonly_role.add_to_policy(self.cloudwatch_logs())
         readonly_role.add_managed_policy(self.datasource_read_policy)
         readonly_role.add_managed_policy(self.featureset_read_policy)
         readonly_role.add_managed_policy(self.model_read_policy)
@@ -1400,11 +1476,11 @@ class WorkbenchCoreStack(Stack):
 
         # Add a subset of policies for the Lambda Role
         lambda_role.add_to_policy(self.parameter_store_full())
+        lambda_role.add_to_policy(self.cloudwatch_logs())
         lambda_role.add_managed_policy(self.datasource_policy)
         lambda_role.add_managed_policy(self.featureset_policy)
         lambda_role.add_managed_policy(self.model_policy)
         lambda_role.add_managed_policy(self.endpoint_policy)
-        lambda_role.add_managed_policy(self.pipeline_policy)
         return lambda_role
 
     def create_glue_role(self) -> iam.Role:
@@ -1420,6 +1496,7 @@ class WorkbenchCoreStack(Stack):
         # Add a subset of policies for the Glue Role
         glue_role.add_to_policy(self.glue_job_logs())
         glue_role.add_to_policy(self.glue_connections())
+        glue_role.add_to_policy(self.cloudwatch_logs())
         glue_role.add_to_policy(self.vpc_discovery())
         glue_role.add_to_policy(self.vpc_network_interface_management())
         glue_role.add_to_policy(self.parameter_store_full())
@@ -1427,7 +1504,6 @@ class WorkbenchCoreStack(Stack):
         glue_role.add_managed_policy(self.featureset_policy)
         glue_role.add_managed_policy(self.model_policy)
         glue_role.add_managed_policy(self.endpoint_policy)
-        glue_role.add_managed_policy(self.pipeline_policy)
         return glue_role
 
     def create_batch_role(self) -> iam.Role:
@@ -1442,13 +1518,13 @@ class WorkbenchCoreStack(Stack):
 
         # Add policies for the Batch Role
         batch_role.add_to_policy(self.batch_job_logs())
+        batch_role.add_to_policy(self.cloudwatch_logs())
         batch_role.add_to_policy(self.parameter_store_full())
         batch_role.add_to_policy(self.dataframe_store_full())
         batch_role.add_managed_policy(self.datasource_policy)
         batch_role.add_managed_policy(self.featureset_policy)
         batch_role.add_managed_policy(self.model_policy)
         batch_role.add_managed_policy(self.endpoint_policy)
-        batch_role.add_managed_policy(self.pipeline_policy)
         return batch_role
 
     def _create_sso_principals(self, base_principals: iam.CompositePrincipal = None) -> iam.CompositePrincipal:

workbench-0.8.170/examples/ml_pipelines/list_models.py

@@ -0,0 +1,14 @@
+import time
+import logging
+from workbench.api import Meta
+
+log = logging.getLogger("workbench")
+
+# List all the models in AWS
+meta = Meta()
+models = meta.models()
+log.info(f"Found {len(models)} models in AWS")
+log.info(models)
+
+# Test out the ml launcher disconnect
+time.sleep(120)

{workbench-0.8.169 → workbench-0.8.170}/src/workbench/api/model.py

@@ -40,6 +40,7 @@ class Model(ModelCore):
         mem_size: int = 2048,
         max_concurrency: int = 5,
         instance: str = "ml.t2.medium",
+        data_capture: bool = False,
     ) -> Endpoint:
         """Create an Endpoint from the Model.
 
@@ -50,6 +51,7 @@ class Model(ModelCore):
             mem_size (int): The memory size for the Endpoint in MB (default: 2048)
             max_concurrency (int): The maximum concurrency for the Endpoint (default: 5)
             instance (str): The instance type to use for Realtime(serverless=False) Endpoints (default: "ml.t2.medium")
+            data_capture (bool): Enable data capture for the Endpoint (default: False)
 
         Returns:
             Endpoint: The Endpoint created from the Model
@@ -73,6 +75,7 @@ class Model(ModelCore):
         model_to_endpoint.transform(
             mem_size=mem_size,
             max_concurrency=max_concurrency,
+            data_capture=data_capture,
         )
 
         # Set the Endpoint Owner and Return the Endpoint

{workbench-0.8.169 → workbench-0.8.170}/src/workbench/core/artifacts/endpoint_core.py

@@ -972,12 +972,23 @@ class EndpointCore(Artifact):
             cls.log.info(f"Deleting Monitoring Schedule {schedule['MonitoringScheduleName']}...")
             cls.sm_client.delete_monitoring_schedule(MonitoringScheduleName=schedule["MonitoringScheduleName"])
 
-        # Recursively delete all endpoint S3 artifacts (inference, data capture, monitoring, etc)
+        # Recursively delete all endpoint S3 artifacts (inference, etc)
+        # Note: We do not want to delete the data_capture/ files since these
+        #       might be used for collection and data drift analysis
         base_endpoint_path = f"{cls.endpoints_s3_path}/{endpoint_name}"
-        s3_objects = wr.s3.list_objects(base_endpoint_path, boto3_session=cls.boto3_session)
-        cls.log.info(f"Deleting S3 Objects at {base_endpoint_path}...")
-        cls.log.info(f"{s3_objects}")
-        wr.s3.delete_objects(s3_objects, boto3_session=cls.boto3_session)
+        all_s3_objects = wr.s3.list_objects(base_endpoint_path, boto3_session=cls.boto3_session)
+
+        # Filter out objects that contain 'data_capture/' in their path
+        s3_objects_to_delete = [obj for obj in all_s3_objects if "/data_capture/" not in obj]
+        cls.log.info(f"Found {len(all_s3_objects)} total objects at {base_endpoint_path}")
+        cls.log.info(f"Filtering out data_capture files, will delete {len(s3_objects_to_delete)} objects...")
+        cls.log.info(f"Objects to delete: {s3_objects_to_delete}")
+
+        if s3_objects_to_delete:
+            wr.s3.delete_objects(s3_objects_to_delete, boto3_session=cls.boto3_session)
+            cls.log.info(f"Successfully deleted {len(s3_objects_to_delete)} objects")
+        else:
+            cls.log.info("No objects to delete (only data_capture files found)")
 
         # Delete any dataframes that were stored in the Dataframe Cache
         cls.log.info("Deleting Dataframe Cache...")

{workbench-0.8.169 → workbench-0.8.170}/src/workbench/core/artifacts/monitor_core.py

@@ -186,11 +186,11 @@ class MonitorCore:
 
         # Log the data capture operation
         self.log.important(f"Enabling Data Capture for {self.endpoint_name} --> {self.data_capture_path}")
-        self.log.important("This normally redeploys the endpoint...")
+        self.log.important("This will redeploy the endpoint...")
 
         # Create and apply the data capture configuration
         data_capture_config = DataCaptureConfig(
-            enable_capture=True,  # Required parameter
+            enable_capture=True,
             sampling_percentage=capture_percentage,
             destination_s3_uri=self.data_capture_path,
         )

{workbench-0.8.169 → workbench-0.8.170}/src/workbench/core/transforms/model_to_endpoint/model_to_endpoint.py

@@ -5,6 +5,7 @@ from sagemaker import ModelPackage
 from sagemaker.serializers import CSVSerializer
 from sagemaker.deserializers import CSVDeserializer
 from sagemaker.serverless import ServerlessInferenceConfig
+from sagemaker.model_monitor import DataCaptureConfig
 
 # Local Imports
 from workbench.core.transforms.transform import Transform, TransformInput, TransformOutput
@@ -51,27 +52,38 @@ class ModelToEndpoint(Transform):
         EndpointCore.managed_delete(self.output_name)
 
         # Get the Model Package ARN for our input model
-        input_model = ModelCore(self.input_name)
-        model_package_arn = input_model.model_package_arn()
+        workbench_model = ModelCore(self.input_name)
 
         # Deploy the model
-        self._deploy_model(model_package_arn, **kwargs)
+        self._deploy_model(workbench_model, **kwargs)
 
         # Add this endpoint to the set of registered endpoints for the model
-        input_model.register_endpoint(self.output_name)
+        workbench_model.register_endpoint(self.output_name)
 
         # This ensures that the endpoint is ready for use
         time.sleep(5)  # We wait for AWS Lag
         end = EndpointCore(self.output_name)
         self.log.important(f"Endpoint {end.name} is ready for use")
 
-    def _deploy_model(self, model_package_arn: str, mem_size: int = 2048, max_concurrency: int = 5):
+    def _deploy_model(
+        self,
+        workbench_model: ModelCore,
+        mem_size: int = 2048,
+        max_concurrency: int = 5,
+        data_capture: bool = False,
+        capture_percentage: int = 100,
+    ):
         """Internal Method: Deploy the Model
 
         Args:
-            model_package_arn(str): The Model Package ARN used to deploy the Endpoint
+            workbench_model(ModelCore): The Workbench ModelCore object to deploy
+            mem_size(int): Memory size for serverless deployment
+            max_concurrency(int): Max concurrency for serverless deployment
+            data_capture(bool): Enable data capture during deployment
+            capture_percentage(int): Percentage of data to capture. Defaults to 100.
         """
         # Grab the specified Model Package
+        model_package_arn = workbench_model.model_package_arn()
         model_package = ModelPackage(
             role=self.workbench_role_arn,
             model_package_arn=model_package_arn,
@@ -95,6 +107,23 @@ class ModelToEndpoint(Transform):
                 max_concurrency=max_concurrency,
             )
 
+        # Configure data capture if requested (and not serverless)
+        data_capture_config = None
+        if data_capture and not self.serverless:
+            # Set up the S3 path for data capture
+            base_endpoint_path = f"{workbench_model.endpoints_s3_path}/{self.output_name}"
+            data_capture_path = f"{base_endpoint_path}/data_capture"
+            self.log.important(f"Configuring Data Capture --> {data_capture_path}")
+            data_capture_config = DataCaptureConfig(
+                enable_capture=True,
+                sampling_percentage=capture_percentage,
+                destination_s3_uri=data_capture_path,
+            )
+        elif data_capture and self.serverless:
+            self.log.warning(
+                "Data capture is not supported for serverless endpoints. Skipping data capture configuration."
+            )
+
         # Deploy the Endpoint
         self.log.important(f"Deploying the Endpoint {self.output_name}...")
         model_package.deploy(
@@ -104,6 +133,7 @@ class ModelToEndpoint(Transform):
             endpoint_name=self.output_name,
             serializer=CSVSerializer(),
             deserializer=CSVDeserializer(),
+            data_capture_config=data_capture_config,
             tags=aws_tags,
         )
 

{workbench-0.8.169 → workbench-0.8.170}/src/workbench.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.169
+Version: 0.8.170
 Summary: Workbench: A Dashboard and Python API for creating and deploying AWS SageMaker Model Pipelines
 Author-email: SuperCowPowers LLC <support@supercowpowers.com>
 License-Expression: MIT

{workbench-0.8.169 → workbench-0.8.170}/src/workbench.egg-info/SOURCES.txt

@@ -262,6 +262,7 @@ examples/glue_jobs/glue_launcher_examples/example_2.py
 examples/meta/meta_list_endpoints.py
 examples/meta/meta_list_models.py
 examples/meta/meta_model_metrics.py
+examples/ml_pipelines/list_models.py
 examples/model_utils/model_metrics.py
 examples/model_utils/model_to_endpoint.py
 examples/model_utils/onboard_model.py