workbench 0.8.234__py3-none-any.whl → 0.8.239__py3-none-any.whl

workbench/scripts/ml_pipeline_launcher.py

@@ -0,0 +1,410 @@
+"""Launch ML pipelines via SQS for testing.
+
+Run this from a directory containing pipeline subdirectories (e.g., ml_pipelines/).
+
+Usage:
+    ml_pipeline_launcher --dt                    # Launch 1 random pipeline group (all scripts in a directory)
+    ml_pipeline_launcher --dt -n 3               # Launch 3 random pipeline groups
+    ml_pipeline_launcher --dt --all              # Launch ALL pipelines
+    ml_pipeline_launcher --dt caco2              # Launch pipelines matching 'caco2'
+    ml_pipeline_launcher --dt caco2 ppb          # Launch pipelines matching 'caco2' or 'ppb'
+    ml_pipeline_launcher --promote --all         # Promote ALL pipelines
+    ml_pipeline_launcher --test-promote --all    # Test-promote ALL pipelines
+    ml_pipeline_launcher --dt --dry-run          # Show what would be launched without launching
+"""
+
+import argparse
+import ast
+import random
+import re
+import subprocess
+import time
+from pathlib import Path
+
+
+def parse_workbench_batch(script_path: Path) -> dict | None:
+    """Parse WORKBENCH_BATCH config from a script file."""
+    content = script_path.read_text()
+    match = re.search(r"WORKBENCH_BATCH\s*=\s*(\{[^}]+\})", content, re.DOTALL)
+    if match:
+        try:
+            return ast.literal_eval(match.group(1))
+        except (ValueError, SyntaxError):
+            return None
+    return None
+
+
+def build_dependency_graph(configs: dict[Path, dict]) -> dict[str, str]:
+    """Build a mapping from each output to its root producer.
+
+    For a chain like A -> B -> C (where B depends on A, C depends on B),
+    this returns {A: A, B: A, C: A} so all are in the same message group.
+    """
+    # Build output -> input mapping (what does each output depend on?)
+    output_to_input = {}
+    for config in configs.values():
+        if not config:
+            continue
+        outputs = config.get("outputs", [])
+        inputs = config.get("inputs", [])
+        for output in outputs:
+            output_to_input[output] = inputs[0] if inputs else None
+
+    # Walk chain to find root
+    def find_root(output: str, visited: set = None) -> str:
+        if visited is None:
+            visited = set()
+        if output in visited:
+            return output
+        visited.add(output)
+        parent = output_to_input.get(output)
+        if parent is None:
+            return output
+        return find_root(parent, visited)
+
+    return {output: find_root(output) for output in output_to_input}
+
+
+def get_group_id(config: dict | None, root_map: dict[str, str]) -> str | None:
+    """Get the root group_id for a pipeline based on its config and root_map."""
+    if not config:
+        return None
+    outputs = config.get("outputs", [])
+    inputs = config.get("inputs", [])
+    # Check inputs first (this script depends on something)
+    if inputs and inputs[0] in root_map:
+        return root_map[inputs[0]]
+    # Check outputs (this script produces something)
+    if outputs and outputs[0] in root_map:
+        return root_map[outputs[0]]
+    return None
+
+
+def sort_by_dependencies(pipelines: list[Path]) -> tuple[list[Path], dict[Path, dict], dict[str, str]]:
+    """Sort pipelines by dependency chains. Returns (sorted_list, configs, root_map)."""
+    # Parse all configs
+    configs = {}
+    for pipeline in pipelines:
+        configs[pipeline] = parse_workbench_batch(pipeline)
+
+    # Build root map for group_id resolution
+    root_map = build_dependency_graph(configs)
+
+    # Build output -> pipeline mapping
+    output_to_pipeline = {}
+    for pipeline, config in configs.items():
+        if config and config.get("outputs"):
+            for output in config["outputs"]:
+                output_to_pipeline[output] = pipeline
+
+    # Build chains by walking from root producers
+    sorted_pipelines = []
+    used = set()
+
+    for pipeline in sorted(pipelines):
+        config = configs.get(pipeline)
+
+        # Skip if already used or has inputs (not a root)
+        if pipeline in used:
+            continue
+        if config and config.get("inputs"):
+            continue
+
+        # Walk the chain from this root
+        chain = [pipeline]
+        used.add(pipeline)
+
+        current = pipeline
+        while True:
+            current_config = configs.get(current)
+            if not current_config or not current_config.get("outputs"):
+                break
+
+            current_output = current_config["outputs"][0]
+            # Find pipeline that consumes this output
+            next_pipeline = None
+            for p, c in configs.items():
+                if p in used or p not in pipelines:
+                    continue
+                if c and c.get("inputs") and current_output in c["inputs"]:
+                    next_pipeline = p
+                    break
+
+            if next_pipeline:
+                chain.append(next_pipeline)
+                used.add(next_pipeline)
+                current = next_pipeline
+            else:
+                break
+
+        sorted_pipelines.extend(chain)
+
+    # Add any remaining pipelines not in chains
+    for pipeline in sorted(pipelines):
+        if pipeline not in used:
+            sorted_pipelines.append(pipeline)
+
+    return sorted_pipelines, configs, root_map
+
+
+def format_dependency_chains(pipelines: list[Path], configs: dict[Path, dict]) -> list[str]:
+    """Format pipelines as dependency chains for display."""
+    # Build output -> pipeline mapping
+    output_to_pipeline = {}
+    for pipeline, config in configs.items():
+        if config and config.get("outputs"):
+            for output in config["outputs"]:
+                output_to_pipeline[output] = pipeline
+
+    # Build chains by walking from root producers
+    chains = []
+    used = set()
+
+    for pipeline in pipelines:
+        config = configs.get(pipeline)
+
+        # Skip if already part of a chain or has inputs (not a root)
+        if pipeline in used:
+            continue
+        if config and config.get("inputs"):
+            continue
+
+        # Start a new chain from this root producer (or standalone)
+        chain = [pipeline]
+        used.add(pipeline)
+
+        # Walk the chain: find who consumes our output
+        current = pipeline
+        while True:
+            current_config = configs.get(current)
+            if not current_config or not current_config.get("outputs"):
+                break
+
+            current_output = current_config["outputs"][0]
+            # Find a pipeline that takes this output as input
+            next_pipeline = None
+            for p, c in configs.items():
+                if p in used or p not in pipelines:
+                    continue
+                if c and c.get("inputs") and current_output in c["inputs"]:
+                    next_pipeline = p
+                    break
+
+            if next_pipeline:
+                chain.append(next_pipeline)
+                used.add(next_pipeline)
+                current = next_pipeline
+            else:
+                break
+
+        chains.append(chain)
+
+    # Add any remaining pipelines not in chains (shouldn't happen but just in case)
+    for pipeline in pipelines:
+        if pipeline not in used:
+            chains.append([pipeline])
+
+    # Format chains as strings
+    lines = []
+    for chain in chains:
+        names = [p.stem for p in chain]
+        lines.append("   " + " --> ".join(names))
+
+    return lines
+
+
+def get_all_pipelines() -> list[Path]:
+    """Get all ML pipeline scripts from subdirectories of current working directory."""
+    cwd = Path.cwd()
+    # Find all .py files in subdirectories (not in cwd itself)
+    pipelines = []
+    for subdir in cwd.iterdir():
+        if subdir.is_dir():
+            pipelines.extend(subdir.rglob("*.py"))
+    return pipelines
+
+
+def get_pipeline_groups(pipelines: list[Path]) -> dict[Path, list[Path]]:
+    """Group pipelines by their parent directory (leaf directories)."""
+    groups = {}
+    for pipeline in pipelines:
+        parent = pipeline.parent
+        groups.setdefault(parent, []).append(pipeline)
+    return groups
+
+
+def select_random_groups(pipelines: list[Path], num_groups: int) -> list[Path]:
+    """Select pipelines from n random leaf directories."""
+    groups = get_pipeline_groups(pipelines)
+    if not groups:
+        return []
+
+    # Select up to num_groups random directories
+    dirs = list(groups.keys())
+    selected_dirs = random.sample(dirs, min(num_groups, len(dirs)))
+
+    # Return all pipelines from those directories
+    selected = []
+    for d in selected_dirs:
+        selected.extend(groups[d])
+    return selected
+
+
+def filter_pipelines_by_patterns(pipelines: list[Path], patterns: list[str]) -> list[Path]:
+    """Filter pipelines by substring patterns matching the basename."""
+    if not patterns:
+        return pipelines
+
+    matched = []
+    for pipeline in pipelines:
+        basename = pipeline.stem.lower()
+        if any(pattern.lower() in basename for pattern in patterns):
+            matched.append(pipeline)
+    return matched
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Launch ML pipelines via SQS for testing")
+    parser.add_argument(
+        "patterns",
+        nargs="*",
+        help="Substring patterns to filter pipelines by basename (e.g., 'caco2' 'ppb')",
+    )
+    parser.add_argument(
+        "-n",
+        "--num-groups",
+        type=int,
+        default=1,
+        help="Number of random pipeline groups to launch (default: 1, ignored if --all or patterns specified)",
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Launch ALL pipelines (ignores -n)",
+    )
+    parser.add_argument(
+        "--realtime",
+        action="store_true",
+        help="Create realtime endpoints (default is serverless)",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Show what would be launched without actually launching",
+    )
+
+    # Mode flags (mutually exclusive)
+    mode_group = parser.add_mutually_exclusive_group(required=True)
+    mode_group.add_argument(
+        "--dt",
+        action="store_true",
+        help="Launch with DT=True (dynamic training mode)",
+    )
+    mode_group.add_argument(
+        "--promote",
+        action="store_true",
+        help="Launch with PROMOTE=True (promotion mode)",
+    )
+    mode_group.add_argument(
+        "--test-promote",
+        action="store_true",
+        help="Launch with TEST_PROMOTE=True (test promotion mode)",
+    )
+
+    args = parser.parse_args()
+
+    # Get all pipelines from subdirectories of current working directory
+    all_pipelines = get_all_pipelines()
+    if not all_pipelines:
+        print(f"No pipeline scripts found in subdirectories of {Path.cwd()}")
+        exit(1)
+
+    # Determine which pipelines to run
+    if args.patterns:
+        # Filter by patterns
+        selected_pipelines = filter_pipelines_by_patterns(all_pipelines, args.patterns)
+        if not selected_pipelines:
+            print(f"No pipelines matching patterns: {args.patterns}")
+            exit(1)
+        selection_mode = f"matching {args.patterns}"
+    elif args.all:
+        # Run all pipelines
+        selected_pipelines = all_pipelines
+        selection_mode = "ALL"
+    else:
+        # Random group selection
+        selected_pipelines = select_random_groups(all_pipelines, args.num_groups)
+        if not selected_pipelines:
+            print("No pipeline groups found")
+            exit(1)
+        # Get the directory names for display
+        groups = get_pipeline_groups(selected_pipelines)
+        group_names = [d.name for d in groups.keys()]
+        selection_mode = f"RANDOM {args.num_groups} group(s): {group_names}"
+
+    # Sort by dependencies (producers before consumers)
+    selected_pipelines, configs, root_map = sort_by_dependencies(selected_pipelines)
+
+    # Determine mode for display and CLI flag
+    if args.dt:
+        mode_name = "DT (Dynamic Training)"
+        mode_flag = "--dt"
+    elif args.promote:
+        mode_name = "PROMOTE"
+        mode_flag = "--promote"
+    else:
+        mode_name = "TEST_PROMOTE"
+        mode_flag = "--test-promote"
+
+    print(f"\n{'=' * 60}")
+    print(f"{'DRY RUN - ' if args.dry_run else ''}LAUNCHING {len(selected_pipelines)} PIPELINES")
+    print(f"{'=' * 60}")
+    print(f"Source: {Path.cwd()}")
+    print(f"Selection: {selection_mode}")
+    print(f"Mode: {mode_name}")
+    print(f"Endpoint: {'Realtime' if args.realtime else 'Serverless'}")
+    print("\nPipeline Chains:")
+    for line in format_dependency_chains(selected_pipelines, configs):
+        print(line)
+    print()
+
+    # Dry run - just show what would be launched
+    if args.dry_run:
+        print("Dry run complete. No pipelines were launched.\n")
+        return
+
+    # Countdown before launching
+    print("Launching in ", end="", flush=True)
+    for i in range(10, 0, -1):
+        print(f"{i}...", end="", flush=True)
+        time.sleep(1)
+    print(" GO!\n")
+
+    # Launch each pipeline using the CLI
+    for i, pipeline in enumerate(selected_pipelines, 1):
+        print(f"\n{'─' * 60}")
+        print(f"Launching pipeline {i}/{len(selected_pipelines)}: {pipeline.name}")
+        print(f"{'─' * 60}")
+
+        # Build the command
+        cmd = ["ml_pipeline_sqs", str(pipeline), mode_flag]
+        if args.realtime:
+            cmd.append("--realtime")
+
+        # Pass root group_id for dependency chain ordering
+        group_id = get_group_id(configs.get(pipeline), root_map)
+        if group_id:
+            cmd.extend(["--group-id", group_id])
+
+        print(f"Running: {' '.join(cmd)}\n")
+        result = subprocess.run(cmd)
+        if result.returncode != 0:
+            print(f"Failed to launch {pipeline.name} (exit code: {result.returncode})")
+
+    print(f"\n{'=' * 60}")
+    print(f"FINISHED LAUNCHING {len(selected_pipelines)} PIPELINES")
+    print(f"{'=' * 60}\n")
+
+
+if __name__ == "__main__":
+    main()

