workbench 0.8.167__tar.gz → 0.8.169__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.167
+Version: 0.8.169
 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/aws_setup/workbench_core/app.py

@@ -0,0 +1,75 @@
+import os
+import boto3
+import logging
+import aws_cdk as cdk
+from workbench_core.workbench_core_stack import WorkbenchCoreStack, WorkbenchCoreStackProps
+
+# Initialize ConfigManager with error handling
+try:
+    # Temporarily disable logging
+    # logging.disable(logging.CRITICAL)
+    from workbench.utils.config_manager import ConfigManager
+
+    cm = ConfigManager()
+except Exception as e:
+    print(f"Workbench ConfigManager initialization failed: {e}")
+    print("Falling back to environment variables only")
+    cm = None
+finally:
+    # Re-enable logging
+    logging.disable(logging.NOTSET)
+
+# Grab the account and region using boto3
+session = boto3.session.Session()
+aws_account = session.client("sts").get_caller_identity().get("Account")
+aws_region = session.region_name
+print(f"Account: {aws_account}")
+print(f"Region: {aws_region}")
+
+# Get configurations for Workbench Bucket, SSO Group, and Additional Buckets
+workbench_bucket = (cm and cm.get_config("WORKBENCH_BUCKET")) or os.getenv("WORKBENCH_BUCKET")
+if not workbench_bucket:
+    print("Error: WORKBENCH_BUCKET is required but not found in config or environment variables.")
+    exit(1)
+
+# 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", ""
+)
+additional_buckets = (
+    [bucket.strip() for bucket in additional_buckets_str.split(",") if bucket.strip()] if additional_buckets_str else []
+)
+existing_vpc_id = (cm and cm.get_config("WORKBENCH_VPC_ID")) or os.getenv("WORKBENCH_VPC_ID")
+subnet_ids = (cm and cm.get_config("WORKBENCH_SUBNET_IDS")) or os.getenv("WORKBENCH_SUBNET_IDS", "")
+
+# Log the configuration for transparency
+print("Configuration:")
+print(f"  WORKBENCH_BUCKET: {workbench_bucket}")
+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}")
+
+# Our CDK App and Environment
+app = cdk.App()
+env = cdk.Environment(account=aws_account, region=aws_region)
+
+# Create the Workbench Core Stack
+sandbox_stack = WorkbenchCoreStack(
+    app,
+    "WorkbenchCore",
+    env=env,
+    props=WorkbenchCoreStackProps(
+        workbench_bucket=workbench_bucket,
+        sso_groups=sso_groups,
+        additional_buckets=additional_buckets,
+        existing_vpc_id=existing_vpc_id,
+        subnet_ids=subnet_ids,
+    ),
+)
+
+app.synth()

workbench-0.8.169/aws_setup/workbench_core/requirements.txt

@@ -0,0 +1,3 @@
+aws-cdk-lib>=2.211.0
+constructs>=10.4.2,<11.0.0
+boto3>=1.34.0

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

@@ -4,19 +4,25 @@ 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
 
 
 class WorkbenchCoreStack(Stack):
@@ -37,8 +43,10 @@ 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
 
         # Workbench Role Names
         self.execution_role_name = "Workbench-ExecutionRole"  # Main role
@@ -93,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    #
@@ -377,12 +386,31 @@ class WorkbenchCoreStack(Stack):
     #   Batch Compute   #
     #####################
     def create_batch_compute_environment(self) -> batch.FargateComputeEnvironment:
-        """Create the Batch compute environment - super simple with Fargate."""
+        """Create the Batch compute environment with Fargate."""
+
+        # Do we have an existing VPC we want to use? Otherwise use a default
+        if self.existing_vpc_id:
+            vpc = ec2.Vpc.from_lookup(self, "ImportedVPC", vpc_id=self.existing_vpc_id)
+        else:
+            vpc = ec2.Vpc.from_lookup(self, "DefaultVpc", is_default=True)
+
+        # Use specific subnets if provided, otherwise let CDK choose
+        vpc_subnets = None
+        if self.subnet_ids:
+            vpc_subnets = ec2.SubnetSelection(
+                subnets=[
+                    ec2.Subnet.from_subnet_id(self, f"BatchSubnet{i}", subnet_id)
+                    for i, subnet_id in enumerate(self.subnet_ids)
+                ]
+            )
+
         return batch.FargateComputeEnvironment(
             self,
             "WorkbenchBatchComputeEnvironment",
             compute_environment_name="workbench-compute-env",
-            vpc=ec2.Vpc.from_lookup(self, "DefaultVpc", is_default=True),
+            vpc=vpc,
+            vpc_subnets=vpc_subnets,
+            # replace_compute_environment=True,
         )
 
     def create_batch_job_queue(self) -> batch.JobQueue:
@@ -396,6 +424,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  #
     #####################
@@ -1380,7 +1442,6 @@ class WorkbenchCoreStack(Stack):
 
         # Add policies for the Batch Role
         batch_role.add_to_policy(self.batch_job_logs())
-        # batch_role.add_to_policy(self.ecr_policy_statement())  # Not sure we need this
         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)
@@ -1402,26 +1463,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.167 → workbench-0.8.169}/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.167 → workbench-0.8.169}/model_docker_images/ml_pipelines/Dockerfile

@@ -14,12 +14,22 @@ RUN pip install --no-cache-dir -r /tmp/requirements.txt
 # Final runtime image
 FROM python:3.12-slim
 
