workbench 0.8.414__tar.gz → 0.8.415__tar.gz

{workbench-0.8.414 → workbench-0.8.415}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.414
+Version: 0.8.415
 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.414 → workbench-0.8.415}/applications/aws_dashboard/Dockerfile

@@ -29,7 +29,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         -r /tmp/requirements.lock
 
 # Layer 2: Workbench only (changes often, ~20 MB)
-ARG WORKBENCH_VERSION=0.8.413
+ARG WORKBENCH_VERSION=0.8.414
 RUN pip install --no-deps "workbench==${WORKBENCH_VERSION}"
 
 # Copy app code and configuration (after pip installs to avoid cache invalidation)

{workbench-0.8.414 → workbench-0.8.415}/examples/models/chemprop_hpo.py

@@ -29,13 +29,13 @@ if recreate or not Model(model_name).exists():
             "uq_version": "v1",
             "hpo": {
                 "backend": "optuna",  # serial; flip to "ray" for parallel trials + ASHA on the 4 GPUs
-                "n_trials": 10,  # small for a first run; raise for a real search
+                "n_trials": 40,  # 5 baseline trials (pruner warmup) + 35 pruned candidates
                 "search_space": "basic",  # "basic" (architecture + dropout) | "basic+lr"
             },
         },
         # For an honest out-of-distribution objective, pass validation_ids=[...]: those
         # rows are held out of training and scored as `holdout_mae`. Without them the
-        # search optimizes `cv_mae` on a random split.
+        # search optimizes `cv_mae` on a single scaffold split of the training rows.
     )
     m.set_owner("BW")
 

{workbench-0.8.414 → workbench-0.8.415}/sagemaker_images/base/inference/Dockerfile

@@ -36,7 +36,7 @@ COPY --from=builder /usr/local/bin /usr/local/bin
 # workbench-bridges==0.2.10 retained for backward compat with existing
 # deployed-endpoint model bundles; drop in a future image rev once no
 # production endpoint references workbench_bridges.
-ARG WORKBENCH_VERSION=0.8.413
+ARG WORKBENCH_VERSION=0.8.414
 RUN pip install --no-deps "workbench==${WORKBENCH_VERSION}" && \
     pip install --no-deps "workbench-bridges==0.2.10"
 

{workbench-0.8.414 → workbench-0.8.415}/sagemaker_images/base/training/Dockerfile

@@ -35,6 +35,6 @@ COPY --from=builder /usr/local/bin /usr/local/bin
 # workbench-bridges==0.2.10 is retained for backward compatibility with
 # existing deployed-endpoint model bundles; drop in a future image rev once
 # no production endpoint references workbench_bridges.
-ARG WORKBENCH_VERSION=0.8.413
+ARG WORKBENCH_VERSION=0.8.414
 RUN pip install --no-deps "workbench==${WORKBENCH_VERSION}" && \
     pip install --no-deps "workbench-bridges==0.2.10"

{workbench-0.8.414 → workbench-0.8.415}/sagemaker_images/ml_pipelines/Dockerfile

@@ -48,7 +48,7 @@ COPY --from=builder /usr/local/bin /usr/local/bin
 # Layer 2: Workbench + Bridges (changes often, small layer)
 # workbench-bridges==0.2.10 pinned for reproducibility while we still depend
 # on it; on track for deprecation once nothing references workbench_bridges.
-ARG WORKBENCH_VERSION=0.8.413
+ARG WORKBENCH_VERSION=0.8.414
 RUN pip install --no-deps "workbench==${WORKBENCH_VERSION}" && \
     pip install --no-deps "workbench-bridges==0.2.10"
 

{workbench-0.8.414 → workbench-0.8.415}/sagemaker_images/pytorch_chem/inference/Dockerfile

@@ -40,7 +40,7 @@ COPY --from=builder /usr/local/bin /usr/local/bin
 # Layer 2: Workbench + bridges — rebuilds per version bump, ~20 MB delta.
 # workbench-bridges==0.2.10 retained for backward compat with existing
 # deployed-endpoint model bundles.
