workbench 0.8.332__tar.gz → 0.8.333__tar.gz

{workbench-0.8.332 → workbench-0.8.333}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.332
+Version: 0.8.333
 Summary: Workbench: A Dashboard and Python API for creating and deploying AWS SageMaker Model Pipelines
 Author-email: SuperCowPowers LLC <support@supercowpowers.com>
 License: MIT License

workbench-0.8.333/examples/models/chemprop_bounded_loss_test.py

@@ -0,0 +1,175 @@
+"""Bounded-Loss Sanity Check (ChemProp)
+
+This script verifies that ChemProp's bounded-loss path (BoundedMAE / BoundedMSE)
+fires correctly when a FeatureSet provides per-target `_gt` / `_lt` censor
+columns and the model is trained with hyperparameters={"bounded_loss": True}.
+
+It builds a small synthetic FeatureSet from public LogP data with artificial
+right-censoring, then trains two ChemProp variants against the same FeatureSet:
+
+    1. logp-bounded-off  -- hyperparameters omitted (bounded_loss defaults to False)
+    2. logp-bounded-on   -- hyperparameters={"bounded_loss": True}
+
+Predictions on the *artificially censored* rows should look very different:
+
+    bounded_loss=False:  predictions cluster near the censor value T because the
+                         labels are exactly T and MAE punishes deviations in
+                         either direction.
+    bounded_loss=True:   predictions are free to go above T without penalty, and
+                         since the uncensored half of the high-LogP chemistry
+                         remains at its true (>T) value, the model can learn
+                         that high-LogP chemistry actually exists.
+
+Test design notes
+-----------------
+- Use a real LogP dataset so the model has actual chemistry signal to learn.
+- Pick T at a percentile that leaves a meaningful population above the cap.
+- Censor ONLY half of the rows above T -- the other half stays at its true
+  value so the bounded model has uncensored examples to learn from. Without
+  any uncensored examples above T, even bounded_loss=True can't pull
+  predictions above T (no information to push them there).
+
+Verifying the bounded-loss effect after training
+------------------------------------------------
+Note: workbench stores boolean FS columns as nullable Int64, so cast logp_gt to
+bool before boolean indexing (or compare with == 1):
+
+    fs = FeatureSet("bounded_loss_test_fs").pull_dataframe()
+    fs["logp_gt"] = fs["logp_gt"].fillna(0).astype(bool)
+    off = Model("logp-bounded-off").get_inference_predictions("full_cross_fold")
+    on  = Model("logp-bounded-on" ).get_inference_predictions("full_cross_fold")
+    off_m = off.merge(fs[["id", "logp_gt", "logp_true"]], on="id")
+    on_m  = on .merge(fs[["id", "logp_gt", "logp_true"]], on="id")
+    for name, m in [("off", off_m), ("on", on_m)]:
+        c = m[m["logp_gt"]]
+        err = (c["prediction"] - c["logp_true"]).abs()
+        print(f"{name}: MAE on censored vs true = {err.mean():.3f}")
+"""
+
+import numpy as np
+
+from workbench.api import Endpoint, FeatureSet, Model, ModelFramework, ModelType, PublicData
+from workbench.core.transforms.pandas_transforms.pandas_to_features import PandasToFeatures
+
+# Recreate Flag in case you want to recreate the artifacts
+recreate = False
+
+# Synthetic dataset tunables -- these define how visible the bounded-loss effect is
+FS_NAME = "bounded_loss_test_fs"
+PUBLIC_KEY = "comp_chem/logp/logp_all"
+N_ROWS = 3000  # subsample size (keeps training fast)
+CENSOR_THRESHOLD = 3.0  # T: cap above which rows can be censored (top ~25% of LogP)
+CENSOR_FRACTION = 0.5  # fraction of >T rows to actually censor
+RANDOM_SEED = 42
+
+
+# =============================================================================
+# Build Synthetic FeatureSet (LogP with artificial right-censoring)
+# =============================================================================
+# Pulls real LogP data from public storage, censors half of the rows above T
+# (recorded logp clipped to T, logp_gt flag set), and keeps the other half at
+# their true values. logp_true is preserved for post-training analysis.
+
+if recreate or not FeatureSet(FS_NAME).exists():
+    df = PublicData().get(PUBLIC_KEY)
+    if df is None:
+        raise RuntimeError(f"Could not load {PUBLIC_KEY} from public data")
+
+    df = df[["smiles", "logp"]].dropna().reset_index(drop=True)
+    df = df.sample(n=min(N_ROWS, len(df)), random_state=RANDOM_SEED).reset_index(drop=True)
+    df["id"] = df.index.astype(int)
+
+    # Preserve the true value so we can analyze post-training
+    df["logp_true"] = df["logp"]
+
+    # Censor half of the rows with logp > T: set recorded logp = T and mark _gt
+    rng = np.random.default_rng(RANDOM_SEED)
+    above = df.index[df["logp_true"] > CENSOR_THRESHOLD]
+    n_to_censor = int(len(above) * CENSOR_FRACTION)
+    censored_idx = rng.choice(above, size=n_to_censor, replace=False)
+
+    df["logp_gt"] = False
+    df.loc[censored_idx, "logp_gt"] = True
+    df.loc[censored_idx, "logp"] = CENSOR_THRESHOLD  # clip the label to the cap
+
+    df = df[["id", "smiles", "logp", "logp_gt", "logp_true"]]
+    n_above = int((df["logp_true"] > CENSOR_THRESHOLD).sum())
+    n_censored = int(df["logp_gt"].sum())
+    print(
+        f"FeatureSet '{FS_NAME}': {len(df)} rows | "
+        f"true logp > {CENSOR_THRESHOLD}: {n_above} | censored (logp_gt=True): {n_censored} | "
+        f"uncensored above cap: {n_above - n_censored}"
+    )
+
+    to_features = PandasToFeatures(FS_NAME)
+    to_features.set_input(df, id_column="id")
+    to_features.set_output_tags(["chemprop", "bounded-loss", "synthetic", "logp"])
+    to_features.transform()
+
+
+# =============================================================================
+# Control Model (bounded_loss=False -- default MAE)
+# =============================================================================
+# Trained without the bounded_loss hyperparameter so the censor labels are
+# treated as exact targets. Predictions on censored rows should cluster near
+# the cap value (T=3.0) because MAE punishes deviations in either direction.
+
+if recreate or not Model("logp-bounded-off").exists():
+    feature_set = FeatureSet(FS_NAME)
+    m = feature_set.to_model(
+        name="logp-bounded-off",
+        model_type=ModelType.UQ_REGRESSOR,
+        model_framework=ModelFramework.CHEMPROP,
+        target_column="logp",
+        feature_list=["smiles"],
+        description=(
+            "Bounded-loss sanity check (control): standard MAE on artificially censored "
+            f"LogP. Censored rows have logp = {CENSOR_THRESHOLD} and logp_gt = True; "
+            "trained without bounded_loss hyperparameter so censor labels are treated as exact."
+        ),
+        tags=["chemprop", "bounded-loss", "synthetic", "logp", "off"],
+    )
+    m.set_owner("BW")
+
+# Create an Endpoint
+if recreate or not Endpoint("logp-bounded-off").exists():
+    m = Model("logp-bounded-off")
+    end = m.to_endpoint(tags=["chemprop", "bounded-loss", "synthetic", "logp", "off"])
+    end.set_owner("BW")
+    end.auto_inference()
+    end.cross_fold_inference()
+
+
+# =============================================================================
+# Treatment Model (bounded_loss=True -- BoundedMAE)
+# =============================================================================
+# Trained with hyperparameters={"bounded_loss": True} so the logp_gt column is
+# read by the template and censor labels become lower bounds. Predictions on
+# censored rows should spread above the cap, using the chemistry signal from
+# the uncensored half of the high-LogP rows.
+
+if recreate or not Model("logp-bounded-on").exists():
+    feature_set = FeatureSet(FS_NAME)
+    m = feature_set.to_model(
+        name="logp-bounded-on",
+        model_type=ModelType.UQ_REGRESSOR,
+        model_framework=ModelFramework.CHEMPROP,
+        target_column="logp",
+        feature_list=["smiles"],
+        description=(
+            "Bounded-loss sanity check (treatment): BoundedMAE on artificially censored "
+            f"LogP. Censored rows have logp = {CENSOR_THRESHOLD} and logp_gt = True; "
+            "trained with bounded_loss=True so censor labels are treated as lower bounds."
+        ),
+        tags=["chemprop", "bounded-loss", "synthetic", "logp", "on"],
+        hyperparameters={"bounded_loss": True},
+    )
+    m.set_owner("BW")
+
+# Create an Endpoint
+if recreate or not Endpoint("logp-bounded-on").exists():
+    m = Model("logp-bounded-on")
+    end = m.to_endpoint(tags=["chemprop", "bounded-loss", "synthetic", "logp", "on"])
+    end.set_owner("BW")
+    end.auto_inference()
+    end.cross_fold_inference()

