workbench 0.8.171__tar.gz → 0.8.172__tar.gz

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

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

@@ -45,6 +45,25 @@ class AWSAccountCheck:
                 self.log.error(f"Unexpected error: {e}")
                 sys.exit(1)
 
+    def ensure_athena_results_bucket(self):
+        """Ensure that the Athena results S3 bucket exists"""
+        bucket_name = f"aws-athena-query-results-{self.aws_clamp.account_id}-{self.aws_clamp.region}"
+        self.log.important(f"Checking if Athena results bucket {bucket_name} exists...")
+
+        try:
+            s3_client = self.aws_clamp.boto3_session.client("s3")
+            s3_client.head_bucket(Bucket=bucket_name)
+            self.log.info(f"Athena results bucket {bucket_name} already exists")
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                self.log.error(f"Athena results bucket {bucket_name} does not exist.")
+                self.log.error("Create the bucket manually using this AWS CLI command:")
+                self.log.error(f"aws s3 mb s3://{bucket_name}")
+                sys.exit(1)
+            else:
+                self.log.error(f"Error checking bucket {bucket_name}: {e}")
+                sys.exit(1)
+
     def check(self):
         """Check if the AWS Account is set up Correctly"""
         self.log.info("*** Caller/Base AWS Identity Check ***")
@@ -71,6 +90,11 @@ class AWSAccountCheck:
             self.ensure_aws_catalog_db(catalog_db)
         print("\n")
 
+        # Check that the Athena Results Bucket exists
+        self.log.info("*** AWS Athena Results Bucket Check ***")
+        self.ensure_athena_results_bucket()
+        print("\n")
+
         self.log.info("AWS Account Clamp: AOK!")
 
 

{workbench-0.8.171 → workbench-0.8.172}/aws_setup/workbench_compute/workbench_compute/workbench_compute_stack.py

@@ -241,6 +241,8 @@ class WorkbenchComputeStack(Stack):
                         script_path = message['script_path']  # s3://bucket/path/to/script.py
                         size = message.get('size', 'small')
                         extra_env = message.get('environment', {})
+                        model_outputs = message.get('MODEL_OUTPUTS', [])
+                        endpoint_outputs = message.get('ENDPOINT_OUTPUTS', [])
 
                         script_name = Path(script_path).stem
                         job_name = f"workbench_{script_name}_{datetime.now():%Y%m%d_%H%M%S}"
@@ -248,16 +250,27 @@ class WorkbenchComputeStack(Stack):
                         # Get job definition name from environment variables
                         job_def_name = JOB_DEFINITIONS.get(size, JOB_DEFINITIONS['small'])
 
+                        # Build environment variables
+                        env_vars = [
+                            {'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()]
+                        ]
+
+                        # Add MODEL_OUTPUTS as comma-separated list if provided
+                        if model_outputs:
+                            env_vars.append({'name': 'MODEL_OUTPUTS', 'value': ','.join(model_outputs)})
+
+                        # Add ENDPOINT_OUTPUTS as comma-separated list if provided
+                        if endpoint_outputs:
+                            env_vars.append({'name': 'ENDPOINT_OUTPUTS', 'value': ','.join(endpoint_outputs)})
+
                         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()]
-                                ]
+                                'environment': env_vars
                             }
                         )
 

{workbench-0.8.171 → workbench-0.8.172}/aws_setup/workbench_core/workbench_core/workbench_core_stack.py

@@ -3,6 +3,7 @@ from aws_cdk import (
     Stack,
     aws_iam as iam,
     aws_logs as logs,
+    aws_athena as athena,
     RemovalPolicy,
     CfnOutput,
 )
@@ -65,6 +66,9 @@ class WorkbenchCoreStack(Stack):
         self.bucket_list = [self.workbench_bucket, athena_bucket, sagemaker_bucket] + self.additional_buckets
         self.bucket_arns = self._bucket_names_to_arns(self.bucket_list)
 
+        # Add our Athena workgroup
+        self.athena_workgroup = self.create_athena_workgroup()
+
         # Create our managed polices
         self.s3_read_policy = self.workbench_s3_read_policy()
         self.s3_full_policy = self.workbench_s3_full_policy()
@@ -84,11 +88,10 @@ class WorkbenchCoreStack(Stack):
         self.inference_store_read_policy = self.workbench_inference_store_read_policy()
         self.inference_store_full_policy = self.workbench_inference_store_full_policy()
 
