workbench 0.8.168__tar.gz → 0.8.170__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.168
+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.168 → workbench-0.8.170}/aws_setup/workbench_core/app.py

@@ -31,7 +31,12 @@ workbench_bucket = (cm and cm.get_config("WORKBENCH_BUCKET")) or os.getenv("WORK
 if not workbench_bucket:
     print("Error: WORKBENCH_BUCKET is required but not found in config or environment variables.")
     exit(1)
-sso_group = (cm and cm.get_config("WORKBENCH_SSO_GROUP")) or os.getenv("WORKBENCH_SSO_GROUP")
+
+# SSO Groups (Groups that can assume the Workbench roles)
+sso_groups_str = (cm and cm.get_config("WORKBENCH_SSO_GROUPS")) or os.getenv("WORKBENCH_SSO_GROUPS", "")
+sso_groups = [group.strip() for group in sso_groups_str.split(",") if group.strip()] if sso_groups_str else []
+
+# Additional Buckets (Extra S3 buckets to grant access to)
 additional_buckets_str = (cm and cm.get_config("WORKBENCH_ADDITIONAL_BUCKETS")) or os.getenv(
     "WORKBENCH_ADDITIONAL_BUCKETS", ""
 )
@@ -44,7 +49,7 @@ subnet_ids = (cm and cm.get_config("WORKBENCH_SUBNET_IDS")) or os.getenv("WORKBE
 # Log the configuration for transparency
 print("Configuration:")
 print(f"  WORKBENCH_BUCKET: {workbench_bucket}")
-print(f"  WORKBENCH_SSO_GROUP: {sso_group}")
+print(f"  WORKBENCH_SSO_GROUPS: {sso_groups}")
 print(f"  WORKBENCH_ADDITIONAL_BUCKETS: {additional_buckets}")
 print(f"  WORKBENCH_VPC_ID: {existing_vpc_id}")
 print(f"  WORKBENCH_SUBNET_IDS: {subnet_ids}")
@@ -60,7 +65,7 @@ sandbox_stack = WorkbenchCoreStack(
     env=env,
     props=WorkbenchCoreStackProps(
         workbench_bucket=workbench_bucket,
-        sso_group=sso_group,
+        sso_groups=sso_groups,
         additional_buckets=additional_buckets,
         existing_vpc_id=existing_vpc_id,
         subnet_ids=subnet_ids,

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

@@ -4,18 +4,22 @@ from aws_cdk import (
     aws_iam as iam,
     aws_logs as logs,
     aws_ec2 as ec2,
+    aws_ecs as ecs,
     aws_batch as batch,
     RemovalPolicy,
+    Duration,
+    Size,
 )
 from constructs import Construct
 from typing import Any, List
 from dataclasses import dataclass, field
+from typing import Dict
 
 
 @dataclass
 class WorkbenchCoreStackProps:
     workbench_bucket: str
-    sso_group: str
+    sso_groups: List[str] = field(default_factory=list)
     additional_buckets: List[str] = field(default_factory=list)
     existing_vpc_id: str = None
     subnet_ids: List[str] = field(default_factory=list)  # Optional subnets ids for the Batch compute environment
@@ -39,7 +43,7 @@ class WorkbenchCoreStack(Stack):
 
         # Grab our properties
         self.workbench_bucket = props.workbench_bucket
-        self.sso_group = props.sso_group
+        self.sso_groups = props.sso_groups
         self.additional_buckets = props.additional_buckets
         self.existing_vpc_id = props.existing_vpc_id
         self.subnet_ids = props.subnet_ids
@@ -75,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()
@@ -97,6 +101,7 @@ class WorkbenchCoreStack(Stack):
         # Batch Compute Environment and Job Queue
         self.batch_compute_environment = self.create_batch_compute_environment()
         self.batch_job_queue = self.create_batch_job_queue()
+        self.batch_job_definitions = self.create_batch_job_definitions()
 
     ####################
     #    S3 Buckets    #
@@ -230,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.
 
@@ -405,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:
@@ -419,6 +433,40 @@ class WorkbenchCoreStack(Stack):
             ],
         )
 
+    def create_batch_job_definitions(self) -> Dict[str, batch.EcsJobDefinition]:
+        """Create ML pipeline job definitions in small/medium/large tiers."""
+
+        ecr_image_uri = f"507740646243.dkr.ecr.{self.region}.amazonaws.com/aws-ml-images/py312-ml-pipelines:0.1"
+        tiers = {
+            "small": (2, 4096),  # 2 vCPU, 4GB RAM
+            "medium": (4, 8192),  # 4 vCPU, 8GB RAM
+            "large": (8, 16384),  # 8 vCPU, 16GB RAM
+        }
+
+        job_definitions = {}
+        for size, (cpu, memory_mib) in tiers.items():
+            job_definitions[size] = batch.EcsJobDefinition(
+                self,
+                f"JobDef{size.capitalize()}",
+                job_definition_name=f"workbench-ml-pipeline-{size}",
+                container=batch.EcsFargateContainerDefinition(
+                    self,
+                    f"Container{size.capitalize()}",
+                    cpu=cpu,
+                    memory=Size.mebibytes(memory_mib),
+                    image=ecs.ContainerImage.from_registry(ecr_image_uri),
+                    job_role=self.workbench_batch_role,
+                    execution_role=self.workbench_batch_role,
+                    environment={
+                        "WORKBENCH_BUCKET": self.workbench_bucket,
+                        "PYTHONUNBUFFERED": "1",
+                    },
+                ),
+                timeout=Duration.hours(3),
+            )
+
+        return job_definitions
+
     #####################
     #  DataFrame Store  #
     #####################
@@ -445,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=[
@@ -462,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
             + [
@@ -476,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/*",
+            ],
         )
 
     ####################
@@ -540,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    #
     ######################
@@ -813,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:
@@ -826,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:
@@ -860,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:
@@ -950,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:
@@ -1191,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(),
         ]
 
@@ -1220,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(),
         ]
 
@@ -1227,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(),
         ]
@@ -1247,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(),
         ]
@@ -1255,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(
@@ -1274,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:
@@ -1313,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:
@@ -1343,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)
@@ -1361,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:
@@ -1381,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())
@@ -1388,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:
@@ -1403,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:
@@ -1424,26 +1539,36 @@ class WorkbenchCoreStack(Stack):
         """
         assumed_by = base_principals or iam.CompositePrincipal()
 
-        # If sso_group is provided, configure trust relationship for AWS SSO integration
+        # If sso_groups are provided, configure trust relationship for AWS SSO integration
         # AWS SSO creates roles with two different ARN patterns depending on configuration:
-        if self.sso_group:
-
-            # Pattern 1: Direct group-based SSO role (no permission set in path)
-            # Used when SSO groups are mapped directly to roles
-            sso_group_arn_1 = (
-                f"arn:aws:iam::{self.account}:role/aws-reserved/sso.amazonaws.com/AWSReservedSSO_{self.sso_group}_*"
-            )
-
-            # Pattern 2: Permission set-based SSO role (includes permission set ID in path)
-            # Used when SSO groups are assigned via permission sets (more common in enterprise setups)
-            # The middle wildcard (*) represents the permission set identifier
-            sso_group_arn_2 = (
-                f"arn:aws:iam::{self.account}:role/aws-reserved/sso.amazonaws.com/*/AWSReservedSSO_{self.sso_group}_*"
-            )
-
+        if self.sso_groups:
+            # Build lists to hold all ARN patterns for all groups
+            sso_group_arns_pattern_1 = []
+            sso_group_arns_pattern_2 = []
+
+            # Generate ARN patterns for each SSO group
+            for sso_group in self.sso_groups:
+                # Pattern 1: Direct group-based SSO role (no permission set in path)
+                # Used when SSO groups are mapped directly to roles
+                sso_group_arn_1 = (
+                    f"arn:aws:iam::{self.account}:role/aws-reserved/sso.amazonaws.com/AWSReservedSSO_{sso_group}_*"
+                )
+                sso_group_arns_pattern_1.append(sso_group_arn_1)
+
+                # Pattern 2: Permission set-based SSO role (includes permission set ID in path)
+                # Used when SSO groups are assigned via permission sets (more common in enterprise setups)
+                # The middle wildcard (*) represents the permission set identifier
+                sso_group_arn_2 = (
+                    f"arn:aws:iam::{self.account}:role/aws-reserved/sso.amazonaws.com/*/AWSReservedSSO_{sso_group}_*"
+                )
+                sso_group_arns_pattern_2.append(sso_group_arn_2)
+
+            # Combine both patterns into a single list for the condition
             # Both patterns are required for maximum compatibility across different SSO configurations
             # This is particularly important for Lake Formation and enterprise SSO integrations
-            condition = {"ArnLike": {"aws:PrincipalArn": [sso_group_arn_1, sso_group_arn_2]}}
+            all_sso_arns = sso_group_arns_pattern_1 + sso_group_arns_pattern_2
+
+            condition = {"ArnLike": {"aws:PrincipalArn": all_sso_arns}}
             assumed_by.add_principals(iam.AccountPrincipal(self.account).with_conditions(condition))
         else:
             # Fallback: Allow any principal in the account to assume this role

{workbench-0.8.168 → workbench-0.8.170}/docs/aws_setup/core_stack.md

@@ -50,7 +50,7 @@ You'll need to set some environmental vars before deploying the stack.
 
 ```
 export WORKBENCH_BUCKET=name-of-workbench-bucket
-export WORKBENCH_SSO_GROUP=DataScientist (or whatever)
+export WORKBENCH_SSO_GROUPS=DataScientists,DataEngineers (no spaces between commas)
 ```
 
 **Optional ENV Vars**

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.168 → workbench-0.8.170}/model_docker_images/ml_pipelines/ml_pipeline_runner.py

@@ -51,14 +51,16 @@ def main():
         log.error("ML_PIPELINE_S3_PATH environment variable not set")
         sys.exit(1)
 
-    # Extract script name for local file
+    # Construct filename for local download
     script_name = os.path.basename(ml_pipeline_s3_path)
     local_script_path = f"/tmp/{script_name}"
     try:
         # Download the ML pipeline script from S3
         download_ml_pipeline_from_s3(ml_pipeline_s3_path, local_script_path)
+
         # Execute the ML pipeline
         exit_code = run_ml_pipeline(local_script_path)
+
         # Clean up
         os.remove(local_script_path)
         sys.exit(exit_code)

{workbench-0.8.168 → 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.168 → 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.168 → 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,
         )