{workbench-0.8.332 → workbench-0.8.333}/sagemaker_images/ml_pipelines/Dockerfile

@@ -39,7 +39,7 @@ COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/pytho
 COPY --from=builder /usr/local/bin /usr/local/bin
 
 # Layer 2: Workbench + Bridges (changes often, small layer)
-ARG WORKBENCH_VERSION=0.8.331
+ARG WORKBENCH_VERSION=0.8.332
 RUN pip install --no-cache-dir --no-deps "workbench==${WORKBENCH_VERSION}" && \
     pip install --no-cache-dir --no-deps workbench-bridges
 

{workbench-0.8.332 → workbench-0.8.333}/src/workbench/algorithms/dataframe/smart_aggregator.py

@@ -10,7 +10,7 @@ import logging
 log = logging.getLogger("workbench")
 
 
-def smart_aggregator(df: pd.DataFrame, target_rows: int = 5000, outlier_column: str = "residual") -> pd.DataFrame:
+def smart_aggregator(df: pd.DataFrame, target_rows: int = 50000, outlier_column: str = "residual") -> pd.DataFrame:
     """
     Reduce DataFrame rows by aggregating similar rows based on numeric column similarity.
 
@@ -36,11 +36,9 @@ def smart_aggregator(df: pd.DataFrame, target_rows: int = 5000, outlier_column:
     # Preserve original column order
     original_columns = df.columns.tolist()
 
-    # If already at or below target, just add the count column and return
+    # If already at or below target, just return the dataframe
     if n_rows <= target_rows:
-        result = df.copy()
-        result["aggregation_count"] = 1
-        return result
+        return df
 
     log.info(f"smart_aggregator: Reducing {n_rows} rows to ~{target_rows} rows")
 
@@ -103,7 +101,7 @@ if __name__ == "__main__":
 
     # Create test data with clusters
     np.random.seed(42)
-    n_samples = 10000
+    n_samples = 100000
 
     # Create 3 distinct clusters
     cluster_1 = np.random.randn(n_samples // 3, 3) + np.array([0, 0, 0])

{workbench-0.8.332 → workbench-0.8.333}/src/workbench/core/artifacts/endpoint_core.py

@@ -609,12 +609,17 @@ class EndpointCore(Artifact):
         target_list = targets if isinstance(targets, list) else [targets]
         primary_target = target_list[0]
 
-        # If we don't have a smiles column, try to merge it from the FeatureSet
-        if "smiles" not in out_of_fold_df.columns and "smiles" in fs.columns:
-            fs_df = fs.query(f'SELECT {fs.id_column}, "smiles" FROM "{fs.athena_table}"')
-            if "smiles" in fs_df.columns:
-                self.log.info("Merging 'smiles' column from FeatureSet into out-of-fold predictions.")
-                out_of_fold_df = out_of_fold_df.merge(fs_df, on=fs.id_column, how="left")
+        # Templates don't save smiles or user-specified meta columns to
+        # validation_predictions.csv — merge them in from the FeatureSet so
+        # cross_fold captures match the live-inference column set.
+        meta_cols = self.workbench_meta().get("inference_meta", []) or []
+        merge_candidates = ["smiles"] + [c for c in meta_cols if c != "smiles"]
+        to_merge = [c for c in merge_candidates if c not in out_of_fold_df.columns and c in fs.columns]
+        if to_merge:
+            quoted = ", ".join(f'"{c}"' for c in to_merge)
+            fs_df = fs.query(f'SELECT {fs.id_column}, {quoted} FROM "{fs.athena_table}"')
+            self.log.info(f"Merging columns from FeatureSet into out-of-fold predictions: {to_merge}")
+            out_of_fold_df = out_of_fold_df.merge(fs_df, on=fs.id_column, how="left")
 
         # Collect UQ columns (q_*, confidence) for additional tracking (used for hashing)
         additional_columns = [col for col in out_of_fold_df.columns if col.startswith("q_") or col == "confidence"]
@@ -1100,6 +1105,12 @@ class EndpointCore(Artifact):
         if "smiles" in cols:
             output_columns.append("smiles")
 
+        # Add user-specified passthrough columns (workbench_meta["inference_meta"])
+        meta_cols = self.workbench_meta().get("inference_meta", []) or []
+        for c in meta_cols:
+            if c in cols and c not in output_columns:
+                output_columns.append(c)
+
         # Write the predictions to S3
         output_file = f"{inference_capture_path}/inference_predictions.csv"
         self.log.info(f"Writing predictions to {output_file}")

{workbench-0.8.332 → workbench-0.8.333}/src/workbench/model_scripts/chemprop/chemprop.template

@@ -15,6 +15,11 @@ import json
 import os
 import time
 
+# Startup diagnostic prints (flush=True so we don't lose them if a downstream
+# step hangs/crashes before its buffer flushes). The bare-import phase below
+# is the most common source of silent inference-container startup hangs.
+print("[startup] importing heavy deps...", flush=True)
+
 import joblib
 import numpy as np
 import pandas as pd
@@ -40,6 +45,8 @@ from uq_harness import (
     predict_intervals,
 )
 
+print("[startup] heavy deps imported", flush=True)
+
 # =============================================================================
 # Default Hyperparameters
 # =============================================================================
@@ -167,15 +174,19 @@ def model_fn(model_dir: str) -> dict:
     Optimized for serverless cold starts - uses direct PyTorch inference
     instead of Lightning Trainer to minimize startup time.
     """
