workbench 0.8.168__py3-none-any.whl → 0.8.170__py3-none-any.whl

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/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/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/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/scripts/ml_pipeline_launcher.py

@@ -27,60 +27,56 @@ def get_batch_role_arn() -> str:
     return f"arn:aws:iam::{account_id}:role/Workbench-BatchRole"
 
 
-def ensure_job_definition():
-    """Register or update the Batch job definition for ML pipeline runner."""
-    batch = AWSAccountClamp().boto3_session.client("batch")
-    name = "workbench-ml-pipeline-runner"
-    response = batch.register_job_definition(
-        jobDefinitionName=name,
-        type="container",
-        platformCapabilities=["FARGATE"],
-        containerProperties={
-            "image": get_ecr_image_uri(),
-            "resourceRequirements": [{"type": "VCPU", "value": "2"}, {"type": "MEMORY", "value": "4096"}],
-            "jobRoleArn": get_batch_role_arn(),
-            "executionRoleArn": get_batch_role_arn(),
-            "environment": [
-                {"name": "WORKBENCH_BUCKET", "value": workbench_bucket},
-                {"name": "PYTHONUNBUFFERED", "value": "1"},
-            ],
-            # "networkConfiguration": {"assignPublicIp": "ENABLED"},  # Required for ECR Image Pull (when not in VPC)
-        },
-        timeout={"attemptDurationSeconds": 10800},  # 3 hours
-    )
-    log.info(f"Job definition ready: {name} (revision {response['revision']})")
-    return name
+def _log_cloudwatch_link(job: dict, message_prefix: str = "View logs") -> None:
+    """
+    Helper method to log CloudWatch logs link with clickable URL and full URL display.
+
+    Args:
+        job: Batch job description dictionary
+        message_prefix: Prefix for the log message (default: "View logs")
+    """
+    log_stream = job.get("container", {}).get("logStreamName")
+    logs_url = get_cloudwatch_logs_url(log_group="/aws/batch/job", log_stream=log_stream)
+    if logs_url:
+        clickable_url = f"\033]8;;{logs_url}\033\\{logs_url}\033]8;;\033\\"
+        log.info(f"{message_prefix}: {clickable_url}")
+    else:
+        log.info("Check AWS Batch console for logs")
 
 
-def run_batch_job(script_path: str) -> int:
+def run_batch_job(script_path: str, size: str = "small") -> int:
     """
     Submit and monitor an AWS Batch job for ML pipeline execution.
-    This function:
-    1. Uploads the ML pipeline script to S3
-    2. Submits a Batch job to run the script in a container
-    3. Monitors job status until completion
-    4. Returns the job's exit code
+
+    Uploads script to S3, submits Batch job, monitors until completion or 2 minutes of RUNNING.
 
     Args:
         script_path: Local path to the ML pipeline script
+        size: Job size tier - "small" (default), "medium", or "large"
+          - small: 2 vCPU, 4GB RAM for lightweight processing
+          - medium: 4 vCPU, 8GB RAM for standard ML workloads
+          - large: 8 vCPU, 16GB RAM for heavy training/inference
 
     Returns:
-        Exit code from the batch job (0 for success, non-zero for failure)
+        Exit code (0 for success/disconnected, non-zero for failure)
     """
+    if size not in ["small", "medium", "large"]:
+        raise ValueError(f"Invalid size '{size}'. Must be 'small', 'medium', or 'large'")
+
     batch = AWSAccountClamp().boto3_session.client("batch")
     script_name = Path(script_path).stem
 
-    # Upload script to S3 for the container to download
+    # Upload script to S3
     s3_path = f"s3://{workbench_bucket}/batch-jobs/{Path(script_path).name}"
     log.info(f"Uploading script to {s3_path}")
     upload_content_to_s3(Path(script_path).read_text(), s3_path)
 
