workbench 0.8.169__tar.gz → 0.8.171__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.169
+Version: 0.8.171
 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.171}/aws_setup/aws_account_check.py

@@ -39,7 +39,7 @@ class AWSAccountCheck:
                 self.log.error(f"AccessDeniedException {e}")
                 self.log.error(f"Access denied while trying to create/access the catalog database '{catalog_db}'.")
                 self.log.error("Create the database manually in the AWS Glue Console, or run this command:")
-                self.log.error('aws glue create-database --database-input \'{"Name": "workbench"}\'')
+                self.log.error(f'aws glue create-database --database-input \'{{"Name": "{catalog_db}"}}\'')
                 sys.exit(1)
             else:
                 self.log.error(f"Unexpected error: {e}")

workbench-0.8.171/aws_setup/workbench_compute/README.md

@@ -0,0 +1,58 @@
+
+# Welcome to your CDK Python project!
+
+This is a blank project for CDK development with Python.
+
+The `cdk.json` file tells the CDK Toolkit how to execute your app.
+
+This project is set up like a standard Python project.  The initialization
+process also creates a virtualenv within this project, stored under the `.venv`
+directory.  To create the virtualenv it assumes that there is a `python3`
+(or `python` for Windows) executable in your path with access to the `venv`
+package. If for any reason the automatic creation of the virtualenv fails,
+you can create the virtualenv manually.
+
+To manually create a virtualenv on MacOS and Linux:
+
+```
+$ python3 -m venv .venv
+```
+
+After the init process completes and the virtualenv is created, you can use the following
+step to activate your virtualenv.
+
+```
+$ source .venv/bin/activate
+```
+
+If you are a Windows platform, you would activate the virtualenv like this:
+
+```
+% .venv\Scripts\activate.bat
+```
+
+Once the virtualenv is activated, you can install the required dependencies.
+
+```
+$ pip install -r requirements.txt
+```
+
+At this point you can now synthesize the CloudFormation template for this code.
+
+```
+$ cdk synth
+```
+
+To add additional dependencies, for example other CDK libraries, just add
+them to your `setup.py` file and rerun the `pip install -r requirements.txt`
+command.
+
+## Useful commands
+
+ * `cdk ls`          list all stacks in the app
+ * `cdk synth`       emits the synthesized CloudFormation template
+ * `cdk deploy`      deploy this stack to your default AWS account/region
+ * `cdk diff`        compare deployed stack with current state
+ * `cdk docs`        open CDK documentation
+
+Enjoy!

workbench-0.8.171/aws_setup/workbench_compute/app.py

@@ -0,0 +1,64 @@
+import boto3
+import aws_cdk as cdk
+from aws_cdk import Fn
+from workbench_compute.workbench_compute_stack import WorkbenchComputeStack, WorkbenchComputeStackProps
+
+# Initialize Workbench ConfigManager
+try:
+    from workbench.utils.config_manager import ConfigManager
+
+    cm = ConfigManager()
+except ImportError:
+    print("Workbench Configuration Manager Not Found...")
+    print("Set the WORKBENCH_CONFIG Env var and run again...")
+    raise SystemExit(1)
+
+# 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 needed for compute stack
+workbench_bucket = cm.get_config("WORKBENCH_BUCKET")
+if not workbench_bucket:
+    print("Error: WORKBENCH_BUCKET is required but not found in config.")
+    exit(1)
+
+
+# Get the Batch and Lambda Roles from WorkbenchCore stack outputs
+batch_role_arn = Fn.import_value("WorkbenchCore-BatchRoleArn")
+lambda_role_arn = Fn.import_value("WorkbenchCore-LambdaRoleArn")
+
+# VPC and subnet configuration (optional)
+existing_vpc_id = cm.get_config("WORKBENCH_VPC_ID")
+subnet_ids = cm.get_config("WORKBENCH_SUBNET_IDS") or []
+
+# Log the configuration for transparency
+print("Configuration:")
+print(f"  WORKBENCH_BUCKET: {workbench_bucket}")
+print(f"  WORKBENCH_BATCH_ROLE_ARN: {batch_role_arn}")
+print(f"  WORKBENCH_LAMBDA_ROLE_ARN: {lambda_role_arn}")
+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 Compute Stack
+compute_stack = WorkbenchComputeStack(
+    app,
+    "WorkbenchCompute",
+    env=env,
+    props=WorkbenchComputeStackProps(
+        workbench_bucket=workbench_bucket,
+        batch_role_arn=batch_role_arn,
+        lambda_role_arn=lambda_role_arn,
+        existing_vpc_id=existing_vpc_id,
+        subnet_ids=subnet_ids,
+    ),
+)
+
+app.synth()