-        # Create our main Workbench API Execution Role
+        # Create our main Workbench API Execution Role and Read Only Role
         self.workbench_execution_role = self.create_execution_role()
-
-        # Create a read-only role for the Workbench API
         self.workbench_readonly_role = self.create_readonly_role()
+        self._create_sso_instructions(self.workbench_execution_role, self.workbench_readonly_role)
 
         # Create additional roles for Lambda, Glue, and Batch
         self.workbench_lambda_role = self.create_lambda_role()
@@ -588,6 +591,24 @@ class WorkbenchCoreStack(Stack):
             ],
         )
 
+    # Create our Athena workgroup
+    # Note: We don't create the Athena results bucket here, just the workgroup
+    # The bucket will be created as part of th aws_account_setup.py script
+    def create_athena_workgroup(self) -> athena.CfnWorkGroup:
+        """Create workbench-specific Athena workgroup with S3 output location."""
+        athena_results_bucket = f"aws-athena-query-results-{self.account}-{self.region}"
+
+        return athena.CfnWorkGroup(
+            self,
+            "WorkbenchAthenaWorkGroup",
+            name="workbench-workgroup",
+            work_group_configuration=athena.CfnWorkGroup.WorkGroupConfigurationProperty(
+                result_configuration=athena.CfnWorkGroup.ResultConfigurationProperty(
+                    output_location=f"s3://{athena_results_bucket}/workbench/"
+                )
+            ),
+        )
+
     ######################
     #    FeatureStore    #
     ######################
@@ -1428,9 +1449,7 @@ class WorkbenchCoreStack(Stack):
         """Create the Workbench Execution Role for API-related tasks"""
         # Define the base assumed by principals with service principals
         base_principals = iam.CompositePrincipal(
-            iam.ServicePrincipal("ecs-tasks.amazonaws.com"),
-            iam.ServicePrincipal("sagemaker.amazonaws.com"),
-            iam.ServicePrincipal("glue.amazonaws.com"),
+            iam.ServicePrincipal("ecs-tasks.amazonaws.com"), iam.ServicePrincipal("sagemaker.amazonaws.com")
         )
 
         # Add SSO configuration to the principals
@@ -1467,9 +1486,7 @@ class WorkbenchCoreStack(Stack):
         """Create the Workbench Read-Only Role for viewing resources"""
 
         # ECS for Dashboard (we might switch to read-only later)
-        base_principals = iam.CompositePrincipal(
-            iam.ServicePrincipal("ecs-tasks.amazonaws.com"),
-        )
+        base_principals = iam.CompositePrincipal(iam.ServicePrincipal("ecs-tasks.amazonaws.com"))
 
         # Add SSO configuration to the principals
         assumed_by = self._create_sso_principals(base_principals)
@@ -1492,6 +1509,25 @@ class WorkbenchCoreStack(Stack):
         readonly_role.add_managed_policy(self.endpoint_read_policy)
         return readonly_role
 
+    def _create_sso_instructions(self, execution_role: iam.Role, readonly_role: iam.Role):
+        """Print SSO setup instructions to console"""
+        if self.sso_groups:
+            # Build ARNs manually since tokens aren't resolved yet
+            execution_arn = f"arn:aws:iam::{self.account}:role/{self.execution_role_name}"
+            readonly_arn = f"arn:aws:iam::{self.account}:role/{self.readonly_role_name}"
+
+            print("\n" + "=" * 60)
+            print("🔧  SSO SETUP REQUIRED (If not already done)  🔧")
+            print("=" * 60)
+            print("Have your SSO Administrator add these roles to your permission set:")
+            print(f"  • {execution_arn}")
+            print(f"  • {readonly_arn}")
+            print()
+            print("For multi-account deployments, you can add both lines for each account:")
+            print("  • arn:aws:iam::<account_id>:role/Workbench-ExecutionRole")
+            print("  • arn:aws:iam::<account_id>:role/Workbench-ReadOnlyRole")
+            print("=" * 60 + "\n")
+
     def create_lambda_role(self) -> iam.Role:
         """Create the Workbench Lambda Role."""
         lambda_assumed_by = iam.ServicePrincipal("lambda.amazonaws.com")

{workbench-0.8.171 → workbench-0.8.172}/docs/aws_setup/core_stack.md

