workbench 0.8.198__py3-none-any.whl → 0.8.203__py3-none-any.whl

workbench/model_scripts/xgb_model/generated_model_script.py

@@ -6,11 +6,13 @@ import numpy as np
 # Model Performance Scores
 from sklearn.metrics import (
     mean_absolute_error,
+    median_absolute_error,
     r2_score,
     root_mean_squared_error,
     precision_recall_fscore_support,
     confusion_matrix,
 )
+from scipy.stats import spearmanr
 
 # Classification Encoder
 from sklearn.preprocessing import LabelEncoder
@@ -26,6 +28,28 @@ import os
 import pandas as pd
 from typing import List, Tuple
 
+# Default Hyperparameters for XGBoost
+DEFAULT_HYPERPARAMETERS = {
+    # Core tree parameters
+    "n_estimators": 200,  # More trees for better signal capture when we have lots of features
+    "max_depth": 6,  # Medium depth
+    "learning_rate": 0.05,  # Lower rate with more estimators for smoother learning
+
+    # Sampling parameters
+    "subsample": 0.7,  # Moderate row sampling to reduce overfitting
+    "colsample_bytree": 0.6,  # More aggressive feature sampling given lots of features
+    "colsample_bylevel": 0.8,  # Additional feature sampling at each tree level
+
+    # Regularization
+    "min_child_weight": 5,  # Higher to prevent overfitting on small groups
+    "gamma": 0.2,  # Moderate pruning - you have real signal so don't over-prune
+    "reg_alpha": 0.5,  # L1 for feature selection (useful with many features)
+    "reg_lambda": 2.0,  # Strong L2 to smooth predictions
+
+    # Random seed
+    "random_state": 42,
+}
+
 # Template Parameters
 TEMPLATE_PARAMS = {
     "model_type": "classifier",
@@ -208,7 +232,7 @@ if __name__ == "__main__":
     model_type = TEMPLATE_PARAMS["model_type"]
     model_metrics_s3_path = TEMPLATE_PARAMS["model_metrics_s3_path"]
     train_all_data = TEMPLATE_PARAMS["train_all_data"]
-    hyperparameters = TEMPLATE_PARAMS["hyperparameters"]
+    hyperparameters = {**DEFAULT_HYPERPARAMETERS, **TEMPLATE_PARAMS["hyperparameters"]}
     validation_split = 0.2
 
     # Script arguments for input/output directories
@@ -325,13 +349,13 @@ if __name__ == "__main__":
                 target: label_names,
                 "precision": scores[0],
                 "recall": scores[1],
-                "fscore": scores[2],
+                "f1": scores[2],
                 "support": scores[3],
             }
         )
 
         # We need to get creative with the Classification Metrics
-        metrics = ["precision", "recall", "fscore", "support"]
+        metrics = ["precision", "recall", "f1", "support"]
         for t in label_names:
             for m in metrics:
                 value = score_df.loc[score_df[target] == t, m].iloc[0]
@@ -348,11 +372,16 @@ if __name__ == "__main__":
         # Calculate various model performance metrics (regression)
         rmse = root_mean_squared_error(y_validate, preds)
         mae = mean_absolute_error(y_validate, preds)
+        medae = median_absolute_error(y_validate, preds)
         r2 = r2_score(y_validate, preds)
-        print(f"RMSE: {rmse:.3f}")
-        print(f"MAE: {mae:.3f}")
-        print(f"R2: {r2:.3f}")
-        print(f"NumRows: {len(df_val)}")
+        spearman_corr = spearmanr(y_validate, preds).correlation
+        support = len(df_val)
+        print(f"rmse: {rmse:.3f}")
+        print(f"mae: {mae:.3f}")
+        print(f"medae: {medae:.3f}")
+        print(f"r2: {r2:.3f}")
+        print(f"spearmanr: {spearman_corr:.3f}")
+        print(f"support: {support}")
 
     # Now save the model to the standard place/name
     joblib.dump(xgb_model, os.path.join(args.model_dir, "xgb_model.joblib"))

workbench/model_scripts/xgb_model/xgb_model.template

