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

workbench/repl/workbench_shell.py

@@ -1,14 +1,25 @@
+# flake8: noqa: E402
+import os
+import sys
+import logging
+import importlib
+import webbrowser
+import readline  # noqa: F401
+
+# Disable OpenMP parallelism to avoid segfaults with PyTorch in iPython
+# This is a known issue on macOS where libomp crashes during thread synchronization
+# Must be set before importing numpy/pandas/torch or any library that uses OpenMP
+os.environ.setdefault("OMP_NUM_THREADS", "1")
+os.environ.setdefault("MKL_NUM_THREADS", "1")
+
+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
-import sys
-import logging
-import importlib
 import botocore
-import webbrowser
 import pandas as pd
-import readline  # noqa
 
 try:
     import matplotlib.pyplot as plt  # noqa
@@ -39,7 +50,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 +81,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 +187,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 +213,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:")
@@ -288,11 +302,6 @@ class WorkbenchShell:
             self.commands["PandasToView"] = importlib.import_module("workbench.core.views.pandas_to_view").PandasToView
             self.commands["Pipeline"] = importlib.import_module("workbench.api.pipeline").Pipeline
 
-            # Algorithms
-            self.commands["FSP"] = importlib.import_module(
-                "workbench.algorithms.dataframe.feature_space_proximity"
-            ).FeatureSpaceProximity
-
             # These are 'nice to have' imports
             self.commands["pd"] = importlib.import_module("pandas")
             self.commands["wr"] = importlib.import_module("awswrangler")
@@ -520,7 +529,7 @@ class WorkbenchShell:
     def get_meta(self):
         return self.meta
 
-    def plot_manager(self, data, plot_type: str = "table", **kwargs):
+    def plot_manager(self, data, plot_type: str = "scatter", **kwargs):
         """Plot Manager for Workbench"""
         from workbench.web_interface.components.plugins import ag_table, graph_plot, scatter_plot
 
@@ -555,14 +564,14 @@ 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)
 
-        # Run the server and open in the browser
-        plugin_test.run()
+        # Open the browser and run the dash server
         url = f"http://127.0.0.1:{plugin_test.port}"
         webbrowser.open(url)
+        plugin_test.run()
 
 
 # Launch Shell Entry Point

workbench/resources/open_source_api.key

@@ -1 +1 @@
-eyJsaWNlbnNlX2lkIjogIk9wZW5fU291cmNlX0xpY2Vuc2UiLCAiY29tcGFueSI6ICIiLCAiYXdzX2FjY291bnRfaWQiOiAiIiwgInRpZXIiOiAiRW50ZXJwcmlzZSBQcm8iLCAiZmVhdHVyZXMiOiBbInBsdWdpbnMiLCAicGFnZXMiLCAidGhlbWVzIiwgInBpcGVsaW5lcyIsICJicmFuZGluZyJdLCAiZXhwaXJlcyI6ICIyMDI2LTAxLTE0In02zCDRy41wKRViRnGmodczFWexLyfXYrJWSuVQQbhbWeRttQRv6zpo9x4O2yBjdRfhb9E7mFUppNiOS_ZGK-bL71nGHt_Mc8niG8jkpvKX9qZ6BqkXF_vzDIOcI8iGiwB3wikeVO4zRLD1AI0U3cgYmIyGXI9QKJ9L7IHyQ0TWqw==
+eyJsaWNlbnNlX2lkIjogIk9wZW5fU291cmNlX0xpY2Vuc2UiLCAiY29tcGFueSI6ICIiLCAiYXdzX2FjY291bnRfaWQiOiAiIiwgInRpZXIiOiAiRW50ZXJwcmlzZSBQcm8iLCAiZmVhdHVyZXMiOiBbInBsdWdpbnMiLCAicGFnZXMiLCAidGhlbWVzIiwgInBpcGVsaW5lcyIsICJicmFuZGluZyJdLCAiZXhwaXJlcyI6ICIyMDI2LTEyLTA1In1IsmpkuybFALADkRj_RfmkQ0LAIsQeXRE7Uoc3DL1UrDr-rSnwu-PDqsKBUkX6jPRFZV3DLxNjBapxPeEIFhfvxvjzz_sc6CwtxNpZ3bPmxSPs2W-j3xZS4-XyEqIilcwSkWh-NU1u27gCuuivn5eiUmIYJGAp0wdVkeE6_Z9dlg==

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-reg-pytorch
+
+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-reg-pytorch)")
+        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/meta_model_sim.py

@@ -0,0 +1,35 @@
+"""MetaModelSimulator: Simulate and analyze ensemble model performance.
+
+This class helps evaluate whether a meta model (ensemble) would outperform
+individual child models by analyzing endpoint inference predictions.
+"""
+
+import argparse
+from workbench.utils.meta_model_simulator import MetaModelSimulator
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Simulate and analyze ensemble model performance using MetaModelSimulator."
+    )
+    parser.add_argument(
+        "models",
+        nargs="+",
+        help="List of model endpoint names to include in the ensemble simulation.",
+    )
+    parser.add_argument(
+        "--id-column",
+        default="molecule_name",
+        help="Name of the ID column (default: molecule_name)",
+    )
+    args = parser.parse_args()
+    models = args.models
+    id_column = args.id_column
+
+    # Create MetaModelSimulator instance and generate report
+    sim = MetaModelSimulator(models, id_column=id_column)
+    sim.report()
+
+
+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()