-ARG WORKBENCH_VERSION=0.8.413
+ARG WORKBENCH_VERSION=0.8.414
 RUN pip install --no-deps "workbench==${WORKBENCH_VERSION}" && \
     pip install --no-deps "workbench-bridges==0.2.10"
 

{workbench-0.8.414 → workbench-0.8.415}/sagemaker_images/pytorch_chem/training/Dockerfile

@@ -34,6 +34,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Layer 2: Workbench + bridges — rebuilds per version bump, ~20 MB delta.
 # workbench-bridges==0.2.10 retained for backward compat with existing
 # deployed-endpoint model bundles.
-ARG WORKBENCH_VERSION=0.8.413
+ARG WORKBENCH_VERSION=0.8.414
 RUN pip install --no-deps "workbench==${WORKBENCH_VERSION}" && \
     pip install --no-deps "workbench-bridges==0.2.10"

{workbench-0.8.414 → workbench-0.8.415}/src/workbench/training/chemprop_hpo.py

@@ -21,18 +21,21 @@ import os
 
 from workbench.training.hpo_harness import Choice, FloatRange, IntRange
 
-# Default per-knob search space, grouped like chemprop's own hpopt keywords. The
-# `basic` group matches chemprop's documented ranges, adapted to our ensemble-tuned
-# defaults: dropout is narrowed (the 5-fold ensemble already regularizes), and
-# ffn_hidden_dim is a categorical that encodes both uniform widths and tapered heads
-# (folding in ffn_num_layers, and exercising the list-valued head the harness
-# supports). uq_version and other knobs stay fixed at their configured values.
+# Default per-knob search space, grouped like chemprop's own hpopt keywords. The `basic`
+# group is chemprop's canonical space verbatim, extended only so ffn_hidden_dim can also
+# express a tapered head. Everything else — uq_version, max_epochs, patience, batch_size,
+# split_strategy, criterion, seed — stays fixed at its configured value.
 _SEARCH_GROUPS = {
     "basic": {
+        # depth / hidden_dim / dropout / ffn_num_layers ranges are chemprop's canonical
+        # search space verbatim. ffn_hidden_dim additionally offers tapered heads (lists),
+        # which chemprop's own space can't express — ffn_num_layers is ignored when a
+        # tapered list is chosen (the list length sets the depth).
         "depth": IntRange(2, 6, 1),
         "hidden_dim": IntRange(300, 2400, 100),
-        "dropout": FloatRange(0.0, 0.3, 0.05),
-        "ffn_hidden_dim": Choice([2000, 1000, 500, [1024, 256, 64], [512, 128]]),
+        "dropout": FloatRange(0.0, 0.4, 0.05),
+        "ffn_num_layers": IntRange(1, 3, 1),
+        "ffn_hidden_dim": Choice([300, 600, 1200, 1800, 2400, [1024, 256, 64], [512, 128]]),
     },
     # Opt-in ("basic+lr"). init_lr/final_lr are tied to max_lr in merge_best_config
     # rather than searched independently (independent search can produce init > max,
@@ -43,6 +46,23 @@ _SEARCH_GROUPS = {
     },
 }
 