-    # Submit the Batch job
+    # Submit job
     job_name = f"workbench_{script_name}_{datetime.now():%Y%m%d_%H%M%S}"
     response = batch.submit_job(
         jobName=job_name,
         jobQueue="workbench-job-queue",
-        jobDefinition=ensure_job_definition(),
+        jobDefinition=f"workbench-ml-pipeline-{size}",
         containerOverrides={
             "environment": [
                 {"name": "ML_PIPELINE_S3_PATH", "value": s3_path},
@@ -89,36 +85,38 @@ def run_batch_job(script_path: str) -> int:
         },
     )
     job_id = response["jobId"]
-    log.info(f"Submitted job: {job_name} ({job_id})")
+    log.info(f"Submitted job: {job_name} ({job_id}) using {size} tier")
 
-    # Monitor job execution
-    last_status = None
+    # Monitor job
+    last_status, running_start = None, None
     while True:
-        # Check job status
         job = batch.describe_jobs(jobs=[job_id])["jobs"][0]
         status = job["status"]
+
         if status != last_status:
             log.info(f"Job status: {status}")
             last_status = status
+            if status == "RUNNING":
+                running_start = time.time()
+
+        # Disconnect after 2 minutes of running
+        if status == "RUNNING" and running_start and (time.time() - running_start >= 120):
+            log.info("✅  ML Pipeline is running successfully!")
+            _log_cloudwatch_link(job, "📊  Monitor logs")
+            return 0
 
-        # Check if job completed
+        # Handle completion
         if status in ["SUCCEEDED", "FAILED"]:
             exit_code = job.get("attempts", [{}])[-1].get("exitCode", 1)
-            if status == "FAILED":
-                log.error(f"Job failed: {job.get('statusReason', 'Unknown reason')}")
-            else:
-                log.info("Job completed successfully")
-
-            # Get CloudWatch logs URL
-            log_stream_name = job.get("container", {}).get("logStreamName")
-            logs_url = get_cloudwatch_logs_url(log_group="/aws/batch/job", log_stream=log_stream_name)
-            if logs_url:
-                # OSC 8 hyperlink format for modern terminals
-                clickable_url = f"\033]8;;{logs_url}\033\\{logs_url}\033]8;;\033\\"
-                log.info(f"View logs: {clickable_url}")
+            msg = (
+                "Job completed successfully"
+                if status == "SUCCEEDED"
+                else f"Job failed: {job.get('statusReason', 'Unknown')}"
+            )
+            log.info(msg) if status == "SUCCEEDED" else log.error(msg)
+            _log_cloudwatch_link(job)
             return exit_code
 
-        # Sleep a bit before next status check
         time.sleep(10)
 
 

{workbench-0.8.168.dist-info → workbench-0.8.170.dist-info}/METADATA

@@ -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.dist-info → workbench-0.8.170.dist-info}/RECORD

@@ -35,7 +35,7 @@ workbench/api/endpoint.py,sha256=RWGqxsCW_pMiENMb_XZlm2ZCldMS4suEBM3F5gT3hYI,381
 workbench/api/feature_set.py,sha256=wzNxNjN0K2FaIC7QUIogMnoHqw2vo0iAHYlGk6fWLCw,6649
 workbench/api/graph_store.py,sha256=LremJyPrQFgsHb7hxsctuCsoxx3p7TKtaY5qALHe6pc,4372
 workbench/api/meta.py,sha256=1_9989cPvf3hd3tA-83hLijOGNnhwXAF8aZF45adeDQ,8596
-workbench/api/model.py,sha256=2hPN8UK4whZ0kDgPtbR7lEknw7XhH5hGYaHA55jmZWQ,4529
+workbench/api/model.py,sha256=RkFVXnlLcMlzNKRUFr_GCmZ7IQJMyhB2lwMwd22HBBo,4691
 workbench/api/monitor.py,sha256=kQHSFiVLRWnHekSdatMKR3QbRj1BBNrVXpZgvV83LPM,5027
 workbench/api/parameter_store.py,sha256=7BObkuATuP6C5AG_46kCWsmuCwuh1vgMJDBSN0gTkwM,4294
 workbench/api/pipeline.py,sha256=MSYGrDSXrRB_oQELtAlOwBfxSBTw3REAkHy5XBHau0Y,6261