+    print(f"[model_fn] loading from {model_dir}", flush=True)
     metadata = joblib.load(os.path.join(model_dir, "ensemble_metadata.joblib"))
+    print(f"[model_fn] metadata loaded, n_ensemble={metadata['n_ensemble']}", flush=True)
 
     # Load all ensemble models (keep on CPU for serverless compatibility)
     # ChemProp handles device placement internally
     ensemble_models = []
     for i in range(metadata["n_ensemble"]):
+        print(f"[model_fn] loading chemprop_model_{i}.pt...", flush=True)
         model = models.MPNN.load_from_file(os.path.join(model_dir, f"chemprop_model_{i}.pt"))
         model.eval()
         ensemble_models.append(model)
+    print(f"[model_fn] all {len(ensemble_models)} models loaded", flush=True)
 
     # Load UQ metadata (regression or classification confidence)
     uq_metadata = None
@@ -630,10 +641,10 @@ if __name__ == "__main__":
             gt_col, lt_col = f"{t_name}_gt", f"{t_name}_lt"
             target_meta: dict = {}
             if gt_col in all_df.columns:
-                gt_arr[:, t_idx] = all_df[gt_col].fillna(False).astype(bool).values
+                gt_arr[:, t_idx] = all_df[gt_col].fillna(0).astype(bool).values
                 target_meta["gt"] = gt_col
             if lt_col in all_df.columns:
-                lt_arr[:, t_idx] = all_df[lt_col].fillna(False).astype(bool).values
+                lt_arr[:, t_idx] = all_df[lt_col].fillna(0).astype(bool).values
                 target_meta["lt"] = lt_col
             if target_meta:
                 resolved_bound_columns[t_name] = target_meta

{workbench-0.8.332 → workbench-0.8.333}/src/workbench/themes/dark/plotly.json

@@ -509,11 +509,30 @@
             ]
         },
         "colorway": [
-            "rgba(100, 100, 200, 0.8)",
-            "rgb(100, 200, 100)",
-            "rgb(170, 170, 80)",
-            "rgb(180, 140, 100)",
-            "rgb(200, 100, 100)"
+            "#FD3216",
+            "#00FE35",
+            "#6A76FC",
+            "#FED4C4",
+            "#FE00CE",
+            "#0DF9FF",
+            "#F6F926",
+            "#FF9616",
+            "#479B55",
+            "#EEA6FB",
+            "#DC587D",
+            "#D626FF",
+            "#6E899C",
+            "#00B5F7",
+            "#B68E00",
+            "#C9FBE5",
+            "#FF0092",
+            "#22FFA7",
+            "#E3EE9E",
+            "#86CE00",
+            "#BC7196",
+            "#7E7DCD",
+            "#FC6955",
+            "#E48F72"
         ],
         "font": {
             "color": "#fff",
@@ -643,11 +662,30 @@
             "gridwidth": 0.5
         },
         "piecolorway": [
-            "rgba(100, 100, 200, 0.8)",
-            "rgb(100, 200, 100)",
-            "rgb(170, 170, 80)",
-            "rgb(180, 140, 100)",
-            "rgb(200, 100, 100)"
+            "#FD3216",
+            "#00FE35",
+            "#6A76FC",
+            "#FED4C4",
+            "#FE00CE",
+            "#0DF9FF",
+            "#F6F926",
+            "#FF9616",
+            "#479B55",
+            "#EEA6FB",
+            "#DC587D",
+            "#D626FF",
+            "#6E899C",
+            "#00B5F7",
+            "#B68E00",
+            "#C9FBE5",
+            "#FF0092",
+            "#22FFA7",
+            "#E3EE9E",
+            "#86CE00",
+            "#BC7196",
+            "#7E7DCD",
+            "#FC6955",
+            "#E48F72"
         ]
     }
 }