workbench-0.8.171/aws_setup/workbench_compute/cdk.json

@@ -0,0 +1,97 @@
+{
+  "app": "python3 app.py",
+  "watch": {
+    "include": [
+      "**"
+    ],
+    "exclude": [
+      "README.md",
+      "cdk*.json",
+      "requirements*.txt",
+      "source.bat",
+      "**/__init__.py",
+      "**/__pycache__",
+      "tests"
+    ]
+  },
+  "context": {
+    "@aws-cdk/aws-lambda:recognizeLayerVersion": true,
+    "@aws-cdk/core:checkSecretUsage": true,
+    "@aws-cdk/core:target-partitions": [
+      "aws",
+      "aws-cn"
+    ],
+    "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true,
+    "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true,
+    "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true,
+    "@aws-cdk/aws-iam:minimizePolicies": true,
+    "@aws-cdk/core:validateSnapshotRemovalPolicy": true,
+    "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true,
+    "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true,
+    "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true,
+    "@aws-cdk/aws-apigateway:disableCloudWatchRole": true,
+    "@aws-cdk/core:enablePartitionLiterals": true,
+    "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true,
+    "@aws-cdk/aws-ecs:disableExplicitDeploymentControllerForCircuitBreaker": true,
+    "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true,
+    "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true,
+    "@aws-cdk/aws-route53-patters:useCertificate": true,
+    "@aws-cdk/customresources:installLatestAwsSdkDefault": false,
+    "@aws-cdk/aws-rds:databaseProxyUniqueResourceName": true,
+    "@aws-cdk/aws-codedeploy:removeAlarmsFromDeploymentGroup": true,
+    "@aws-cdk/aws-apigateway:authorizerChangeDeploymentLogicalId": true,
+    "@aws-cdk/aws-ec2:launchTemplateDefaultUserData": true,
+    "@aws-cdk/aws-secretsmanager:useAttachedSecretResourcePolicyForSecretTargetAttachments": true,
+    "@aws-cdk/aws-redshift:columnId": true,
+    "@aws-cdk/aws-stepfunctions-tasks:enableEmrServicePolicyV2": true,
+    "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": true,
+    "@aws-cdk/aws-apigateway:requestValidatorUniqueId": true,
+    "@aws-cdk/aws-kms:aliasNameRef": true,
+    "@aws-cdk/aws-kms:applyImportedAliasPermissionsToPrincipal": true,
+    "@aws-cdk/aws-autoscaling:generateLaunchTemplateInsteadOfLaunchConfig": true,
+    "@aws-cdk/core:includePrefixInUniqueNameGeneration": true,
+    "@aws-cdk/aws-efs:denyAnonymousAccess": true,
+    "@aws-cdk/aws-opensearchservice:enableOpensearchMultiAzWithStandby": true,
+    "@aws-cdk/aws-lambda-nodejs:useLatestRuntimeVersion": true,
+    "@aws-cdk/aws-efs:mountTargetOrderInsensitiveLogicalId": true,
+    "@aws-cdk/aws-rds:auroraClusterChangeScopeOfInstanceParameterGroupWithEachParameters": true,
+    "@aws-cdk/aws-appsync:useArnForSourceApiAssociationIdentifier": true,
+    "@aws-cdk/aws-rds:preventRenderingDeprecatedCredentials": true,
+    "@aws-cdk/aws-codepipeline-actions:useNewDefaultBranchForCodeCommitSource": true,
+    "@aws-cdk/aws-cloudwatch-actions:changeLambdaPermissionLogicalIdForLambdaAction": true,
+    "@aws-cdk/aws-codepipeline:crossAccountKeysDefaultValueToFalse": true,
+    "@aws-cdk/aws-codepipeline:defaultPipelineTypeToV2": true,
+    "@aws-cdk/aws-kms:reduceCrossAccountRegionPolicyScope": true,
+    "@aws-cdk/aws-eks:nodegroupNameAttribute": true,
+    "@aws-cdk/aws-ec2:ebsDefaultGp3Volume": true,
+    "@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm": true,
+    "@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault": false,
+    "@aws-cdk/aws-s3:keepNotificationInImportedBucket": false,
+    "@aws-cdk/core:explicitStackTags": true,
+    "@aws-cdk/aws-ecs:enableImdsBlockingDeprecatedFeature": false,
+    "@aws-cdk/aws-ecs:disableEcsImdsBlocking": true,
+    "@aws-cdk/aws-ecs:reduceEc2FargateCloudWatchPermissions": true,
+    "@aws-cdk/aws-dynamodb:resourcePolicyPerReplica": true,
+    "@aws-cdk/aws-ec2:ec2SumTImeoutEnabled": true,
+    "@aws-cdk/aws-appsync:appSyncGraphQLAPIScopeLambdaPermission": true,
+    "@aws-cdk/aws-rds:setCorrectValueForDatabaseInstanceReadReplicaInstanceResourceId": true,
+    "@aws-cdk/core:cfnIncludeRejectComplexResourceUpdateCreatePolicyIntrinsics": true,
+    "@aws-cdk/aws-lambda-nodejs:sdkV3ExcludeSmithyPackages": true,
+    "@aws-cdk/aws-stepfunctions-tasks:fixRunEcsTaskPolicy": true,
+    "@aws-cdk/aws-ec2:bastionHostUseAmazonLinux2023ByDefault": true,
+    "@aws-cdk/aws-route53-targets:userPoolDomainNameMethodWithoutCustomResource": true,
+    "@aws-cdk/aws-elasticloadbalancingV2:albDualstackWithoutPublicIpv4SecurityGroupRulesDefault": true,
+    "@aws-cdk/aws-iam:oidcRejectUnauthorizedConnections": true,
+    "@aws-cdk/core:enableAdditionalMetadataCollection": true,
+    "@aws-cdk/aws-lambda:createNewPoliciesWithAddToRolePolicy": false,
+    "@aws-cdk/aws-s3:setUniqueReplicationRoleName": true,
+    "@aws-cdk/aws-events:requireEventBusPolicySid": true,
+    "@aws-cdk/core:aspectPrioritiesMutating": true,
+    "@aws-cdk/aws-dynamodb:retainTableReplica": true,
+    "@aws-cdk/aws-stepfunctions:useDistributedMapResultWriterV2": true,
+    "@aws-cdk/s3-notifications:addS3TrustKeyPolicyForSnsSubscriptions": true,
+    "@aws-cdk/aws-ec2:requirePrivateSubnetsForEgressOnlyInternetGateway": true,
+    "@aws-cdk/aws-s3:publicAccessBlockedByDefault": true,
+    "@aws-cdk/aws-lambda:useCdkManagedLogGroup": true
+  }
+}