+# Install required X11 libraries for RDKit Draw module
+RUN apt-get update && apt-get install -y \
+    libxrender1 \
+    libxext6 \
+    libsm6 \
+    libxinerama1 \
+    libfontconfig1 \
+    libcairo2 \
+    && rm -rf /var/lib/apt/lists/*
+
 # Copy installed packages from builder
 COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
 COPY --from=builder /usr/local/bin /usr/local/bin
 
 # Install workbench only
-RUN pip install --no-cache-dir workbench==0.8.166
+RUN pip install --no-cache-dir workbench==0.8.167
 
 # Add the script runner
 COPY ml_pipeline_runner.py /app/ml_pipeline_runner.py

{workbench-0.8.167 → workbench-0.8.169}/model_docker_images/ml_pipelines/ml_pipeline_runner.py

@@ -15,51 +15,45 @@ def download_ml_pipeline_from_s3(s3_path: str, local_path: str):
     parsed = urlparse(s3_path)
     bucket = parsed.netloc
     key = parsed.path.lstrip("/")
-
     log.info(f"Downloading {s3_path} to {local_path}")
-
     s3_client = boto3.client("s3")
     s3_client.download_file(bucket, key, local_path)
 
-    # Make it executable
-    os.chmod(local_path, 0o755)
-    log.info(f"Downloaded and made executable: {local_path}")
-
 
 def run_ml_pipeline(script_path: str):
     """Execute the ML pipeline script."""
     log.info(f"Executing ML pipeline: {script_path}")
-
     try:
         # Run the script with python
-        result = subprocess.run([sys.executable, script_path], capture_output=True, text=True, check=True)
-
-        log.info("ML pipeline completed successfully")
-        log.info(f"STDOUT: {result.stdout}")
-        return 0
-
-    except subprocess.CalledProcessError as e:
-        log.error(f"ML pipeline failed with exit code {e.returncode}")
-        log.error(f"STDOUT: {e.stdout}")
-        log.error(f"STDERR: {e.stderr}")
-        return e.returncode
+        result = subprocess.run(
+            [sys.executable, script_path], check=False, text=True  # Don't raise exception on non-zero exit
+        )
+        if result.returncode == 0:
+            log.info("ML pipeline completed successfully")
+        else:
+            log.error(f"ML pipeline failed with exit code {result.returncode}")
+
+        return result.returncode
+    except Exception as e:
+        log.error(f"ML pipeline execution error: {e}")
+        return 1
 
 
 def main():
+    """Main entry point for the ML pipeline runner."""
+
     # Report the version of the workbench package
     log.info(f"Workbench version: {workbench.__version__}")
 
     # Get the ML Pipeline S3 path from environment variable
     ml_pipeline_s3_path = os.environ.get("ML_PIPELINE_S3_PATH")
-
     if not ml_pipeline_s3_path:
         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)
@@ -69,9 +63,7 @@ def main():
 
         # Clean up
         os.remove(local_script_path)
-
         sys.exit(exit_code)
-
     except Exception as e:
         log.error(f"Error in run_script: {e}")
         sys.exit(1)

workbench-0.8.169/src/workbench/scripts/ml_pipeline_launcher.py

@@ -0,0 +1,137 @@
+import argparse
+import logging
+import time
+from datetime import datetime
+from pathlib import Path
+
+# Workbench Imports
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+from workbench.utils.config_manager import ConfigManager
+from workbench.utils.s3_utils import upload_content_to_s3
+from workbench.utils.cloudwatch_utils import get_cloudwatch_logs_url
+
+log = logging.getLogger("workbench")
+cm = ConfigManager()
+workbench_bucket = cm.get_config("WORKBENCH_BUCKET")
+
+
+def get_ecr_image_uri() -> str:
+    """Get the ECR image URI for the current region."""
+    region = AWSAccountClamp().region
+    return f"507740646243.dkr.ecr.{region}.amazonaws.com/aws-ml-images/py312-ml-pipelines:0.1"
+
+
+def get_batch_role_arn() -> str:
+    """Get the Batch execution role ARN."""
+    account_id = AWSAccountClamp().account_id
+    return f"arn:aws:iam::{account_id}:role/Workbench-BatchRole"
+
+
+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, size: str = "small") -> int:
+    """
+    Submit and monitor an AWS Batch job for ML pipeline execution.
+
+    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 (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
+    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 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=f"workbench-ml-pipeline-{size}",
+        containerOverrides={
+            "environment": [
+                {"name": "ML_PIPELINE_S3_PATH", "value": s3_path},
+                {"name": "WORKBENCH_BUCKET", "value": workbench_bucket},
+            ]
+        },
+    )
+    job_id = response["jobId"]
+    log.info(f"Submitted job: {job_name} ({job_id}) using {size} tier")
+
+    # Monitor job
+    last_status, running_start = None, None
+    while True:
+        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
+
+        # Handle completion
+        if status in ["SUCCEEDED", "FAILED"]:
+            exit_code = job.get("attempts", [{}])[-1].get("exitCode", 1)
+            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
+
+        time.sleep(10)
+
+
+def main():
+    """CLI entry point for running ML pipelines on AWS Batch."""
+    parser = argparse.ArgumentParser(description="Run ML pipeline script on AWS Batch")
+    parser.add_argument("script_file", help="Local path to ML pipeline script")
+    args = parser.parse_args()
+    try:
+        exit_code = run_batch_job(args.script_file)
+        exit(exit_code)
+    except Exception as e:
+        log.error(f"Error: {e}")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()