+# Knobs outside the default space — the working list for when this is revisited:
+#
+#   * learning rate (`max_lr`, `warmup_epochs`) — already available as the opt-in "lr"
+#     group. Chemprop's maintainers describe their recommended search as focusing on
+#     learning rate and batch size, which makes this the first candidate to promote into
+#     the default.
+#   * `batch_size` — searchable in chemprop, not here. It moves memory and throughput as
+#     well as accuracy, so a range has to be chosen against the training instance's GPU
+#     memory rather than picked from the literature.
+#   * `aggregation` (mean/sum/norm) — searchable in chemprop but unreachable from here:
+#     `chemprop_core.build_mpnn_model` constructs `NormAggregation` directly, so tuning it
+#     requires a model-construction change first.
+#   * `activation`, `aggregation_norm` — the remainder of chemprop's "all" keyword.
+#
+# Each added knob costs trials: the default space is already 5-dimensional, and pruning
+# reserves the first PRUNE_STARTUP_TRIALS trials as un-pruned baselines.
+
 
 def chemprop_search_space(groups=("basic",)) -> dict:
     """Build the default chemprop search space for the named knob ``groups``.
@@ -219,6 +239,13 @@ def _make_trial_fn(train_df, eval_df, base_hyperparameters, target_columns, smil
         target_scaler = train_ds.normalize_targets()
         eval_ds.normalize_targets(target_scaler)
         output_transform = nn.UnscaleTransform.from_standard_scaler(target_scaler)
+        # `val_loss` is the criterion on *normalized* targets, but the value we return is
+        # unscaled MAE. Rescale per-epoch reports into the objective's units so the pruning
+        # signal, the recorded per-trial values, and the final value are all one quantity —
+        # otherwise pruned trials record a smaller number than completed ones and look
+        # better than they are. Exact for the default MAE criterion (MAE is linear in the
+        # target scale); a monotone proxy for others.
+        target_scale = float(getattr(target_scaler, "scale_", [1.0])[0])
         train_ds.cache = eval_ds.cache = True
 
         batch_size = hp.get("batch_size", 64)
@@ -230,9 +257,13 @@ def _make_trial_fn(train_df, eval_df, base_hyperparameters, target_columns, smil
         class _ReportPruning(pl.Callback):
             # Report per-epoch held-out val_loss so the harness can prune weak trials (ASHA).
             def on_validation_epoch_end(self, trainer, module):
+                # Lightning's pre-training sanity-check pass fires this at epoch 0 too;
+                # reporting it would duplicate step 0 with an untrained model's value.
+                if trainer.sanity_checking:
+                    return
                 v = trainer.callback_metrics.get("val_loss")
                 if v is not None:
-                    report(step=trainer.current_epoch, **{metric: float(v)})
+                    report(step=trainer.current_epoch, **{metric: float(v) * target_scale})
 
         trainer = pl.Trainer(
             accelerator="auto",

{workbench-0.8.414 → workbench-0.8.415}/src/workbench/training/hpo_harness.py

@@ -66,6 +66,14 @@ class Choice:
 Spec = Union[IntRange, FloatRange, Choice]
 SearchSpace = dict
 
+# Pruning grace. Successive-halving/ASHA defaults judge a trial almost immediately,
+# which wrecks a small search (tens of trials): configs get killed before they've
+# trained enough to rank. No trial is eligible for pruning until it has reported this
+# many steps (epochs), and Optuna additionally needs this many completed trials as
+# baselines before it prunes anything.
+PRUNE_WARMUP_STEPS = 20
+PRUNE_STARTUP_TRIALS = 5
+
 
 @dataclass
 class HpoResult:
@@ -165,7 +173,14 @@ def _run_optuna(trial_fn, search_space, *, n_trials, max_parallel, metric, mode,
 
     optuna.logging.set_verbosity(optuna.logging.WARNING)
     sampler = optuna.samplers.TPESampler(seed=seed)
-    pruner = optuna.pruners.SuccessiveHalvingPruner() if pruning else optuna.pruners.NopPruner()
+    # MedianPruner rather than successive halving: at our trial budgets, halving prunes
+    # before configs are rankable. Median prunes a trial only once baselines exist and it
+    # is worse than the median at the same step.
+    pruner = (
+        optuna.pruners.MedianPruner(n_startup_trials=PRUNE_STARTUP_TRIALS, n_warmup_steps=PRUNE_WARMUP_STEPS)
+        if pruning
+        else optuna.pruners.NopPruner()
+    )
     study = optuna.create_study(direction="minimize" if mode == "min" else "maximize", sampler=sampler, pruner=pruner)
 
     def objective(trial):
@@ -256,7 +271,13 @@ def _run_ray(
         tune.report({metric: value, time_attr: last_step + 1})
 
     trainable_res = tune.with_resources(trainable, resources_per_trial) if resources_per_trial else trainable
-    scheduler = ASHAScheduler(metric=metric, mode=mode, time_attr=time_attr) if pruning else None
+    # grace_period defaults to 1 (prune after a single report) — far too eager; give each
+    # trial the same warmup the Optuna path gets.
+    scheduler = (
+        ASHAScheduler(metric=metric, mode=mode, time_attr=time_attr, grace_period=PRUNE_WARMUP_STEPS)
+        if pruning
+        else None
+    )
     tuner = tune.Tuner(
         trainable_res,
         param_space=param_space,

{workbench-0.8.414 → workbench-0.8.415}/src/workbench.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.414
+Version: 0.8.415
 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.414 → workbench-0.8.415}/src/workbench.egg-info/scm_version.json

@@ -1,7 +1,7 @@
 {
-  "tag": "0.8.414",
+  "tag": "0.8.415",
   "distance": 0,
-  "node": "gcf5e6946b61d17ae042ef40394cb61066ab025c4",
+  "node": "gd26599192a9c2b2e08473a6852cca6e837bff4ae",
   "dirty": false,
   "branch": "HEAD",
   "node_date": "2026-07-23"

{workbench-0.8.414 → workbench-0.8.415}/tests/training/chemprop_hpo_test.py

@@ -10,11 +10,13 @@ from workbench.training.hpo_harness import Choice, FloatRange, IntRange
 
 
 def test_basic_group_is_default():
-    """The default space is the `basic` group with our chosen knobs/ranges."""
+    """The default space is the `basic` group, matching chemprop's canonical ranges."""
     space = chemprop_search_space()