{workbench-0.8.332 → workbench-0.8.333}/src/workbench/themes/light/plotly.json

@@ -544,11 +544,30 @@
             ]
         },
         "colorway": [
-            "rgba(100, 100, 200, 0.8)",
-            "rgb(100, 200, 100)",
-            "rgb(170, 170, 80)",
-            "rgb(180, 140, 100)",
-            "rgb(200, 100, 100)"
+            "#2E91E5",
+            "#E15F99",
+            "#1CA71C",
+            "#FB0D0D",
+            "#DA16FF",
+            "#222A2A",
+            "#B68100",
+            "#750D86",
+            "#EB663B",
+            "#511CFB",
+            "#00A08B",
+            "#FB00D1",
+            "#FC0080",
+            "#B2828D",
+            "#6C7C32",
+            "#778AAE",
+            "#862A16",
+            "#A777F1",
+            "#620042",
+            "#1616A7",
+            "#DA60CA",
+            "#6C4516",
+            "#0D2A63",
+            "#AF0038"
         ],
         "font": {
             "color": "#212529",
@@ -668,11 +687,30 @@
             "gridwidth": 0.5
         },
         "piecolorway": [
-            "rgba(100, 100, 200, 0.8)",
-            "rgb(100, 200, 100)",
-            "rgb(170, 170, 80)",
-            "rgb(180, 140, 100)",
-            "rgb(200, 100, 100)"
+            "#2E91E5",
+            "#E15F99",
+            "#1CA71C",
+            "#FB0D0D",
+            "#DA16FF",
+            "#222A2A",
+            "#B68100",
+            "#750D86",
+            "#EB663B",
+            "#511CFB",
+            "#00A08B",
+            "#FB00D1",
+            "#FC0080",
+            "#B2828D",
+            "#6C7C32",
+            "#778AAE",
+            "#862A16",
+            "#A777F1",
+            "#620042",
+            "#1616A7",
+            "#DA60CA",
+            "#6C4516",
+            "#0D2A63",
+            "#AF0038"
         ]
     }
 }