@@ -53,10 +53,10 @@ workbench/core/artifacts/athena_source.py,sha256=RNmCe7s6uH4gVHpcdJcL84aSbF5Q1ah
 workbench/core/artifacts/cached_artifact_mixin.py,sha256=ngqFLZ4cQx_TFouXZgXZQsv_7W6XCvxVGXXSfzzaft8,3775
 workbench/core/artifacts/data_source_abstract.py,sha256=5IRCzFVK-17cd4NXPMRfx99vQAmQ0WHE5jcm5RfsVTg,10619
 workbench/core/artifacts/data_source_factory.py,sha256=YL_tA5fsgubbB3dPF6T4tO0rGgz-6oo3ge4i_YXVC-M,2380
-workbench/core/artifacts/endpoint_core.py,sha256=6uDOl-VKrTbLMlHZEYFY80XwrCP5H0W36JoHySjhl7M,48163
+workbench/core/artifacts/endpoint_core.py,sha256=CtLo_eqONpDvyyYtjRkrrujVVv6cJV-vRQxHef6MYdk,48841
 workbench/core/artifacts/feature_set_core.py,sha256=055VdSYR09HP4ygAuYvIYtHQ7Ec4XxsZygpgEl5H5jQ,29136
 workbench/core/artifacts/model_core.py,sha256=U0dSkpZMrsIgbUglVkPwAgN0gji7Oa7glOjqMQJDAzE,50927
-workbench/core/artifacts/monitor_core.py,sha256=AIgceuO_YqAmttDdJWcdFcrvE93icfSsTr6WgXWGjYo,37684
+workbench/core/artifacts/monitor_core.py,sha256=BvJ8gMxZXYZeMzAC25PVTXWiyXKtxD1qK6LqDcIByzs,37657
 workbench/core/cloud_platform/cloud_meta.py,sha256=-g4-LTC3D0PXb3VfaXdLR1ERijKuHdffeMK_zhD-koQ,8809
 workbench/core/cloud_platform/aws/README.md,sha256=QT5IQXoUHbIA0qQ2wO6_2P2lYjYQFVYuezc22mWY4i8,97
 workbench/core/cloud_platform/aws/aws_account_clamp.py,sha256=OzFknZXKW7VTvnDGGX4BXKoh0i1gQ7yaEBhkLCyHFSs,6310
@@ -103,7 +103,7 @@ workbench/core/transforms/features_to_features/heavy/glue/Readme.md,sha256=TuyCa
 workbench/core/transforms/features_to_model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workbench/core/transforms/features_to_model/features_to_model.py,sha256=gwqdQZJUIfZv1M7uGhzzBxUwRS0thJE_o_H2IUsBT40,19789
 workbench/core/transforms/model_to_endpoint/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-workbench/core/transforms/model_to_endpoint/model_to_endpoint.py,sha256=hbLsdOVlfAH4XCVNUfr3SFH8rKjxIs0QyYrNwjCh7SM,4970
+workbench/core/transforms/model_to_endpoint/model_to_endpoint.py,sha256=TIYXvuK0s383PwJ4iS6fCRhuif6oIxsoWb4CpMGJjY4,6358
 workbench/core/transforms/pandas_transforms/__init__.py,sha256=xL4MT8-fZ1SFqDbTLc8XyxjupHtB1YR6Ej0AC2nwd7I,894
 workbench/core/transforms/pandas_transforms/data_to_pandas.py,sha256=sJHPeuNF8Q8aQqgRnkdWkyvur5cbggdUVIwR-xF3Dlo,3621
 workbench/core/transforms/pandas_transforms/features_to_pandas.py,sha256=af6xdPt2V4zhh-SzQa_UYxdmNMzMLXbrbsznV5QoIJg,3441