-    assert set(space) == {"depth", "hidden_dim", "dropout", "ffn_hidden_dim"}
-    assert space["depth"] == IntRange(2, 6, 1)
-    assert space["dropout"] == FloatRange(0.0, 0.3, 0.05)
+    assert set(space) == {"depth", "hidden_dim", "dropout", "ffn_num_layers", "ffn_hidden_dim"}
+    assert space["depth"] == IntRange(2, 6, 1)  # chemprop {2,3,4,5,6}
+    assert space["hidden_dim"] == IntRange(300, 2400, 100)  # chemprop {300,400,...,2400}
+    assert space["dropout"] == FloatRange(0.0, 0.4, 0.05)  # chemprop {0.0,0.05,...,0.4}
+    assert space["ffn_num_layers"] == IntRange(1, 3, 1)  # chemprop {1,2,3}
     assert isinstance(space["ffn_hidden_dim"], Choice)
     assert [1024, 256, 64] in space["ffn_hidden_dim"].options  # tapered head is a choice
 
@@ -36,8 +38,10 @@ def test_unknown_group_raises():
 
 def test_resolve_search_space_shorthands():
     """String/iterable/dict/None all resolve to a {knob: Spec} space."""
-    assert set(resolve_search_space("basic")) == {"depth", "hidden_dim", "dropout", "ffn_hidden_dim"}
+    assert set(resolve_search_space("basic")) == set(chemprop_search_space(("basic",)))
     assert "max_lr" in resolve_search_space("basic+lr")
+    # basic+lr is a superset of basic
+    assert set(resolve_search_space("basic")) < set(resolve_search_space("basic+lr"))
     assert set(resolve_search_space(None)) == set(chemprop_search_space())
     custom = {"depth": IntRange(3, 5, 1)}
     assert resolve_search_space(custom) is custom  # ready dict passes through

{workbench-0.8.414 → workbench-0.8.415}/tests/training/hpo_harness_test.py

@@ -87,6 +87,16 @@ def test_invalid_backend_raises():
         run_search(_quadratic_objective, SPACE, n_trials=1, backend="optuaa")
 
 
+def test_default_pruning_respects_warmup():
+    """With pruning on, trials reporting fewer steps than the warmup are never pruned.
+
+    Guards the grace period: an over-eager pruner kills configs before they're rankable
+    (which collapsed a real 10-trial search down to 2 usable evaluations).
+    """
+    result = run_search(_quadratic_objective, SPACE, n_trials=12, backend="optuna")
+    assert all(t["state"] == "COMPLETE" for t in result.trials)
+
+
 def test_pruning_disabled_runs_all_trials():
     """With pruning off, every trial completes (none pruned)."""
     result = run_search(_quadratic_objective, SPACE, n_trials=15, backend="optuna", pruning=False)