workbench-0.8.171/aws_setup/workbench_compute/requirements.txt

@@ -0,0 +1,2 @@
+aws-cdk-lib==2.211.0
+constructs>=10.0.0,<11.0.0

workbench-0.8.171/aws_setup/workbench_compute/tests/unit/test_workbench_compute_stack.py

@@ -0,0 +1,17 @@
+import aws_cdk as core
+import aws_cdk.assertions as assertions
+
+from workbench_compute.workbench_compute_stack import WorkbenchComputeStack
+
+
+# example tests. To run these tests, uncomment this file along with the example
+# resource in workbench_compute/workbench_compute_stack.py
+def test_sqs_queue_created():
+    app = core.App()
+    stack = WorkbenchComputeStack(app, "workbench-compute")
+    template = assertions.Template.from_stack(stack)
+
+
+#     template.has_resource_properties("AWS::SQS::Queue", {
+#         "VisibilityTimeout": 300
+#     })

workbench-0.8.171/aws_setup/workbench_compute/workbench_compute/workbench_compute_stack.py

@@ -0,0 +1,272 @@
+from aws_cdk import (
+    Environment,
+    Stack,
+    aws_iam as iam,
+    aws_ec2 as ec2,
+    aws_ecs as ecs,
+    aws_ecr as ecr,
+    aws_batch as batch,
+    aws_sqs as sqs,
+    aws_lambda as lambda_,
+    aws_lambda_event_sources as lambda_events,
+    Duration,
+    Size,
+)
+from constructs import Construct
+from typing import Any, List, Dict
+from dataclasses import dataclass, field
+from textwrap import dedent
+
+
+@dataclass
+class WorkbenchComputeStackProps:
+    workbench_bucket: str
+    batch_role_arn: str
+    lambda_role_arn: str
+    existing_vpc_id: str = None
+    subnet_ids: List[str] = field(default_factory=list)  # Optional subnets ids for the Batch compute environment
+
+
+class WorkbenchComputeStack(Stack):
+    def __init__(
+        self,
+        scope: Construct,
+        construct_id: str,
+        env: Environment,
+        props: WorkbenchComputeStackProps,
+        **kwargs: Any,
+    ) -> None:
+        desc = "Workbench Compute: Defines the compute infrastructure for Workbench including Batch resources."
+        super().__init__(scope, construct_id, env=env, description=desc, **kwargs)
+
+        # Print the environment details
+        print("Environment")
+        print(env)
+
+        # Grab our properties
+        self.workbench_bucket = props.workbench_bucket
+        self.batch_role_arn = props.batch_role_arn
+        self.existing_vpc_id = props.existing_vpc_id
+        self.subnet_ids = props.subnet_ids
+
+        # Import the Batch role
+        self.workbench_batch_role = iam.Role.from_role_arn(self, "ImportedBatchRole", self.batch_role_arn)
+
+        # Import the Lambda role
+        self.workbench_lambda_role = iam.Role.from_role_arn(self, "ImportedLambdaRole", props.lambda_role_arn)
+
+        # Create Batch resources first (Lambda will reference them)
+        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()
+
+        # Create ML Pipeline SQS Queue and Lambda (after Batch resources)
+        self.ml_pipeline_queue = self.create_ml_pipeline_queue()
+        self.batch_trigger_lambda = self.create_batch_trigger_lambda()
+
+    #####################
+    #   Batch Compute   #
+    #####################
+    def create_batch_compute_environment(self) -> batch.FargateComputeEnvironment:
+        """Create the Batch compute environment with Fargate."""
+
+        # Check if we have an existing VPC to use
+        if self.existing_vpc_id:
+            vpc = ec2.Vpc.from_lookup(self, "ImportedVPC", vpc_id=self.existing_vpc_id)
+            vpc_subnets = None
+
+            # Use specific subnets if provided
+            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)
+                    ]
+                )
+        else:
+            raise ValueError(
+                'Please provide the Workbench Config entry: "WORKBENCH_VPC_ID":"vpc-123456789abcdef0". '
+                "Default VPC networks cannot pull ECR images for Batch jobs."
+            )
+
+        return batch.FargateComputeEnvironment(
+            self,
+            "WorkbenchBatchComputeEnvironment",
+            compute_environment_name="workbench-compute-env",
+            vpc=vpc,
+            vpc_subnets=vpc_subnets,
+        )
+
+    def create_batch_job_queue(self) -> batch.JobQueue:
+        """Create the Batch job queue."""
+        return batch.JobQueue(
+            self,
+            "WorkbenchBatchJobQueue",
+            job_queue_name="workbench-job-queue",
+            compute_environments=[
+                batch.OrderedComputeEnvironment(compute_environment=self.batch_compute_environment, order=1)
+            ],
+        )
+
+    def create_batch_job_definitions(self) -> Dict[str, batch.EcsJobDefinition]:
+        """Create ML pipeline job definitions in small/medium/large tiers."""
+
+        # Using ECR image for ML pipelines
+        ecr_image = ecs.ContainerImage.from_ecr_repository(
+            repository=ecr.Repository.from_repository_arn(
+                self,
+                "MLPipelineRepo",
+                repository_arn=f"arn:aws:ecr:{self.region}:507740646243:repository/aws-ml-images/py312-ml-pipelines",
+            ),
+            tag="0.1",
+        )
+
+        # Job Definition Tiers
+        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-batch-{size}",
+                container=batch.EcsFargateContainerDefinition(
+                    self,
+                    f"Container{size.capitalize()}",
+                    cpu=cpu,
+                    memory=Size.mebibytes(memory_mib),
+                    image=ecr_image,
+                    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
+
+    ###########################
+    #  ML Pipeline SQS Queue  #
+    ###########################
+    def create_ml_pipeline_queue(self) -> sqs.Queue:
+        """Create SQS FIFO queue for ML pipeline orchestration with deduplication."""
+
+        # Dead letter queue for failed messages (also FIFO)
+        dlq = sqs.Queue(
+            self,
+            "MLPipelineDLQ",
+            queue_name="workbench-ml-pipeline-dlq.fifo",  # Must end with .fifo
+            fifo=True,
+            retention_period=Duration.days(14),
+        )
+
+        # Main FIFO queue with deduplication
+        return sqs.Queue(
+            self,
+            "MLPipelineQueue",
+            queue_name="workbench-ml-pipeline-queue.fifo",  # Must end with .fifo
+            fifo=True,
+            content_based_deduplication=True,  # Auto-dedupe based on message content
+            visibility_timeout=Duration.minutes(15),
+            retention_period=Duration.days(1),
+            dead_letter_queue=sqs.DeadLetterQueue(max_receive_count=1, queue=dlq),
+        )
+
+    def create_batch_trigger_lambda(self) -> lambda_.Function:
+        """Create Lambda function to process SQS messages and trigger Batch jobs."""
+
+        # Create a mapping of job definition names to pass to Lambda
+        job_def_names = {size: f"workbench-batch-{size}" for size in ["small", "medium", "large"]}
+
+        batch_trigger_lambda = lambda_.Function(
+            self,
+            "BatchTriggerLambda",
+            function_name="workbench-batch-trigger",
+            runtime=lambda_.Runtime.PYTHON_3_12,
+            handler="index.lambda_handler",
+            code=lambda_.Code.from_inline(self._get_lambda_code()),
+            timeout=Duration.minutes(5),
+            environment={
+                "WORKBENCH_BUCKET": self.workbench_bucket,
+                "JOB_QUEUE": self.batch_job_queue.job_queue_name,
+                "JOB_DEF_SMALL": job_def_names["small"],
+                "JOB_DEF_MEDIUM": job_def_names["medium"],
+                "JOB_DEF_LARGE": job_def_names["large"],
+            },
+            role=self.workbench_lambda_role,
+        )
+
+        # Connect SQS to Lambda
+        batch_trigger_lambda.add_event_source(
+            lambda_events.SqsEventSource(
+                self.ml_pipeline_queue, batch_size=1, max_concurrency=5  # Limit parallel Batch job submissions
+            )
+        )
+        return batch_trigger_lambda
+
+    @staticmethod
+    def _get_lambda_code() -> str:
+        """Return the Lambda function code as a string."""
+        return dedent(
+            '''
+            import json
+            import boto3
+            import os
+            from datetime import datetime
+            from pathlib import Path
+
+            batch = boto3.client('batch')
+            s3 = boto3.client('s3')
+
+            WORKBENCH_BUCKET = os.environ['WORKBENCH_BUCKET']
+            JOB_QUEUE = os.environ['JOB_QUEUE']
+            JOB_DEFINITIONS = {
+                'small': os.environ['JOB_DEF_SMALL'],
+                'medium': os.environ['JOB_DEF_MEDIUM'],
+                'large': os.environ['JOB_DEF_LARGE'],
+            }
+
+            def lambda_handler(event, context):
+                """Process SQS messages and submit Batch jobs."""
+
+                for record in event['Records']:
+                    try:
+                        message = json.loads(record['body'])
+                        script_path = message['script_path']  # s3://bucket/path/to/script.py
+                        size = message.get('size', 'small')
+                        extra_env = message.get('environment', {})
+
+                        script_name = Path(script_path).stem
+                        job_name = f"workbench_{script_name}_{datetime.now():%Y%m%d_%H%M%S}"
+
+                        # Get job definition name from environment variables
+                        job_def_name = JOB_DEFINITIONS.get(size, JOB_DEFINITIONS['small'])
+
+                        response = batch.submit_job(
+                            jobName=job_name,
+                            jobQueue=JOB_QUEUE,
+                            jobDefinition=job_def_name,
+                            containerOverrides={
+                                'environment': [
+                                    {'name': 'ML_PIPELINE_S3_PATH', 'value': script_path},
+                                    {'name': 'WORKBENCH_BUCKET', 'value': WORKBENCH_BUCKET},
+                                    *[{'name': k, 'value': v} for k, v in extra_env.items()]
+                                ]
+                            }
+                        )
+
+                        print(f"Submitted job: {job_name} ({response['jobId']})")
+
+                    except Exception as e:
+                        print(f"Error processing message: {e}")
+                        raise  # Let SQS retry via DLQ
+
+                return {'statusCode': 200}
+        '''
+        ).strip()