workbench 0.8.162__py3-none-any.whl → 0.8.202__py3-none-any.whl

workbench/scripts/endpoint_test.py

@@ -0,0 +1,162 @@
+"""
+Local test harness for SageMaker model scripts.
+
+Usage:
+    python model_script_harness.py <local_script.py> <model_name>
+
+Example:
+    python model_script_harness.py pytorch.py aqsol-pytorch-reg
+
+This allows you to test LOCAL changes to a model script against deployed model artifacts.
+Evaluation data is automatically pulled from the FeatureSet (training = FALSE rows).
+
+Optional: testing/env.json with additional environment variables
+"""
+
+import os
+import sys
+import json
+import importlib.util
+import tempfile
+import shutil
+import pandas as pd
+import torch
+
+# Workbench Imports
+from workbench.api import Model, FeatureSet
+from workbench.utils.pytorch_utils import download_and_extract_model
+
+# Force CPU mode BEFORE any PyTorch imports to avoid MPS/CUDA issues on Mac
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+torch.set_default_device("cpu")
+# Disable MPS entirely
+if hasattr(torch.backends, "mps"):
+    torch.backends.mps.is_available = lambda: False
+
+
+def get_eval_data(workbench_model: Model) -> pd.DataFrame:
+    """Get evaluation data from the FeatureSet associated with this model."""
+    # Get the FeatureSet
+    fs_name = workbench_model.get_input()
+    fs = FeatureSet(fs_name)
+    if not fs.exists():
+        raise ValueError(f"No FeatureSet found: {fs_name}")
+
+    # Get evaluation data (training = FALSE)
+    table = workbench_model.training_view().table
+    print(f"Querying evaluation data from {table}...")
+    eval_df = fs.query(f'SELECT * FROM "{table}" WHERE training = FALSE')
+    print(f"Retrieved {len(eval_df)} evaluation rows")
+
+    return eval_df
+
+
+def load_model_script(script_path: str):
+    """Dynamically load the model script module."""
+    if not os.path.exists(script_path):
+        raise FileNotFoundError(f"Script not found: {script_path}")
+
+    spec = importlib.util.spec_from_file_location("model_script", script_path)
+    module = importlib.util.module_from_spec(spec)
+
+    # Add to sys.modules so imports within the script work
+    sys.modules["model_script"] = module
+
+    spec.loader.exec_module(module)
+    return module
+
+
+def main():
+    if len(sys.argv) < 3:
+        print("Usage: python model_script_harness.py <local_script.py> <model_name>")
+        print("\nArguments:")
+        print("  local_script.py  - Path to your LOCAL model script to test")
+        print("  model_name       - Workbench model name (e.g., aqsol-pytorch-reg)")
+        print("\nOptional: testing/env.json with additional environment variables")
+        sys.exit(1)
+
+    script_path = sys.argv[1]
+    model_name = sys.argv[2]
+
+    # Validate local script exists
+    if not os.path.exists(script_path):
+        print(f"Error: Local script not found: {script_path}")
+        sys.exit(1)
+
+    # Initialize Workbench model
+    print(f"Loading Workbench model: {model_name}")
+    workbench_model = Model(model_name)
+    print(f"Model Framework: {workbench_model.model_framework}")
+    print()
+
+    # Create a temporary model directory
+    model_dir = tempfile.mkdtemp(prefix="model_harness_")
+    print(f"Using model directory: {model_dir}")
+
+    try:
+        # 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()
+
+        # Set up SageMaker environment variables
+        os.environ["SM_MODEL_DIR"] = model_dir
+        print(f"Set SM_MODEL_DIR = {model_dir}")
+
+        # Download and extract model artifacts
+        s3_uri = workbench_model.model_data_url()
+        download_and_extract_model(s3_uri, model_dir)
+        print()
+
+        # Load the LOCAL model script
+        print(f"Loading LOCAL model script: {script_path}")
+        module = load_model_script(script_path)
+        print()
+
+        # Check for required functions
+        if not hasattr(module, "model_fn"):
+            raise AttributeError("Model script must have a model_fn function")
+        if not hasattr(module, "predict_fn"):
+            raise AttributeError("Model script must have a predict_fn function")
+
+        # Load the model
+        print("Calling model_fn...")
+        print("-" * 50)
+        model = module.model_fn(model_dir)
+        print("-" * 50)
+        print(f"Model loaded: {type(model)}")
+        print()
+
+        # Get evaluation data from FeatureSet
+        print("Pulling evaluation data from FeatureSet...")
+        df = get_eval_data(workbench_model)
+        print(f"Input shape: {df.shape}")
+        print(f"Columns: {df.columns.tolist()}")
+        print()
+
+        print("Calling predict_fn...")
+        print("-" * 50)
+        result = module.predict_fn(df, model)
+        print("-" * 50)
+        print()
+
+        print("Prediction result:")
+        print(f"Output shape: {result.shape}")
+        print(f"Output columns: {result.columns.tolist()}")
+        print()
+        print(result.head(10).to_string())
+
+    finally:
+        # Cleanup
+        print(f"\nCleaning up model directory: {model_dir}")
+        shutil.rmtree(model_dir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    main()

workbench/scripts/lambda_test.py

@@ -0,0 +1,73 @@
+"""
+Local test harness for AWS Lambda scripts.
+
+Usage:
+    lambda_test <lambda_script.py>
+
+Required: testing/event.json with the event definition
+Options: testing/env.json with a set of ENV vars
+"""
+
+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()