workbench 0.8.174__py3-none-any.whl → 0.8.227__py3-none-any.whl

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
@@ -293,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")
@@ -525,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
 
@@ -560,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_sqs.py

@@ -1,6 +1,8 @@
 import argparse
+import ast
 import logging
 import json
+import re
 from pathlib import Path
 
 # Workbench Imports
@@ -13,27 +15,105 @@ cm = ConfigManager()
 workbench_bucket = cm.get_config("WORKBENCH_BUCKET")
 
 
-def submit_to_sqs(script_path: str, size: str = "small") -> None:
+def parse_workbench_batch(script_content: str) -> dict | None:
+    """Parse WORKBENCH_BATCH config from a script.
+
+    Looks for a dictionary assignment like:
+        WORKBENCH_BATCH = {
+            "outputs": ["feature_set_xyz"],
+        }
+    or:
+        WORKBENCH_BATCH = {
+            "inputs": ["feature_set_xyz"],
+        }
+
+    Args:
+        script_content: The Python script content as a string
+
+    Returns:
+        The parsed dictionary or None if not found
+    """
+    pattern = r"WORKBENCH_BATCH\s*=\s*(\{[^}]+\})"
+    match = re.search(pattern, script_content, re.DOTALL)
+    if match:
+        try:
+            return ast.literal_eval(match.group(1))
+        except (ValueError, SyntaxError) as e:
+            print(f"⚠️  Warning: Failed to parse WORKBENCH_BATCH: {e}")
+            return None
+    return None
+
+
+def get_message_group_id(batch_config: dict | None) -> str:
+    """Derive MessageGroupId from outputs or inputs.
+
+    - Scripts with outputs use first output as group
+    - Scripts with inputs use first input as group
+    - Default to "ml-pipeline-jobs" if no config
+    """
+    if not batch_config:
+        return "ml-pipeline-jobs"
+
+    outputs = batch_config.get("outputs", [])
+    inputs = batch_config.get("inputs", [])
+
+    if outputs:
+        return outputs[0]
+    elif inputs:
+        return inputs[0]
+    else:
+        return "ml-pipeline-jobs"
+
+
+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}")
 
+    # Read script content and parse WORKBENCH_BATCH config
+    script_content = script_file.read_text()
+    batch_config = parse_workbench_batch(script_content)
+    group_id = get_message_group_id(batch_config)
+    outputs = (batch_config or {}).get("outputs", [])
+    inputs = (batch_config or {}).get("inputs", [])
+
     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}")
+    if outputs:
+        print(f"📤  Outputs: {outputs}")
+    if inputs:
+        print(f"📥  Inputs: {inputs}")
+    print(f"📦  Batch Group: {group_id}")
     sqs = AWSAccountClamp().boto3_session.client("sqs")
     script_name = script_file.name
 
@@ -59,7 +139,7 @@ def submit_to_sqs(script_path: str, size: str = "small") -> None:
     print(f"   Destination: {s3_path}")
 
     try:
-        upload_content_to_s3(script_file.read_text(), s3_path)
+        upload_content_to_s3(script_content, s3_path)
         print("✅  Script uploaded successfully")
     except Exception as e:
         print(f"❌  Upload failed: {e}")
@@ -88,14 +168,21 @@ def submit_to_sqs(script_path: str, size: str = "small") -> None:
 
     # Prepare message
     message = {"script_path": s3_path, "size": size}
-    print("\n📨  Sending message to SQS...")
+
+    # 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
+            MessageGroupId=group_id,  # From WORKBENCH_BATCH or default
         )
         message_id = response["MessageId"]
         print("✅  Message sent successfully!")
@@ -110,6 +197,14 @@ def submit_to_sqs(script_path: str, size: str = "small") -> None:
     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}")
+    if outputs:
+        print(f"📤  Outputs: {outputs}")
+    if inputs:
+        print(f"📥  Inputs: {inputs}")
+    print(f"📦  Batch Group: {group_id}")
     print(f"🆔  Message ID: {message_id}")
     print("\n🔍  MONITORING LOCATIONS:")
     print(f"   • SQS Queue: AWS Console → SQS → {queue_name}")
@@ -126,9 +221,30 @@ def main():
     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)
+        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}")

workbench/scripts/training_test.py