{workbench-0.8.332 → workbench-0.8.333}/src/workbench/themes/midnight_blue/plotly.json

@@ -509,11 +509,30 @@
             ]
         },
         "colorway": [
-            "rgba(100, 100, 200, 0.8)",
-            "rgb(100, 200, 100)",
-            "rgb(170, 170, 80)",
-            "rgb(180, 140, 100)",
-            "rgb(200, 100, 100)"
+            "#FD3216",
+            "#00FE35",
+            "#6A76FC",
+            "#FED4C4",
+            "#FE00CE",
+            "#0DF9FF",
+            "#F6F926",
+            "#FF9616",
+            "#479B55",
+            "#EEA6FB",
+            "#DC587D",
+            "#D626FF",
+            "#6E899C",
+            "#00B5F7",
+            "#B68E00",
+            "#C9FBE5",
+            "#FF0092",
+            "#22FFA7",
+            "#E3EE9E",
+            "#86CE00",
+            "#BC7196",
+            "#7E7DCD",
+            "#FC6955",
+            "#E48F72"
         ],
         "font": {
             "color": "#dee2e6",
@@ -643,11 +662,30 @@
             "gridwidth": 0.5
         },
         "piecolorway": [
-            "rgba(100, 100, 200, 0.8)",
-            "rgb(100, 200, 100)",
-            "rgb(170, 170, 80)",
-            "rgb(180, 140, 100)",
-            "rgb(200, 100, 100)"
+            "#FD3216",
+            "#00FE35",
+            "#6A76FC",
+            "#FED4C4",
+            "#FE00CE",
+            "#0DF9FF",
+            "#F6F926",
+            "#FF9616",
+            "#479B55",
+            "#EEA6FB",
+            "#DC587D",
+            "#D626FF",
+            "#6E899C",
+            "#00B5F7",
+            "#B68E00",
+            "#C9FBE5",
+            "#FF0092",
+            "#22FFA7",
+            "#E3EE9E",
+            "#86CE00",
+            "#BC7196",
+            "#7E7DCD",
+            "#FC6955",
+            "#E48F72"
         ]
     }
 }

{workbench-0.8.332 → workbench-0.8.333}/src/workbench/web_interface/components/plugins/ag_table.py

@@ -21,7 +21,8 @@ class AGTable(PluginInterface):
     max_height = 500
     header_height = 36
     row_height = 25
-    thin_columns = {"Health", "Owner"}
+    # Thin columns: small by default (auto-size cap) but user can drag wider
+    thin_columns = {"Health": 100, "Owner": 100}
 
     def create_component(self, component_id: str, max_height: int = 500, row_selection: str = "single") -> AgGrid:
         """Create a Table Component without any data.
@@ -38,12 +39,17 @@ class AGTable(PluginInterface):
         self.max_height = max_height
 
         # AG Grid configuration for tighter rows and columns
+        column_limits = [{"colId": col, "maxWidth": w} for col, w in self.thin_columns.items()]
         grid_options = {
             "rowSelection": row_selection,
             "suppressCellFocus": True,
             "rowHeight": self.row_height,
             "defaultColDef": {"sortable": True, "filter": True, "resizable": True},
-            "autoSizeStrategy": {"type": "fitCellContents", "defaultMaxWidth": 200},
+            "autoSizeStrategy": {
+                "type": "fitCellContents",
+                "defaultMaxWidth": 400,
+                "columnLimits": column_limits,
+            },
         }
 
         self.container = AgGrid(
@@ -86,12 +92,7 @@ class AGTable(PluginInterface):
         table_data = table_df.to_dict("records")
 
         # Create column definitions for AG Grid
-        column_defs = []
-        for col in table_df.columns:
-            col_def = {"headerName": col, "field": col}
-            if col in self.thin_columns:
-                col_def["maxWidth"] = 60
-            column_defs.append(col_def)
+        column_defs = [{"headerName": col, "field": col} for col in table_df.columns]
 
         # Dynamically adjust height based on row count
         row_count = len(table_df)