@@ -167,7 +167,7 @@ workbench/resources/open_source_api.key,sha256=3S0OTblsmC0msUPdE_dbBmI83xJNmYscu
 workbench/resources/signature_verify_pub.pem,sha256=V3-u-3_z2PH-805ybkKvzDOBwAbvHxcKn0jLBImEtzM,272
 workbench/scripts/check_double_bond_stereo.py,sha256=p5hnL54Weq77ES0HCELq9JeoM-PyUGkvVSeWYF2dKyo,7776
 workbench/scripts/glue_launcher.py,sha256=bIKQvfGxpAhzbeNvTnHfRW_5kQhY-169_868ZnCejJk,10692
-workbench/scripts/ml_pipeline_launcher.py,sha256=RUKUBERL7RE-uNs_ttkPUa6Rf-QJERYWIhp_XLOtF78,5083
+workbench/scripts/ml_pipeline_launcher.py,sha256=fjI35SXi9CDSQ6Lan7qGcLAHkVCDioyhbPlo0eDHDxQ,4913
 workbench/scripts/monitor_cloud_watch.py,sha256=s7MY4bsHts0nup9G0lWESCvgJZ9Mw1Eo-c8aKRgLjMw,9235
 workbench/scripts/redis_expire.py,sha256=DxI_RKSNlrW2BsJZXcsSbaWGBgPZdPhtzHjV9SUtElE,1120
 workbench/scripts/redis_report.py,sha256=iaJSuGPyLCs6e0TMcZDoT0YyJ43xJ1u74YD8FLnnUg4,990
@@ -276,9 +276,9 @@ workbench/web_interface/page_views/main_page.py,sha256=X4-KyGTKLAdxR-Zk2niuLJB2Y
 workbench/web_interface/page_views/models_page_view.py,sha256=M0bdC7bAzLyIaE2jviY12FF4abdMFZmg6sFuOY_LaGI,2650
 workbench/web_interface/page_views/page_view.py,sha256=Gh6YnpOGlUejx-bHZAf5pzqoQ1H1R0OSwOpGhOBO06w,455
 workbench/web_interface/page_views/pipelines_page_view.py,sha256=v2pxrIbsHBcYiblfius3JK766NZ7ciD2yPx0t3E5IJo,2656
-workbench-0.8.168.dist-info/licenses/LICENSE,sha256=z4QMMPlLJkZjU8VOKqJkZiQZCEZ--saIU2Z8-p3aVc0,1080
-workbench-0.8.168.dist-info/METADATA,sha256=PTFR16ft5NCrG-_umsJKrIJSa3eLnpju1EkXZmafxxM,9210
-workbench-0.8.168.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-workbench-0.8.168.dist-info/entry_points.txt,sha256=V_v6hQ4DYoCJnTnqbm036reCri_CXkA_ONcRSuF5OKg,305
-workbench-0.8.168.dist-info/top_level.txt,sha256=Dhy72zTxaA_o_yRkPZx5zw-fwumnjGaeGf0hBN3jc_w,10
-workbench-0.8.168.dist-info/RECORD,,
+workbench-0.8.170.dist-info/licenses/LICENSE,sha256=z4QMMPlLJkZjU8VOKqJkZiQZCEZ--saIU2Z8-p3aVc0,1080
+workbench-0.8.170.dist-info/METADATA,sha256=GbS745jAMPDykgLqfLcwjb9nRnczT-uV9Q11GbPBAX8,9210
+workbench-0.8.170.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+workbench-0.8.170.dist-info/entry_points.txt,sha256=V_v6hQ4DYoCJnTnqbm036reCri_CXkA_ONcRSuF5OKg,305
+workbench-0.8.170.dist-info/top_level.txt,sha256=Dhy72zTxaA_o_yRkPZx5zw-fwumnjGaeGf0hBN3jc_w,10
+workbench-0.8.170.dist-info/RECORD,,