workbench/scripts/ml_pipeline_sqs.py

@@ -71,6 +71,8 @@ def submit_to_sqs(
     realtime: bool = False,
     dt: bool = False,
     promote: bool = False,
+    test_promote: bool = False,
+    group_id: str | None = None,
 ) -> None:
     """
     Upload script to S3 and submit message to SQS queue for processing.
@@ -81,6 +83,8 @@ def submit_to_sqs(
         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)
+        test_promote: If True, sets TEST_PROMOTE=True in environment (default: False)
+        group_id: Optional MessageGroupId override for dependency chains (default: derived from script)
 
     Raises:
         ValueError: If size is invalid or script file not found
@@ -99,7 +103,8 @@ def submit_to_sqs(
     # 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)
+    if group_id is None:
+        group_id = get_message_group_id(batch_config)
     outputs = (batch_config or {}).get("outputs", [])
     inputs = (batch_config or {}).get("inputs", [])
 
@@ -108,6 +113,7 @@ def submit_to_sqs(
     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"🧪  Test Promote: {test_promote}")
     print(f"🪣  Bucket: {workbench_bucket}")
     if outputs:
         print(f"📤  Outputs: {outputs}")
@@ -174,6 +180,7 @@ def submit_to_sqs(
         "SERVERLESS": "False" if realtime else "True",
         "DT": str(dt),
         "PROMOTE": str(promote),
+        "TEST_PROMOTE": str(test_promote),
     }
 
     # Send the message to SQS
@@ -200,6 +207,7 @@ def submit_to_sqs(
     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"🧪  Test Promote: {test_promote}")
     if outputs:
         print(f"📤  Outputs: {outputs}")
     if inputs:
@@ -234,7 +242,17 @@ def main():
     parser.add_argument(
         "--promote",
         action="store_true",
-        help="Set Promote=True (models and endpoints will use promoted naming",
+        help="Set Promote=True (models and endpoints will use promoted naming)",
+    )
+    parser.add_argument(
+        "--test-promote",
+        action="store_true",
+        help="Set TEST_PROMOTE=True (creates test endpoint with '-test' suffix)",
+    )
+    parser.add_argument(
+        "--group-id",
+        default=None,
+        help="Override MessageGroupId for SQS (used for dependency chain ordering)",
     )
     args = parser.parse_args()
     try:
@@ -244,6 +262,8 @@ def main():
             realtime=args.realtime,
             dt=args.dt,
             promote=args.promote,
+            test_promote=args.test_promote,
+            group_id=args.group_id,
         )
     except Exception as e:
         print(f"\n❌  ERROR: {e}")

workbench/themes/dark/custom.css

@@ -135,6 +135,35 @@ div:has(> [class*="ag-theme-"]) {
     --bs-border-color: rgb(60, 60, 60);
 }
 
+/* React-select dropdown styling - target actual rendered elements */
+.Select-control {
+    background-color: rgb(35, 35, 35) !important;
+    border-color: rgb(60, 60, 60) !important;
+}
+
+.Select-value-label, .Select-input input {
+    color: rgb(210, 210, 210) !important;
+}
+
+.Select-placeholder {
+    color: rgb(150, 150, 150) !important;
+}
+
+.Select-menu-outer {
+    background-color: rgb(35, 35, 35) !important;
+    border-color: rgb(60, 60, 60) !important;
+}
+
+.VirtualizedSelectOption {
+    background-color: rgb(35, 35, 35) !important;
+    color: rgb(210, 210, 210) !important;
+}
+
+.VirtualizedSelectFocusedOption {
+    background-color: rgb(60, 60, 60) !important;
+    color: rgb(230, 230, 230) !important;
+}
+
 /* Bootstrap form controls (dbc components) */
 .form-select, .form-control {
     background-color: rgb(35, 35, 35) !important;

workbench/themes/light/custom.css

@@ -180,6 +180,35 @@ div:has(> [class*="ag-theme-"]) {
     --bs-border-color: var(--wb-accent);
 }
 
+/* React-select dropdown styling - target actual rendered elements */
+.Select-control {
+    background-color: var(--wb-dropdown-bg) !important;
+    border-color: var(--wb-accent) !important;
+}
+
+.Select-value-label, .Select-input input {
+    color: var(--wb-text-primary) !important;
+}
+
+.Select-placeholder {
+    color: var(--wb-text-muted) !important;
+}
+
+.Select-menu-outer {
+    background-color: var(--wb-dropdown-bg) !important;
+    border-color: var(--wb-accent) !important;
+}
+
+.VirtualizedSelectOption {
+    background-color: var(--wb-dropdown-bg) !important;
+    color: var(--wb-text-primary) !important;
+}
+
+.VirtualizedSelectFocusedOption {
+    background-color: var(--wb-dropdown-hover) !important;
+    color: var(--wb-text-primary) !important;
+}
+
 /* Bootstrap form controls (dbc components) */
 .form-select, .form-control {
     background-color: var(--wb-dropdown-bg) !important;

workbench/themes/midnight_blue/custom.css

@@ -133,6 +133,34 @@ div:has(> [class*="ag-theme-"]) {
     --bs-border-color: rgb(80, 85, 115);
 }
 
+/* React-select dropdown styling - target actual rendered elements */
+.Select-control {
+    background-color: rgb(55, 60, 90) !important;
+    border-color: rgb(80, 85, 115) !important;
+}
+
+.Select-value-label, .Select-input input {
+    color: rgb(210, 210, 210) !important;
+}
+
+.Select-placeholder {
+    color: rgb(150, 150, 170) !important;
+}
+
+.Select-menu-outer {
+    background-color: rgb(55, 60, 90) !important;
+    border-color: rgb(80, 85, 115) !important;
+}
+
+.VirtualizedSelectOption {
+    background-color: rgb(55, 60, 90) !important;
+    color: rgb(210, 210, 210) !important;
+}
+
+.VirtualizedSelectFocusedOption {
+    background-color: rgb(70, 75, 110) !important;
+    color: rgb(230, 230, 230) !important;
+}
 
 /* Bootstrap form controls (dbc components) */
 .form-select, .form-control {

workbench/utils/model_utils.py

@@ -459,6 +459,12 @@ def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
     # Spearman correlation for robustness
     interval_to_error_corr = spearmanr(width_68, abs_residuals)[0]
 
+    # --- Confidence to Error Correlation ---
+    # If confidence column exists, compute correlation (should be negative: high confidence = low error)
+    confidence_to_error_corr = None
+    if "confidence" in df.columns:
+        confidence_to_error_corr = spearmanr(df["confidence"], abs_residuals)[0]
+
     # Collect results
     results = {
         "coverage_68": coverage_68,
@@ -472,6 +478,7 @@ def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
         "median_width_90": median_width_90,
         "median_width_95": median_width_95,
         "interval_to_error_corr": interval_to_error_corr,
+        "confidence_to_error_corr": confidence_to_error_corr,
         "n_samples": len(df),
     }
 
@@ -489,6 +496,8 @@ def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
     print(f"CRPS: {mean_crps:.3f} (lower is better)")
     print(f"Interval Score 95%: {mean_is_95:.3f} (lower is better)")
     print(f"Interval/Error Corr: {interval_to_error_corr:.3f} (higher is better, target: >0.5)")
+    if confidence_to_error_corr is not None:
+        print(f"Confidence/Error Corr: {confidence_to_error_corr:.3f} (lower is better, target: <-0.5)")
     print(f"Samples: {len(df)}")
     return results