@@ -6,11 +6,13 @@ import numpy as np
 # Model Performance Scores
 from sklearn.metrics import (
     mean_absolute_error,
+    median_absolute_error,
     r2_score,
     root_mean_squared_error,
     precision_recall_fscore_support,
     confusion_matrix,
 )
+from scipy.stats import spearmanr
 
 # Classification Encoder
 from sklearn.preprocessing import LabelEncoder
@@ -26,6 +28,28 @@ import os
 import pandas as pd
 from typing import List, Tuple
 
+# Default Hyperparameters for XGBoost
+DEFAULT_HYPERPARAMETERS = {
+    # Core tree parameters
+    "n_estimators": 200,  # More trees for better signal capture when we have lots of features
+    "max_depth": 6,  # Medium depth
+    "learning_rate": 0.05,  # Lower rate with more estimators for smoother learning
+
+    # Sampling parameters
+    "subsample": 0.7,  # Moderate row sampling to reduce overfitting
+    "colsample_bytree": 0.6,  # More aggressive feature sampling given lots of features
+    "colsample_bylevel": 0.8,  # Additional feature sampling at each tree level
+
+    # Regularization
+    "min_child_weight": 5,  # Higher to prevent overfitting on small groups
+    "gamma": 0.2,  # Moderate pruning - you have real signal so don't over-prune
+    "reg_alpha": 0.5,  # L1 for feature selection (useful with many features)
+    "reg_lambda": 2.0,  # Strong L2 to smooth predictions
+
+    # Random seed
+    "random_state": 42,
+}
+
 # Template Parameters
 TEMPLATE_PARAMS = {
     "model_type": "{{model_type}}",
@@ -208,7 +232,7 @@ if __name__ == "__main__":
     model_type = TEMPLATE_PARAMS["model_type"]
     model_metrics_s3_path = TEMPLATE_PARAMS["model_metrics_s3_path"]
     train_all_data = TEMPLATE_PARAMS["train_all_data"]
-    hyperparameters = TEMPLATE_PARAMS["hyperparameters"]
+    hyperparameters = {**DEFAULT_HYPERPARAMETERS, **TEMPLATE_PARAMS["hyperparameters"]}
     validation_split = 0.2
 
     # Script arguments for input/output directories
@@ -325,13 +349,13 @@ if __name__ == "__main__":
                 target: label_names,
                 "precision": scores[0],
                 "recall": scores[1],
-                "fscore": scores[2],
+                "f1": scores[2],
                 "support": scores[3],
             }
         )
 
         # We need to get creative with the Classification Metrics
-        metrics = ["precision", "recall", "fscore", "support"]
+        metrics = ["precision", "recall", "f1", "support"]
         for t in label_names:
             for m in metrics:
                 value = score_df.loc[score_df[target] == t, m].iloc[0]
@@ -348,11 +372,16 @@ if __name__ == "__main__":
         # Calculate various model performance metrics (regression)
         rmse = root_mean_squared_error(y_validate, preds)
         mae = mean_absolute_error(y_validate, preds)
+        medae = median_absolute_error(y_validate, preds)
         r2 = r2_score(y_validate, preds)
-        print(f"RMSE: {rmse:.3f}")
-        print(f"MAE: {mae:.3f}")
-        print(f"R2: {r2:.3f}")
-        print(f"NumRows: {len(df_val)}")
+        spearman_corr = spearmanr(y_validate, preds).correlation
+        support = len(df_val)
+        print(f"rmse: {rmse:.3f}")
+        print(f"mae: {mae:.3f}")
+        print(f"medae: {medae:.3f}")
+        print(f"r2: {r2:.3f}")
+        print(f"spearmanr: {spearman_corr:.3f}")
+        print(f"support: {support}")
 
     # Now save the model to the standard place/name
     joblib.dump(xgb_model, os.path.join(args.model_dir, "xgb_model.joblib"))

workbench/repl/workbench_shell.py

@@ -1,16 +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

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-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_launcher.py → lambda_test.py}

@@ -1,3 +1,13 @@
+"""
+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