workbench 0.8.161__py3-none-any.whl → 0.8.192__py3-none-any.whl

workbench/repl/workbench_shell.py

@@ -1,4 +1,6 @@
+import IPython
 from IPython import start_ipython
+from distutils.version import LooseVersion
 from IPython.terminal.prompts import Prompts
 from IPython.terminal.ipapp import load_default_config
 from pygments.token import Token
@@ -39,7 +41,7 @@ from workbench.cached.cached_meta import CachedMeta
 try:
     import rdkit  # noqa
     import mordred  # noqa
-    from workbench.utils import chem_utils
+    from workbench.utils.chem_utils import vis
 
     HAVE_CHEM_UTILS = True
 except ImportError:
@@ -70,7 +72,7 @@ if not ConfigManager().config_okay():
 
 # Set the log level to important
 log = logging.getLogger("workbench")
-log.setLevel(IMPORTANT_LEVEL_NUM)
+log.setLevel(logging.INFO)
 log.addFilter(
     lambda record: not (
         record.getMessage().startswith("Async: Metadata") or record.getMessage().startswith("Updated Metadata")
@@ -176,12 +178,12 @@ class WorkbenchShell:
 
         # Add cheminformatics utils if available
         if HAVE_CHEM_UTILS:
-            self.commands["show"] = chem_utils.show
+            self.commands["show"] = vis.show
 
     def start(self):
         """Start the Workbench IPython shell"""
         cprint("magenta", "\nWelcome to Workbench!")
-        if self.aws_status is False:
+        if not self.aws_status:
             cprint("red", "AWS Account Connection Failed...Review/Fix the Workbench Config:")
             cprint("red", f"Path: {self.cm.site_config_path}")
             self.show_config()
@@ -202,7 +204,10 @@ class WorkbenchShell:
 
         # Start IPython with the config and commands in the namespace
         try:
-            ipython_argv = ["--no-tip", "--theme", "linux"]
+            if LooseVersion(IPython.__version__) >= LooseVersion("9.0.0"):
+                ipython_argv = ["--no-tip", "--theme", "linux"]
+            else:
+                ipython_argv = []
             start_ipython(ipython_argv, user_ns=locs, config=config)
         finally:
             spinner = self.spinner_start("Goodbye to AWS:")
@@ -555,7 +560,7 @@ class WorkbenchShell:
         from workbench.web_interface.components.plugin_unit_test import PluginUnitTest
 
         # Get kwargs
-        theme = kwargs.get("theme", "dark")
+        theme = kwargs.get("theme", "midnight_blue")
 
         plugin_test = PluginUnitTest(plugin_class, theme=theme, input_data=data, **kwargs)
 

workbench/scripts/lambda_launcher.py

@@ -0,0 +1,63 @@
+import sys
+import os
+import json
+import importlib.util
+
+
+def main():
+    if len(sys.argv) != 2:
+        print("Usage: lambda_launcher <handler_module_name>")
+        print("\nOptional: testing/event.json with test event")
+        print("Optional: testing/env.json with environment variables")
+        sys.exit(1)
+
+    handler_file = sys.argv[1]
+
+    # Add .py if not present
+    if not handler_file.endswith(".py"):
+        handler_file += ".py"
+
+    # Check if file exists
+    if not os.path.exists(handler_file):
+        print(f"Error: File '{handler_file}' not found")
+        sys.exit(1)
+
+    # Load environment variables from env.json if it exists
+    if os.path.exists("testing/env.json"):
+        print("Loading environment variables from testing/env.json")
+        with open("testing/env.json") as f:
+            env_vars = json.load(f)
+            for key, value in env_vars.items():
+                os.environ[key] = value
+                print(f"  Set {key} = {value}")
+        print()
+
+    # Load event configuration
+    if os.path.exists("testing/event.json"):
+        print("Loading event from testing/event.json")
+        with open("testing/event.json") as f:
+            event = json.load(f)
+    else:
+        print("No testing/event.json found, using empty event")
+        event = {}
+
+    # Load the module dynamically
+    spec = importlib.util.spec_from_file_location("lambda_module", handler_file)
+    lambda_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(lambda_module)
+
+    # Call the lambda_handler
+    print(f"Invoking lambda_handler from {handler_file}...")
+    print("-" * 50)
+    print(f"Event: {json.dumps(event, indent=2)}")
+    print("-" * 50)
+
+    result = lambda_module.lambda_handler(event, {})
+
+    print("-" * 50)
+    print("Result:")
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

workbench/scripts/ml_pipeline_batch.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-batch-{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()

workbench/scripts/ml_pipeline_sqs.py

@@ -0,0 +1,186 @@
+import argparse
+import logging
+import json
+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
+
+log = logging.getLogger("workbench")
+cm = ConfigManager()
+workbench_bucket = cm.get_config("WORKBENCH_BUCKET")
+
+
+def submit_to_sqs(
+    script_path: str,
+    size: str = "small",
+    realtime: bool = False,
+    dt: bool = False,
+    promote: bool = False,
+) -> None:
+    """
+    Upload script to S3 and submit message to SQS queue for processing.
+
+    Args:
+        script_path: Local path to the ML pipeline script
+        size: Job size tier - "small" (default), "medium", or "large"
+        realtime: If True, sets serverless=False for real-time processing (default: False)
+        dt: If True, sets DT=True in environment (default: False)
+        promote: If True, sets PROMOTE=True in environment (default: False)
+
+    Raises:
+        ValueError: If size is invalid or script file not found
+    """
+    print(f"\n{'=' * 60}")
+    print("🚀  SUBMITTING ML PIPELINE JOB")
+    print(f"{'=' * 60}")
+    if size not in ["small", "medium", "large"]:
+        raise ValueError(f"Invalid size '{size}'. Must be 'small', 'medium', or 'large'")
+
+    # Validate script exists
+    script_file = Path(script_path)
+    if not script_file.exists():
+        raise FileNotFoundError(f"Script not found: {script_path}")
+
+    print(f"📄  Script: {script_file.name}")
+    print(f"📏  Size tier: {size}")
+    print(f"⚡  Mode: {'Real-time' if realtime else 'Serverless'} (serverless={'False' if realtime else 'True'})")
+    print(f"🔄  DynamicTraining: {dt}")
+    print(f"🆕  Promote: {promote}")
+    print(f"🪣  Bucket: {workbench_bucket}")
+    sqs = AWSAccountClamp().boto3_session.client("sqs")
+    script_name = script_file.name
+
+    # List Workbench queues
+    print("\n📋  Listing Workbench SQS queues...")
+    try:
+        queues = sqs.list_queues(QueueNamePrefix="workbench-")
+        queue_urls = queues.get("QueueUrls", [])
+        if queue_urls:
+            print(f"✅  Found {len(queue_urls)} workbench queue(s):")
+            for url in queue_urls:
+                queue_name = url.split("/")[-1]
+                print(f"   • {queue_name}")
+        else:
+            print("⚠️  No workbench queues found")
+    except Exception as e:
+        print(f"❌  Error listing queues: {e}")
+
+    # Upload script to S3
+    s3_path = f"s3://{workbench_bucket}/batch-jobs/{script_name}"
+    print("\n📤  Uploading script to S3...")
+    print(f"   Source: {script_path}")
+    print(f"   Destination: {s3_path}")
+
+    try:
+        upload_content_to_s3(script_file.read_text(), s3_path)
+        print("✅  Script uploaded successfully")
+    except Exception as e:
+        print(f"❌  Upload failed: {e}")
+        raise
+    # Get queue URL and info
+    queue_name = "workbench-ml-pipeline-queue.fifo"
+    print("\n🎯  Getting queue information...")
+    print(f"   Queue name: {queue_name}")
+
+    try:
+        queue_url = sqs.get_queue_url(QueueName=queue_name)["QueueUrl"]
+        print(f"   Queue URL: {queue_url}")
+
+        # Get queue attributes for additional info
+        attrs = sqs.get_queue_attributes(
+            QueueUrl=queue_url, AttributeNames=["ApproximateNumberOfMessages", "ApproximateNumberOfMessagesNotVisible"]
+        )
+        messages_available = attrs["Attributes"].get("ApproximateNumberOfMessages", "0")
+        messages_in_flight = attrs["Attributes"].get("ApproximateNumberOfMessagesNotVisible", "0")
+        print(f"   Messages in queue: {messages_available}")
+        print(f"   Messages in flight: {messages_in_flight}")
+
+    except Exception as e:
+        print(f"❌  Error accessing queue: {e}")
+        raise
+
+    # Prepare message
+    message = {"script_path": s3_path, "size": size}
+
+    # Set environment variables
+    message["environment"] = {
+        "SERVERLESS": "False" if realtime else "True",
+        "DT": str(dt),
+        "PROMOTE": str(promote),
+    }
+
+    # Send the message to SQS
+    try:
+        print("\n📨  Sending message to SQS...")
+        response = sqs.send_message(
+            QueueUrl=queue_url,
+            MessageBody=json.dumps(message, indent=2),
+            MessageGroupId="ml-pipeline-jobs",  # Required for FIFO
+        )
+        message_id = response["MessageId"]
+        print("✅  Message sent successfully!")
+        print(f"   Message ID: {message_id}")
+    except Exception as e:
+        print(f"❌  Failed to send message: {e}")
+        raise
+
+    # Success summary
+    print(f"\n{'=' * 60}")
+    print("✅  JOB SUBMISSION COMPLETE")
+    print(f"{'=' * 60}")
+    print(f"📄  Script: {script_name}")
+    print(f"📏  Size: {size}")
+    print(f"⚡  Mode: {'Real-time' if realtime else 'Serverless'} (SERVERLESS={'False' if realtime else 'True'})")
+    print(f"🔄  DynamicTraining: {dt}")
+    print(f"🆕  Promote: {promote}")
+    print(f"🆔  Message ID: {message_id}")
+    print("\n🔍  MONITORING LOCATIONS:")
+    print(f"   • SQS Queue: AWS Console → SQS → {queue_name}")
+    print("   • Lambda Logs: AWS Console → Lambda → Functions")
+    print("   • Batch Jobs: AWS Console → Batch → Jobs")
+    print("   • CloudWatch: AWS Console → CloudWatch → Log groups")
+    print("\n⏳  Your job should start processing soon...")
+
+
+def main():
+    """CLI entry point for submitting ML pipelines via SQS."""
+    parser = argparse.ArgumentParser(description="Submit ML pipeline to SQS queue for Batch processing")
+    parser.add_argument("script_file", help="Local path to ML pipeline script")
+    parser.add_argument(
+        "--size", default="small", choices=["small", "medium", "large"], help="Job size tier (default: small)"
+    )
+    parser.add_argument(
+        "--realtime",
+        action="store_true",
+        help="Create realtime endpoints (default is serverless)",
+    )
+    parser.add_argument(
+        "--dt",
+        action="store_true",
+        help="Set DT=True (models and endpoints will have '-dt' suffix)",
+    )
+    parser.add_argument(
+        "--promote",
+        action="store_true",
+        help="Set Promote=True (models and endpoints will use promoted naming",
+    )
+    args = parser.parse_args()
+    try:
+        submit_to_sqs(
+            args.script_file,
+            args.size,
+            realtime=args.realtime,
+            dt=args.dt,
+            promote=args.promote,
+        )
+    except Exception as e:
+        print(f"\n❌  ERROR: {e}")
+        log.error(f"Error: {e}")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

workbench/scripts/monitor_cloud_watch.py

@@ -4,8 +4,10 @@ import sys
 import time
 import argparse
 from datetime import datetime, timedelta, timezone
-from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+# Workbench Imports
 from workbench.utils.repl_utils import cprint, Spinner
+from workbench.utils.cloudwatch_utils import get_cloudwatch_client, get_active_log_streams, stream_log_events
 
 # Define the log levels to include all log levels above the specified level
 log_level_map = {
@@ -33,64 +35,6 @@ def date_display(dt):
         return dt.strftime("%Y-%m-%d %I:%M%p") + "(UTC)"
 
 
-def get_cloudwatch_client():
-    """Get the CloudWatch Logs client using the Workbench assumed role session."""
-    session = AWSAccountClamp().boto3_session
-    return session.client("logs")
-
-
-def get_active_log_streams(client, log_group_name, start_time_ms, stream_filter=None):
-    """Retrieve log streams that have events after the specified start time."""
-
-    # Get all the streams in the log group
-    active_streams = []
-    stream_params = {
-        "logGroupName": log_group_name,
-        "orderBy": "LastEventTime",
-        "descending": True,
-    }
-
-    # Loop to retrieve all log streams (maximum 50 per call)
-    while True:
-        response = client.describe_log_streams(**stream_params)
-        log_streams = response.get("logStreams", [])
-
-        for log_stream in log_streams:
-            log_stream_name = log_stream["logStreamName"]
-            last_event_timestamp = log_stream.get("lastEventTimestamp")
-
-            # Include streams with events since the specified start time
-            # Note: There's some issue where the last event timestamp is 'off'
-            #       so we're going to add 60 minutes from the last event timestamp
-            last_event_timestamp += 60 * 60 * 1000
-            if last_event_timestamp >= start_time_ms:
-                active_streams.append(log_stream_name)
-            else:
-                break  # Stop if we reach streams older than the start time
-
-        # Check if there are more streams to retrieve
-        if "nextToken" in response:
-            stream_params["nextToken"] = response["nextToken"]
-        else:
-            break
-
-    # Sort and report the active log streams
-    active_streams.sort()
-    if active_streams:
-        print("Active log streams:", len(active_streams))
-
-    # Filter the active streams by a substring if provided
-    if stream_filter and active_streams:
-        print(f"Filtering active log streams by '{stream_filter}'...")
-        active_streams = [stream for stream in active_streams if stream_filter in stream]
-
-    for stream in active_streams:
-        print(f"\t - {stream}")
-
-    # Return the active log streams
-    return active_streams
-
-
 def get_latest_log_events(client, log_group_name, start_time, end_time=None, stream_filter=None):
     """Retrieve the latest log events from the active/filtered log streams in a CloudWatch Logs group."""
 
@@ -99,11 +43,15 @@ def get_latest_log_events(client, log_group_name, start_time, end_time=None, str
         get_latest_log_events.first_run = True
 
     log_events = []
-    start_time_ms = int(start_time.timestamp() * 1000)  # Convert start_time to milliseconds
+    start_time_ms = int(start_time.timestamp() * 1000)
+
+    # Use the util function to get active streams
+    active_streams = get_active_log_streams(log_group_name, start_time_ms, stream_filter, client)
 
-    # Get the active log streams with events since start_time
-    active_streams = get_active_log_streams(client, log_group_name, start_time_ms, stream_filter)
     if active_streams:
+        print(f"Active log streams: {len(active_streams)}")
+        for stream in active_streams:
+            print(f"\t - {stream}")
         print(f"Processing log events from {date_display(start_time)} on {len(active_streams)} active log streams...")
         get_latest_log_events.first_run = False
     else:
@@ -114,50 +62,22 @@ def get_latest_log_events(client, log_group_name, start_time, end_time=None, str
             print("Monitoring for new events...")
         return log_events
 
-    # Iterate over the active streams and fetch log events
+    # Use the util function to stream events from each log stream
     for log_stream_name in active_streams:
-        params = {
-            "logGroupName": log_group_name,
-            "logStreamName": log_stream_name,
-            "startTime": start_time_ms,  # Use start_time in milliseconds
-            "startFromHead": True,  # Start from the nearest event to start_time
-        }
-        next_event_token = None
-        if end_time is not None:
-            params["endTime"] = int(end_time.timestamp() * 1000)
-
-        # Process the log events from this log stream
         spinner = Spinner("lightpurple", f"Pulling events from {log_stream_name}:")
         spinner.start()
         log_stream_events = 0
 
-        # Get the log events for the active log stream
-        while True:
-            if next_event_token:
-                params["nextToken"] = next_event_token
-                params.pop("startTime", None)  # Remove startTime when using nextToken
+        # Stream events using the util function
+        for event in stream_log_events(
+            log_group_name, log_stream_name, start_time, end_time, follow=False, client=client
+        ):
+            log_stream_events += 1
+            log_events.append(event)
 
-            # Fetch the log events (this call takes a while: optimize if we can)
-            events_response = client.get_log_events(**params)
+        spinner.stop()
+        print(f"Processed {log_stream_events} events from {log_stream_name} (Total: {len(log_events)})")
 
-            events = events_response.get("events", [])
-            for event in events:
-                event["logStreamName"] = log_stream_name
-
-            # Add the log stream events to our list of all log events
-            log_stream_events += len(events)
-            log_events.extend(events)
-
-            # Handle pagination for log events
-            next_event_token = events_response.get("nextForwardToken")
-
-            # Break the loop if there are no more events to fetch
-            if not next_event_token or next_event_token == params.get("nextToken"):
-                spinner.stop()
-                print(f"Processed {log_stream_events} events from {log_stream_name} (Total: {len(log_events)})")
-                break
-
-    # Return the log events
     return log_events
 
 
@@ -206,6 +126,7 @@ def monitor_log_group(
     print(f"Monitoring log group: {log_group_name} from {date_display(start_time)}")
     print(f"Log levels: {log_levels}")
     print(f"Search terms: {search_terms}")
+
     while True:
         # Get the latest log events with stream filtering if provided
         all_log_events = get_latest_log_events(client, log_group_name, start_time, end_time, stream_filter)
@@ -218,7 +139,6 @@ def monitor_log_group(
 
                 # Check the search terms
                 if not search_terms or any(term in event["message"].lower() for term in search_terms):
-
                     # Calculate the start and end index for this match
                     start_index = max(i - before, 0)
                     end_index = min(i + after, len(all_log_events) - 1)

workbench/utils/aws_utils.py

@@ -55,7 +55,8 @@ def aws_throttle(func=None, retry_intervals=None):
     if func is None:
         return lambda f: aws_throttle(f, retry_intervals=retry_intervals)
 
-    service_hold_time = 2  # Seconds to wait before calling AWS function
+    # This is currently commented out (we might want to use it later)
+    # service_hold_time = 2  # Seconds to wait before calling AWS function
     default_intervals = [2**i for i in range(1, 9)]  # Default exponential backoff: 2, 4, 8... 256 seconds
     intervals = retry_intervals or default_intervals
 
@@ -64,8 +65,8 @@ def aws_throttle(func=None, retry_intervals=None):
         for attempt, delay in enumerate(intervals, start=1):
             try:
                 # Add sleep before calling AWS func if running as a service
-                if cm.running_as_service:
-                    time.sleep(service_hold_time)
+                # if cm.running_as_service:
+                #    time.sleep(service_hold_time)
                 return func(*args, **kwargs)
             except ClientError as e:
                 if e.response["Error"]["Code"] == "ThrottlingException":