@@ -0,0 +1,85 @@
+"""
+Local test harness for SageMaker training scripts.
+
+Usage:
+    python training_test.py <model_script.py> <featureset_name>
+
+Example:
+    python training_test.py ../model_scripts/pytorch_model/generated_model_script.py caco2-class-features
+"""
+
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+
+import pandas as pd
+
+from workbench.api import FeatureSet
+
+
+def get_training_data(featureset_name: str) -> pd.DataFrame:
+    """Get training data from the FeatureSet."""
+    fs = FeatureSet(featureset_name)
+    return fs.pull_dataframe()
+
+
+def main():
+    if len(sys.argv) < 3:
+        print("Usage: python training_test.py <model_script.py> <featureset_name>")
+        sys.exit(1)
+
+    script_path = sys.argv[1]
+    featureset_name = sys.argv[2]
+
+    if not os.path.exists(script_path):
+        print(f"Error: Script not found: {script_path}")
+        sys.exit(1)
+
+    # Create temp directories
+    model_dir = tempfile.mkdtemp(prefix="training_model_")
+    train_dir = tempfile.mkdtemp(prefix="training_data_")
+    output_dir = tempfile.mkdtemp(prefix="training_output_")
+
+    print(f"Model dir: {model_dir}")
+    print(f"Train dir: {train_dir}")
+
+    try:
+        # Get training data and save to CSV
+        print(f"Loading FeatureSet: {featureset_name}")
+        df = get_training_data(featureset_name)
+        print(f"Data shape: {df.shape}")
+
+        train_file = os.path.join(train_dir, "training_data.csv")
+        df.to_csv(train_file, index=False)
+
+        # Set up environment
+        env = os.environ.copy()
+        env["SM_MODEL_DIR"] = model_dir
+        env["SM_CHANNEL_TRAIN"] = train_dir
+        env["SM_OUTPUT_DATA_DIR"] = output_dir
+
+        print("\n" + "=" * 60)
+        print("Starting training...")
+        print("=" * 60 + "\n")
+
+        # Run the script
+        cmd = [sys.executable, script_path, "--model-dir", model_dir, "--train", train_dir]
+        result = subprocess.run(cmd, env=env)
+
+        print("\n" + "=" * 60)
+        if result.returncode == 0:
+            print("Training completed successfully!")
+        else:
+            print(f"Training failed with return code: {result.returncode}")
+        print("=" * 60)
+
+    finally:
+        shutil.rmtree(model_dir, ignore_errors=True)
+        shutil.rmtree(train_dir, ignore_errors=True)
+        shutil.rmtree(output_dir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    main()

workbench/themes/dark/custom.css

@@ -110,6 +110,40 @@ a:hover {
     color: rgb(100, 255, 100);
 }
 
+/* Dropdown styling (dcc.Dropdown) - override Bootstrap's variables */
+.dash-dropdown {
+    --bs-body-bg: rgb(35, 35, 35);
+    --bs-body-color: rgb(210, 210, 210);
+    --bs-border-color: rgb(60, 60, 60);
+}
+
+/* Bootstrap form controls (dbc components) */
+.form-select, .form-control {
+    background-color: rgb(35, 35, 35) !important;
+    border: 1px solid rgb(60, 60, 60) !important;
+    color: rgb(210, 210, 210) !important;
+}
+
+.form-select:focus, .form-control:focus {
+    background-color: rgb(45, 45, 45) !important;
+    border-color: rgb(80, 80, 80) !important;
+    box-shadow: 0 0 0 0.2rem rgba(80, 80, 80, 0.25) !important;
+}
+
+.dropdown-menu {
+    background-color: rgb(35, 35, 35) !important;
+    border: 1px solid rgb(60, 60, 60) !important;
+}
+
+.dropdown-item {
+    color: rgb(210, 210, 210) !important;
+}
+
+.dropdown-item:hover, .dropdown-item:focus {
+    background-color: rgb(50, 50, 50) !important;
+    color: rgb(230, 230, 230) !important;
+}
+
 /* Table styling */
 table {
     width: 100%;
@@ -128,4 +162,29 @@ td {
     padding: 5px;
     border: 0.5px solid #444;
     text-align: center !important;
+}
+
+/* AG Grid table header colors - gradient theme */
+/* Data Sources tables - red gradient */
+#main_data_sources .ag-header,
+#data_sources_table .ag-header {
+    background: linear-gradient(180deg, rgb(140, 60, 60) 0%, rgb(80, 35, 35) 100%) !important;
+}
+
+/* Feature Sets tables - yellow/olive gradient */
+#main_feature_sets .ag-header,
+#feature_sets_table .ag-header {
+    background: linear-gradient(180deg, rgb(120, 115, 55) 0%, rgb(70, 65, 30) 100%) !important;
+}
+
+/* Models tables - green gradient */
+#main_models .ag-header,
+#models_table .ag-header {
+    background: linear-gradient(180deg, rgb(55, 110, 55) 0%, rgb(30, 60, 30) 100%) !important;
+}
+
+/* Endpoints tables - purple gradient */
+#main_endpoints .ag-header,
+#endpoints_table .ag-header {
+    background: linear-gradient(180deg, rgb(100, 60, 120) 0%, rgb(55, 30, 70) 100%) !important;
 }