@@ -90,13 +90,13 @@ Now that the `Workbench-ExecutionRole` has been deployed via AWS Stack. These gu
 - [Set up IAM Users](iam_assume_role.md) (not recommend, but contact us we'll help you out)
 
 
-## AWS Account Setup Check
-After deploying the Workbench Core Stack and setting up users to assume that Role, you can run this test/checking script. If the results ends with `INFO AWS Account Clamp: AOK!` you're in good shape. If not feel free to contact us on [Discord](https://discord.gg/WHAJuz8sw8) and we'll get it straightened out for you :)
+## AWS Account Setup 
+After deploying the Workbench Core Stack and setting up users to assume that Role, you can run this account setup script. If the results ends with `INFO AWS Account Clamp: AOK!` you're in good shape. If not feel free to contact us on [Discord](https://discord.gg/WHAJuz8sw8) and we'll get it straightened out for you :)
 
 ```bash
 pip install workbench (if not already installed)
 cd workbench/aws_setup
-python aws_account_check.py
+python aws_account_setup.py
 <lot of print outs for various checks>
 INFO AWS Account Clamp: AOK!
 ```

{workbench-0.8.171 → workbench-0.8.172}/docs/aws_setup/sso_assume_role.md

@@ -35,12 +35,26 @@ This guide provides step-by-step instructions to configure AWS SSO users to assu
                 "Effect": "Allow",
                 "Action": "sts:AssumeRole",
                 "Resource": "arn:aws:iam::<account-id>:role/Workbench-ExecutionRole"
+                "Resource": "arn:aws:iam::<account-id>:role/Workbench-ReadOnlyRole"
             }
         ]
     }
     ```
 
-    - Replace `<account-id>` with your AWS account ID.
+Replace `<account-id>` with your AWS account ID, OR for SSO Groups that span multiple accounts you can add both lines for each account.
+
+```
+    "Resource": "arn:aws:iam::<account-1>:role/Workbench-ExecutionRole"
+    "Resource": "arn:aws:iam::<account-1>:role/Workbench-ReadOnlyRole"
+
+    "Resource": "arn:aws:iam::<account-2>:role/Workbench-ExecutionRole"
+    "Resource": "arn:aws:iam::<account-2>:role/Workbench-ReadOnlyRole"\
+
+    "Resource": "arn:aws:iam::<account-3>:role/Workbench-ExecutionRole"
+    "Resource": "arn:aws:iam::<account-3>:role/Workbench-ReadOnlyRole"
+```
+
+**Please consult with your AWS SSO Administrator for guidance on this process.**
 
 ### 4. Save Changes
 
@@ -50,18 +64,22 @@ This guide provides step-by-step instructions to configure AWS SSO users to assu
 
 ## Verifying Access for SSO Users
 
-1. Log in to the AWS Management Console as a user in the configured group.
+1. Activate an AWS Profile for the configured SSO group.
 2. Use the following CLI command to test access:
 
     ```bash
     aws sts assume-role \
         --role-arn arn:aws:iam::<account-id>:role/Workbench-ExecutionRole \
         --role-session-name TestSession
+     
+    aws sts assume-role \
+        --role-arn arn:aws:iam::<account-id>:role/Workbench-ReadOnlyRole \
+        --role-session-name TestSession
     ```
 
-    - Replace `<account-id>` with your AWS account ID.
+    Replace `<account-id>` with your AWS account ID.
 
-3. If successful, you will receive temporary credentials for the `Workbench-ExecutionRole`.
+3. If successful, you will receive temporary credentials.
 
 
 ## Troubleshooting

workbench-0.8.171/examples/models/uq/mapie_xgb.py → workbench-0.8.172/examples/models/uq/mapie.py

@@ -10,20 +10,20 @@ my_features = FeatureSet("aqsol_features")
 model = Model("aqsol-regression")
 
 # Get the custom script path for the proximity model
-script_path = get_custom_script_path("uq_models", "mapie_xgb.template")
+script_path = get_custom_script_path("uq_models", "mapie.template")
 
 # Get Feature and Target Columns from the existing given Model
 features = model.features()
 target = model.target()
 
-# Create the Mapie XGB Model from our FeatureSet
+# Create the Mapie Model from our FeatureSet
 fs = FeatureSet(model.get_input())
 my_model = fs.to_model(
     name="aqsol-mapie-reg",
-    model_type=ModelType.REGRESSOR,
+    model_type=ModelType.UQ_REGRESSOR,
     feature_list=features,
     target_column=target,
-    description=f"Mapie XGB Model for AQSol",
+    description=f"Mapie + XGB Model for AQSol",
     tags=["aqsol", "mapie", "regression"],
     custom_script=script_path,
 )