Skip to content

Meta Endpoints

What's a MetaEndpoint?

A MetaEndpoint is a deployed Endpoint backed by a directed acyclic graph (DAG) of other endpoints + aggregation nodes. From the caller's perspective it's just an Endpoint — endpoint.inference(df) returns a DataFrame. The DAG machinery is server-side.

MetaEndpoint lets you compose multiple deployed endpoints into a single inference target. Two canonical shapes:

  • Feature pipelines — fan out to several feature endpoints in parallel, merge the columns, optionally feed a predictor.
  • Ensembles — fan out to several predictor endpoints, aggregate their predictions (mean, weighted mean, vote, calibrated confidence weighting, …) into a single output.

The same DAG abstraction covers both.

Every endpoint is a df-in → df-out unit, so the same DAG abstraction expresses both canonical shapes. A feature pipeline fans out to feature endpoints and concatenates their columns:

Feature MetaEndpoint: SMILES df fans out to smiles-to-2d-v1 (sync) and smiles-to-3d-v2 (async), a Concat node merges the columns, returning one df with 339 features. The whole thing is a single async endpoint.
Feature pipeline — fan out to the 2D + 3D endpoints, Concat the columns. One async endpoint, ~339 features.

An ensemble fans out to predictor endpoints and aggregates their predictions (mean, vote, weighted, calibrated confidence) into one output:

Ensemble MetaEndpoint: feature df fans out to three logd predictor endpoints (xgb, pytorch, chemprop), a Mean node aggregates them into a single prediction with std and confidence.
Ensemble — fan out to predictor endpoints, aggregate into one prediction. A MetaEndpoint is itself an endpoint, so DAGs nest arbitrarily.

One endpoint, ~339 features

smiles-to-2d-3d-v2 is a deployed MetaEndpoint that fans out to the 2D descriptor endpoint and the 3D descriptor endpoint in parallel and concatenates the results — ~313 RDKit/Mordred 2D features + 26 curated GFN2-xTB 3D features, merged per molecule. Callers just do endpoint.inference(df); the fan-out and merge happen server-side.

Our flagship feature MetaEndpoint combines the 2D and 3D descriptor endpoints into a single inference target. Because one of its children (smiles-to-3d-v2) is async, the whole MetaEndpoint is automatically deployed async (see Async Auto-Detection) — so a single call returns the complete 2D + 3D feature set without the caller juggling two endpoints or two invocation modes.

import pandas as pd
from workbench.api import MetaEndpoint

# Use the deployed endpoint like any other — fan-out + merge is server-side
end = MetaEndpoint("smiles-to-2d-3d-v2")
df = pd.DataFrame({"smiles": ["CCO", "c1ccccc1"]})
result = end.inference(df)
# result = input columns + ~313 2D feature columns + 26 3D feature columns

It is built from this DAG (see feature_endpoints/smiles_to_2d_3d_v2.py for the deploy script):

from workbench.api import MetaEndpoint
from workbench.utils.meta_endpoint_dag import MetaEndpointDAG
from workbench.utils.aggregation_nodes import Concat

dag = MetaEndpointDAG()
dag.add_endpoint("smiles-to-2d-v1")   # sync, RDKit + Mordred 2D
dag.add_endpoint("smiles-to-3d-v2")   # async, curated GFN2-xTB 3D
dag.add_aggregation(Concat(name="combine"))
dag.add_edge("smiles-to-2d-v1", "combine")
dag.add_edge("smiles-to-3d-v2", "combine")
dag.set_input_node("smiles-to-2d-v1", "smiles-to-3d-v2")
dag.set_output_node("combine")

end = MetaEndpoint.create(
    name="smiles-to-2d-3d-v2",
    dag=dag,
    description="SMILES → RDKit/Mordred 2D + Boltzmann 3D molecular descriptors",
    tags=["meta", "features"],
    min_instances=0,   # scale to zero when idle
    max_instances=1,
)

Quick Start: Ensemble

Combine three predictor endpoints (XGBoost + PyTorch + ChemProp) into one ensemble:

from workbench.api import MetaEndpoint
from workbench.utils.meta_endpoint_dag import MetaEndpointDAG
from workbench.utils.aggregation_nodes import Mean

dag = MetaEndpointDAG()
dag.add_endpoint("logd-reg-xgb")
dag.add_endpoint("logd-reg-pytorch")
dag.add_endpoint("logd-reg-chemprop")
dag.add_aggregation(Mean(name="ensemble"))
dag.add_edge("logd-reg-xgb", "ensemble")
dag.add_edge("logd-reg-pytorch", "ensemble")
dag.add_edge("logd-reg-chemprop", "ensemble")
dag.set_input_node("logd-reg-xgb", "logd-reg-pytorch", "logd-reg-chemprop")
dag.set_output_node("ensemble")

end = MetaEndpoint.create(name="logd-meta", dag=dag, tags=["meta", "logd"])

The output has the standard prediction / prediction_std (ensemble disagreement) / confidence columns alongside whatever pass-through columns the input had.

Async Auto-Detection

If any child endpoint in the DAG is deployed as async (e.g. smiles-to-3d-v2), the MetaEndpoint is automatically deployed as async too — its 60-minute invocation budget needs to accommodate the slowest child. You don't specify this; MetaEndpoint.create() detects it via dag.has_async_endpoint() and chooses the deploy mode. This is exactly why smiles-to-2d-3d-v2 above is async: its sync 2D child and async 3D child are composed transparently, and the container dispatches each to fast_inference (sync) or async_inference (async) as appropriate.

DAG Building Blocks

Endpoint nodes

Every endpoint node refers to a deployed Workbench endpoint by name. Endpoint nodes can be:

  • Input nodes — receive the caller's input DataFrame directly. Declared with dag.set_input_node(...).
  • Downstream endpoint nodes — take their input from a single upstream parent (e.g. a Concat aggregation feeding a predictor).
dag.add_endpoint("smiles-to-2d-v1")                       # node name = endpoint name (default)
dag.add_endpoint("smiles-to-2d-v1", node_name="left_2d")  # explicit node name (for aliasing)

Aggregation nodes

Class Use case Output
Concat Column-union of feature outputs from parallel branches All upstream columns merged
Mean Equal-weight average of predictions prediction, prediction_std, confidence
WeightedMean(weights=[…]) Static-weight average Same
Vote Majority vote for classifiers prediction (label), confidence (winner share)
ConfidenceWeighted(model_weights=…) Per-row weights from upstream confidences Same as Mean, with calibrated confidence
InverseMaeWeighted(model_weights=…) Static weights from inverse-MAE Same
ScaledConfidenceWeighted(model_weights=…) Static MAE × per-row confidence Same
CalibratedConfidenceWeighted(model_weights=…, corr_scale=…) Confidence × |conf-error correlation| Same

Edges

Edges declare data flow:

dag.add_edge("smiles-to-2d-v1", "combine")    # 2D output flows into combine
dag.add_edge("combine", "predictor")           # combined features flow into predictor

Endpoint nodes accept at most one inbound edge (one source for their input DataFrame). Aggregation nodes can have any number of inbound edges.

Validation

Call dag.validate() to fail loud on misconfiguration before any inference round-trips. Checks include cycle detection, dangling endpoint nodes, and reachability from input to output.

Row Alignment

The walker injects a synthetic __dag_row_id column at the start of every run() and strips it before returning. Aggregation nodes use it as the join key so callers do not need to supply any id column on their input data — and any id-like column they do supply just flows through as a regular pass-through column.

This gives MetaEndpoints a clean contract: any DataFrame with the columns your input-node endpoints expect is a valid input.

How It Works

Creation flow

When you call MetaEndpoint.create(name, dag, ...):

  1. Validate — dag.validate() fails loud on cycles, dangling nodes, etc.
  2. Resolve async flags — looks up each child's workbench_meta to record per-endpoint async status.
  3. Lineage anchor — backtraces the first input endpoint to a FeatureSet/target/feature_list (Workbench Models need to point at a FeatureSet).
  4. Build + register the Model — runs the standard FeatureSet.to_model() flow, passing the DAG dict + region + bucket as custom_args. The meta_endpoint.template substitutes those placeholders, the SageMaker training job persists meta_endpoint_config.json as the model artifact, and the model package is registered with the standard inference image.
  5. Deploy — model.to_endpoint(...) deploys an Endpoint, async if any child is async (with max_instances=1 and 5-minute idle drain). inference_batch_size is auto-set to the minimum across DAG children.

Inference flow

When the deployed MetaEndpoint receives a request:

  1. The container deserializes the DAG from the model artifact.
  2. The walker traverses nodes in topological order.
  3. Endpoint nodes call fast_inference (sync child) or async_inference (async child) via workbench.endpoints — the transport is decided per child by the async flags captured at deploy time.
  4. Aggregation nodes apply their combination logic on the upstream outputs.
  5. The output node's DataFrame is returned to the caller (with the synthetic __dag_row_id stripped).

Failure policy is fail-fast: any exception in any node propagates out and the request fails. (Future: per-node failure policies for partial-result aggregation.)

CLI: Ensemble Simulator

Before committing to an ensemble shape, the ensemble_sim CLI lets you evaluate aggregation strategies offline against captured cross-fold predictions:

ensemble_sim logd-reg-xgb logd-reg-pytorch logd-reg-chemprop

This reports per-strategy MAE / RMSE / R² so you can pick the aggregation node that performs best before building the DAG. The same simulator will also be wired into MetaEndpoint.create() for auto-tuning when a DAG includes a strategy-tunable aggregation node (planned).

API Reference

MetaEndpoint: An Endpoint backed by a directed acyclic graph (DAG) of child endpoints and aggregation nodes.

A MetaEndpoint behaves identically to a regular Endpoint at runtime — callers do endpoint.inference(df) and get a DataFrame back. The DAG machinery is server-side: the deployed container loads the serialized DAG, dispatches each child invocation to fast_inference (sync) or async_inference (async), and runs aggregation nodes locally.

Common usage::

from workbench.api import MetaEndpoint
from workbench.utils.meta_endpoint_dag import MetaEndpointDAG
from workbench.utils.aggregation_nodes import Concat

dag = MetaEndpointDAG()
dag.add_endpoint("smiles-to-2d-v1")
dag.add_endpoint("smiles-to-3d-v2")
dag.add_aggregation(Concat(name="combine"))
dag.add_edge("smiles-to-2d-v1", "combine")
dag.add_edge("smiles-to-3d-v2", "combine")
dag.set_input_node("smiles-to-2d-v1", "smiles-to-3d-v2")
dag.set_output_node("combine")

end = MetaEndpoint.create(name="my-features-meta", dag=dag)
# Input does not need any id column — the DAG handles row alignment internally.
df = end.inference(input_df)

If any child endpoint in the DAG is async (e.g. smiles-to-3d-v2), the MetaEndpoint is automatically deployed as async too — its invocation budget needs to accommodate the slowest child.

create(serverless=False) deploys the meta on a realtime instance and requires every child to be realtime as well — a serverless child's cold start would defeat the point of a warm meta.

MetaEndpoint

Bases: Endpoint

Endpoint backed by a :class:MetaEndpointDAG.

Constructor wraps an existing deployed MetaEndpoint by name, identical to :class:Endpoint. Use :meth:create to build and deploy a new one from a DAG.

Source code in src/workbench/api/meta_endpoint.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
class MetaEndpoint(Endpoint):
    """Endpoint backed by a :class:`MetaEndpointDAG`.

    Constructor wraps an existing deployed MetaEndpoint by name, identical
    to :class:`Endpoint`. Use :meth:`create` to build and deploy a new one
    from a DAG.
    """

    @classmethod
    def create(
        cls,
        name: str,
        dag: MetaEndpointDAG,
        description: str | None = None,
        tags: list[str] | None = None,
        serverless: bool = True,
        min_instances: int = 0,
        max_instances: int = 1,
    ) -> "MetaEndpoint":
        """Build, register, and deploy a MetaEndpoint from a DAG.

        Steps:
          1. Validate the DAG; populate per-endpoint async flags.
          2. Backtrace lineage from a primary endpoint to satisfy
             Workbench's Model machinery (FeatureSet, target, features).
          3. Run the standard ``FeatureSet.to_model()`` flow, passing the
             DAG / region / bucket as ``custom_args`` so the meta-endpoint
             template fills them in at training time.
          4. Deploy the endpoint (async if any DAG child is async) and
             stash the serialized DAG on it.

        Args:
            name: Endpoint / Model name.
            dag: A :class:`MetaEndpointDAG` describing the data flow.
            description: Optional description for the registered model.
            tags: Optional list of Workbench tags.
            serverless: Deploy a sync meta as serverless (default: True). False
                deploys it on a realtime instance and raises if any child
                endpoint is serverless. Async metas are always instance-backed.
            min_instances: Autoscaler floor (default: 0 — scale to zero when
                idle; realtime deployments floor at 1). Set to 1 in production
                to keep an async meta warm. Ignored for serverless deployments.
            max_instances: Autoscaler ceiling (default: 1). The meta is a thin
                orchestrator that fans out to children, so it rarely needs more
                than one instance. Ignored for serverless deployments.

        Returns:
            The deployed MetaEndpoint, ready for ``.inference()``.
        """
        if not Artifact.is_name_valid(name, delimiter="-", lower_case=False):
            raise ValueError(f"Invalid MetaEndpoint name: '{name}' (use only alphanumerics and '-')")

        log.important(f"Validating DAG for MetaEndpoint '{name}'...")
        dag.validate()
        dag.populate_child_metadata()
        is_async = dag.has_async_endpoint()
        log.important(
            f"DAG: {len(dag.endpoints)} endpoints, {len(dag.aggregations)} aggregation nodes "
            f"({'async' if is_async else 'sync'} deployment)"
        )

        # A realtime meta stays warm only if its children do; reject serverless children.
        if not serverless:
            serverless_children = sorted(ep for ep in set(dag.endpoints.values()) if Endpoint(ep).is_serverless())
            if serverless_children:
                raise ValueError(
                    f"MetaEndpoint '{name}': realtime deployment requested, but child endpoint(s) "
                    f"{serverless_children} are serverless. Redeploy them realtime first."
                )

        # An async meta sizes its batch from the child fleet (see the batch
        # computation below), which requires every async child to have advertised
        # its `inference_batch_size` + `max_instances` — both stamped on the
        # child's `workbench_meta` when its deploy script completes. Fail fast,
        # before deploying anything, if a child isn't fully deployed yet: the
        # children must come up first.
        if is_async:
            unstamped = [
                ep
                for ep, child_is_async in dag.endpoint_async_flags.items()
                if child_is_async and (ep not in dag.endpoint_batch_sizes or ep not in dag.endpoint_max_instances)
            ]
            if unstamped:
                raise RuntimeError(
                    f"MetaEndpoint '{name}': async child endpoint(s) {unstamped} have not advertised "
                    f"'inference_batch_size' / 'max_instances' — they aren't fully deployed yet. Deploy "
                    f"the child endpoint(s) first (their deploy script stamps these on completion), then "
                    f"create this meta. The meta sizes its async batch from the child fleet, so it cannot "
                    f"be built before its children."
                )

        # Backtrace lineage from a primary endpoint to satisfy Workbench Model
        # machinery (every Model needs a FeatureSet to hang off of).
        feature_list, feature_set_name, target_column = cls._derive_lineage(dag)

        # Build the model via the standard FeatureSet → Model flow. The
        # meta-endpoint template's `{{dag}}`, `{{aws_region}}`, `{{s3_bucket}}`
        # placeholders are filled from custom_args.
        aws_clamp = AWSAccountClamp()
        sm_session = aws_clamp.sagemaker_session()
        workbench_bucket = ConfigManager().get_config("WORKBENCH_BUCKET")

        feature_set = FeatureSet(feature_set_name)
        feature_set.to_model(
            name=name,
            model_type=cls._derive_model_type(dag),
            model_framework=ModelFramework.META,
            tags=tags or [name],
            description=description or f"MetaEndpoint DAG over: {', '.join(dag.endpoints.values())}",
            target_column=target_column,
            feature_list=feature_list,
            custom_args={
                "dag": dag.to_dict(),
                "aws_region": sm_session.boto_region_name,
                "s3_bucket": workbench_bucket,
            },
        )

        # MetaEndpoint containers are thin orchestrators — they receive the
        # full input as one invocation and let each child's own async_inference
        # do the row-level fan-out against the child fleet.
        mode = "async" if is_async else "serverless" if serverless else "realtime"
        log.important(f"Deploying MetaEndpoint '{name}' ({mode})...")
        model = Model(name)
        uses_autoscaler = is_async or not serverless
        endpoint = model.to_endpoint(
            tags=tags or [name],
            serverless=serverless and not is_async,
            async_endpoint=is_async,
            min_instances=min_instances if uses_autoscaler else 0,
            max_instances=max_instances if uses_autoscaler else None,
        )

        # The meta delegates row-level batching to its children, who saturate
        # their own fleets; size the async meta batch to the slowest child's
        # fleet: smallest async-child batch × that fleet's max_instances × 4
        # (≈4 rounds through the fleet). Smaller batches → shorter invocations
        # (less likely to hit the 60-min async cap) and finer-grained failure/
        # retry. All-sync metas stay large (rows are cheap). Child stamps are
        # guaranteed present here — validated right after populate_child_metadata.
        meta_batch = 500
        if is_async:
            async_children = [ep for ep, child_is_async in dag.endpoint_async_flags.items() if child_is_async]
            async_batches = [dag.endpoint_batch_sizes[ep] for ep in async_children]
            async_fleets = [dag.endpoint_max_instances[ep] for ep in async_children]
            meta_batch = min(async_batches) * max(async_fleets) * 4
        endpoint.upsert_workbench_meta(
            {
                "inference_batch_size": meta_batch,
                "meta_endpoint_dag": dag.to_dict(),
            }
        )

        # The DAG already knows its column contract, so register it directly.
        # Auto-discovery would fan out a live inference across every child just
        # to diff what came back. A failure here costs only the warm cache —
        # `Endpoint.output_columns()` re-derives on demand — so it must not
        # take down an endpoint that's already deployed and serving.
        from workbench.utils.endpoint_utils import register_output_columns

        try:
            register_output_columns(endpoint, dag.output_columns())
        except Exception as e:
            log.warning(f"MetaEndpoint '{name}': could not register output columns ({e})")

        log.important(f"MetaEndpoint '{name}' created successfully!")
        return cls(name)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @classmethod
    def _derive_model_type(cls, dag: MetaEndpointDAG) -> ModelType:
        """Pick the most accurate :class:`ModelType` for the DAG's output.

        - Output node is a terminal endpoint → borrow that endpoint's declared
          type (e.g., a downstream predictor endpoint contributes its own type).
        - Output node is :class:`~workbench.utils.aggregation_nodes.Concat` →
          ``TRANSFORMER`` (column-union of feature outputs).
        - Output node is :class:`~workbench.utils.aggregation_nodes.Vote` →
          ``CLASSIFIER`` (majority vote of class labels).
        - Output node is any other prediction aggregator → ``REGRESSOR``.
        """
        from workbench.utils.aggregation_nodes import Concat, Vote

        output_name = dag.output_node
        if output_name in dag.endpoints:
            return Model(dag.endpoints[output_name]).model_type

        agg = dag.aggregations[output_name]
        if isinstance(agg, Concat):
            return ModelType.TRANSFORMER
        if isinstance(agg, Vote):
            return ModelType.CLASSIFIER
        return ModelType.REGRESSOR

    @classmethod
    def _derive_lineage(cls, dag: MetaEndpointDAG) -> tuple[list[str], str, str | list[str] | None]:
        """Derive a (feature_list, feature_set_name, target_column) tuple for the meta.

        Workbench Models need to trace back to a FeatureSet, and downstream
        tooling (``test_inference``, ``register_input_columns`` /
        ``register_output_columns``) all anchor on
        ``model.features()`` + ``model.get_input()``. For DAG-based
        MetaEndpoints we satisfy that contract by:

          - ``feature_list`` / ``feature_set_name`` — borrowed from the first
            input endpoint, since the meta consumes that endpoint's input
            columns and that endpoint's FeatureSet is guaranteed to contain
            usable smoke-test data.
          - ``target_column`` — :meth:`MetaEndpointDAG.terminal_target`,
            which walks back from the output node to the closest predictor
            endpoint(s) so mixed DAGs (smiles → features → predictor) report
            what the meta actually predicts, not what the input endpoint
            happens to produce. A panel over several predictors reports the
            list of targets it covers.
        """
        if not dag.input_nodes:
            raise ValueError("DAG has no input nodes — cannot derive lineage")
        primary_endpoint_name = dag.endpoints[dag.input_nodes[0]]

        ep = Endpoint(primary_endpoint_name)
        if not ep.exists():
            raise ValueError(f"Primary endpoint '{primary_endpoint_name}' does not exist")

        primary_model = Model(ep.get_input())
        feature_list = primary_model.features() or list(dag.input_columns())
        feature_set_name = primary_model.get_input()
        # Target reflects what the DAG ultimately predicts (walks back from
        # the output node), not what the input endpoint happens to produce.
        # For mixed DAGs (smiles → features → predictor), this surfaces the
        # downstream predictor's target instead of the feature endpoint's None.
        target_column = dag.terminal_target()

        log.info(
            f"Lineage anchor: {primary_endpoint_name} -> {primary_model.name} -> {feature_set_name} "
            f"(target: {target_column})"
        )
        return feature_list, feature_set_name, target_column

    def get_dag(self) -> MetaEndpointDAG:
        """Reconstruct the MetaEndpointDAG from this endpoint's stored metadata."""
        meta = self.workbench_meta() or {}
        dag_dict = meta.get("meta_endpoint_dag")
        if not dag_dict:
            raise ValueError(
                f"MetaEndpoint '{self.name}' has no DAG in workbench_meta. Recreate via MetaEndpoint.create()."
            )
        return MetaEndpointDAG.from_dict(dag_dict)

    def run_dag_test(self, input_df: pd.DataFrame) -> pd.DataFrame:
        """Execute the DAG client-side against ``input_df``.

        Bypasses the deployed meta endpoint entirely: each child endpoint is
        invoked directly via the regular ``Endpoint(name).inference()`` API
        from this process. Useful for debugging the DAG topology, isolating
        which child is misbehaving, or running the DAG when the deployed
        meta endpoint is unavailable.

        Result is identical to ``self.inference(input_df)`` modulo transport
        and any container-only side effects (data capture, etc.).
        """
        return self.get_dag().run(input_df)

    def purge_async_queue(self) -> int:
        """Cancel queued async invocations for the meta and every async child.

        When a meta-level batch job is killed, orphaned work can be sitting
        at two layers: invocations queued at the meta endpoint, and child
        invocations queued by in-flight meta predict_fns. This override
        purges both so the entire DAG drains.

        Returns:
            int: Total number of staged input objects deleted across the
            meta and its async children.
        """
        total = super().purge_async_queue()
        dag = self.get_dag()
        for child_name in set(dag.endpoints.values()):
            if dag.endpoint_async_flags.get(child_name):
                total += Endpoint(child_name).purge_async_queue()
        return total

create(name, dag, description=None, tags=None, serverless=True, min_instances=0, max_instances=1) classmethod

Build, register, and deploy a MetaEndpoint from a DAG.

Steps
  1. Validate the DAG; populate per-endpoint async flags.
  2. Backtrace lineage from a primary endpoint to satisfy Workbench's Model machinery (FeatureSet, target, features).
  3. Run the standard FeatureSet.to_model() flow, passing the DAG / region / bucket as custom_args so the meta-endpoint template fills them in at training time.
  4. Deploy the endpoint (async if any DAG child is async) and stash the serialized DAG on it.

Parameters:

Name Type Description Default
name str

Endpoint / Model name.

 required
dag MetaEndpointDAG

A :class:MetaEndpointDAG describing the data flow.

 required
description str | None

Optional description for the registered model.

 None
tags list[str] | None

Optional list of Workbench tags.

 None
serverless bool

Deploy a sync meta as serverless (default: True). False deploys it on a realtime instance and raises if any child endpoint is serverless. Async metas are always instance-backed.

 True
min_instances int

Autoscaler floor (default: 0 — scale to zero when idle; realtime deployments floor at 1). Set to 1 in production to keep an async meta warm. Ignored for serverless deployments.

 0
max_instances int

Autoscaler ceiling (default: 1). The meta is a thin orchestrator that fans out to children, so it rarely needs more than one instance. Ignored for serverless deployments.

 1

Returns:

Type Description
'MetaEndpoint'

The deployed MetaEndpoint, ready for .inference().
Source code in src/workbench/api/meta_endpoint.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
@classmethod
def create(
    cls,
    name: str,
    dag: MetaEndpointDAG,
    description: str | None = None,
    tags: list[str] | None = None,
    serverless: bool = True,
    min_instances: int = 0,
    max_instances: int = 1,
) -> "MetaEndpoint":
    """Build, register, and deploy a MetaEndpoint from a DAG.

    Steps:
      1. Validate the DAG; populate per-endpoint async flags.
      2. Backtrace lineage from a primary endpoint to satisfy
         Workbench's Model machinery (FeatureSet, target, features).
      3. Run the standard ``FeatureSet.to_model()`` flow, passing the
         DAG / region / bucket as ``custom_args`` so the meta-endpoint
         template fills them in at training time.
      4. Deploy the endpoint (async if any DAG child is async) and
         stash the serialized DAG on it.

    Args:
        name: Endpoint / Model name.
        dag: A :class:`MetaEndpointDAG` describing the data flow.
        description: Optional description for the registered model.
        tags: Optional list of Workbench tags.
        serverless: Deploy a sync meta as serverless (default: True). False
            deploys it on a realtime instance and raises if any child
            endpoint is serverless. Async metas are always instance-backed.
        min_instances: Autoscaler floor (default: 0 — scale to zero when
            idle; realtime deployments floor at 1). Set to 1 in production
            to keep an async meta warm. Ignored for serverless deployments.
        max_instances: Autoscaler ceiling (default: 1). The meta is a thin
            orchestrator that fans out to children, so it rarely needs more
            than one instance. Ignored for serverless deployments.

    Returns:
        The deployed MetaEndpoint, ready for ``.inference()``.
    """
    if not Artifact.is_name_valid(name, delimiter="-", lower_case=False):
        raise ValueError(f"Invalid MetaEndpoint name: '{name}' (use only alphanumerics and '-')")

    log.important(f"Validating DAG for MetaEndpoint '{name}'...")
    dag.validate()
    dag.populate_child_metadata()
    is_async = dag.has_async_endpoint()
    log.important(
        f"DAG: {len(dag.endpoints)} endpoints, {len(dag.aggregations)} aggregation nodes "
        f"({'async' if is_async else 'sync'} deployment)"
    )

    # A realtime meta stays warm only if its children do; reject serverless children.
    if not serverless:
        serverless_children = sorted(ep for ep in set(dag.endpoints.values()) if Endpoint(ep).is_serverless())
        if serverless_children:
            raise ValueError(
                f"MetaEndpoint '{name}': realtime deployment requested, but child endpoint(s) "
                f"{serverless_children} are serverless. Redeploy them realtime first."
            )

    # An async meta sizes its batch from the child fleet (see the batch
    # computation below), which requires every async child to have advertised
    # its `inference_batch_size` + `max_instances` — both stamped on the
    # child's `workbench_meta` when its deploy script completes. Fail fast,
    # before deploying anything, if a child isn't fully deployed yet: the
    # children must come up first.
    if is_async:
        unstamped = [
            ep
            for ep, child_is_async in dag.endpoint_async_flags.items()
            if child_is_async and (ep not in dag.endpoint_batch_sizes or ep not in dag.endpoint_max_instances)
        ]
        if unstamped:
            raise RuntimeError(
                f"MetaEndpoint '{name}': async child endpoint(s) {unstamped} have not advertised "
                f"'inference_batch_size' / 'max_instances' — they aren't fully deployed yet. Deploy "
                f"the child endpoint(s) first (their deploy script stamps these on completion), then "
                f"create this meta. The meta sizes its async batch from the child fleet, so it cannot "
                f"be built before its children."
            )

    # Backtrace lineage from a primary endpoint to satisfy Workbench Model
    # machinery (every Model needs a FeatureSet to hang off of).
    feature_list, feature_set_name, target_column = cls._derive_lineage(dag)

    # Build the model via the standard FeatureSet → Model flow. The
    # meta-endpoint template's `{{dag}}`, `{{aws_region}}`, `{{s3_bucket}}`
    # placeholders are filled from custom_args.
    aws_clamp = AWSAccountClamp()
    sm_session = aws_clamp.sagemaker_session()
    workbench_bucket = ConfigManager().get_config("WORKBENCH_BUCKET")

    feature_set = FeatureSet(feature_set_name)
    feature_set.to_model(
        name=name,
        model_type=cls._derive_model_type(dag),
        model_framework=ModelFramework.META,
        tags=tags or [name],
        description=description or f"MetaEndpoint DAG over: {', '.join(dag.endpoints.values())}",
        target_column=target_column,
        feature_list=feature_list,
        custom_args={
            "dag": dag.to_dict(),
            "aws_region": sm_session.boto_region_name,
            "s3_bucket": workbench_bucket,
        },
    )

    # MetaEndpoint containers are thin orchestrators — they receive the
    # full input as one invocation and let each child's own async_inference
    # do the row-level fan-out against the child fleet.
    mode = "async" if is_async else "serverless" if serverless else "realtime"
    log.important(f"Deploying MetaEndpoint '{name}' ({mode})...")
    model = Model(name)
    uses_autoscaler = is_async or not serverless
    endpoint = model.to_endpoint(
        tags=tags or [name],
        serverless=serverless and not is_async,
        async_endpoint=is_async,
        min_instances=min_instances if uses_autoscaler else 0,
        max_instances=max_instances if uses_autoscaler else None,
    )

    # The meta delegates row-level batching to its children, who saturate
    # their own fleets; size the async meta batch to the slowest child's
    # fleet: smallest async-child batch × that fleet's max_instances × 4
    # (≈4 rounds through the fleet). Smaller batches → shorter invocations
    # (less likely to hit the 60-min async cap) and finer-grained failure/
    # retry. All-sync metas stay large (rows are cheap). Child stamps are
    # guaranteed present here — validated right after populate_child_metadata.
    meta_batch = 500
    if is_async:
        async_children = [ep for ep, child_is_async in dag.endpoint_async_flags.items() if child_is_async]
        async_batches = [dag.endpoint_batch_sizes[ep] for ep in async_children]
        async_fleets = [dag.endpoint_max_instances[ep] for ep in async_children]
        meta_batch = min(async_batches) * max(async_fleets) * 4
    endpoint.upsert_workbench_meta(
        {
            "inference_batch_size": meta_batch,
            "meta_endpoint_dag": dag.to_dict(),
        }
    )

    # The DAG already knows its column contract, so register it directly.
    # Auto-discovery would fan out a live inference across every child just
    # to diff what came back. A failure here costs only the warm cache —
    # `Endpoint.output_columns()` re-derives on demand — so it must not
    # take down an endpoint that's already deployed and serving.
    from workbench.utils.endpoint_utils import register_output_columns

    try:
        register_output_columns(endpoint, dag.output_columns())
    except Exception as e:
        log.warning(f"MetaEndpoint '{name}': could not register output columns ({e})")

    log.important(f"MetaEndpoint '{name}' created successfully!")
    return cls(name)

get_dag()

Reconstruct the MetaEndpointDAG from this endpoint's stored metadata.

Source code in src/workbench/api/meta_endpoint.py 
297
298
299
300
301
302
303
304
305
def get_dag(self) -> MetaEndpointDAG:
    """Reconstruct the MetaEndpointDAG from this endpoint's stored metadata."""
    meta = self.workbench_meta() or {}
    dag_dict = meta.get("meta_endpoint_dag")
    if not dag_dict:
        raise ValueError(
            f"MetaEndpoint '{self.name}' has no DAG in workbench_meta. Recreate via MetaEndpoint.create()."
        )
    return MetaEndpointDAG.from_dict(dag_dict)

purge_async_queue()

Cancel queued async invocations for the meta and every async child.

When a meta-level batch job is killed, orphaned work can be sitting at two layers: invocations queued at the meta endpoint, and child invocations queued by in-flight meta predict_fns. This override purges both so the entire DAG drains.

Returns:

Name Type Description
int int

Total number of staged input objects deleted across the
int

meta and its async children.
Source code in src/workbench/api/meta_endpoint.py 
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
def purge_async_queue(self) -> int:
    """Cancel queued async invocations for the meta and every async child.

    When a meta-level batch job is killed, orphaned work can be sitting
    at two layers: invocations queued at the meta endpoint, and child
    invocations queued by in-flight meta predict_fns. This override
    purges both so the entire DAG drains.

    Returns:
        int: Total number of staged input objects deleted across the
        meta and its async children.
    """
    total = super().purge_async_queue()
    dag = self.get_dag()
    for child_name in set(dag.endpoints.values()):
        if dag.endpoint_async_flags.get(child_name):
            total += Endpoint(child_name).purge_async_queue()
    return total

run_dag_test(input_df)

Execute the DAG client-side against input_df.

Bypasses the deployed meta endpoint entirely: each child endpoint is invoked directly via the regular Endpoint(name).inference() API from this process. Useful for debugging the DAG topology, isolating which child is misbehaving, or running the DAG when the deployed meta endpoint is unavailable.

Result is identical to self.inference(input_df) modulo transport and any container-only side effects (data capture, etc.).

Source code in src/workbench/api/meta_endpoint.py 
307
308
309
310
311
312
313
314
315
316
317
318
319
def run_dag_test(self, input_df: pd.DataFrame) -> pd.DataFrame:
    """Execute the DAG client-side against ``input_df``.

    Bypasses the deployed meta endpoint entirely: each child endpoint is
    invoked directly via the regular ``Endpoint(name).inference()`` API
    from this process. Useful for debugging the DAG topology, isolating
    which child is misbehaving, or running the DAG when the deployed
    meta endpoint is unavailable.

    Result is identical to ``self.inference(input_df)`` modulo transport
    and any container-only side effects (data capture, etc.).
    """
    return self.get_dag().run(input_df)

MetaEndpointDAG — a directed acyclic graph of endpoints and aggregation nodes describing an inference-time data flow.

A DAG has two kinds of nodes:

  • Endpoint nodes — references to deployed Workbench Endpoint instances by name. The DAG defers actual Endpoint instantiation until execution / column-contract resolution.

  • Aggregation nodes — instances of :class:AggregationNode subclasses that combine outputs from upstream nodes.

DAG construction is explicit::

dag = MetaEndpointDAG()
dag.add_endpoint("smiles-to-2d-v1")
dag.add_endpoint("smiles-to-3d-v2")
dag.add_aggregation(Concat(name="combine"))
dag.add_edge("smiles-to-2d-v1", "combine")
dag.add_edge("smiles-to-3d-v2", "combine")
dag.set_input_node("smiles-to-2d-v1", "smiles-to-3d-v2")
dag.set_output_node("combine")
dag.validate()

An endpoint can declare which prediction head(s) it owns, and then contributes only those columns::

dag.add_endpoint("cl-mouse-reg-1", targets=["log_cl"])

A multi-task child then contributes one head's worth of columns instead of all of them. Omit targets to take the endpoint's whole output. Two nodes claiming the same target is an error rather than a silent first-branch-wins.

Every endpoint also exposes its primary head unprefixed (prediction, confidence, ...). primary takes just those two under a panel name, so a champion panel stays compact and never depends on what each model called its target::

dag.add_endpoint("logd-value-reg-1", primary="logd")      # -> logd_pred, logd_conf
dag.add_endpoint("pka-a1-value-reg-1", primary="pka_a1")  # -> pka_a1_pred, pka_a1_conf

:meth:MetaEndpointDAG.output_columns gives the resulting column contract without deploying anything.

Row-alignment across parallel branches: the walker injects a synthetic :data:DAG_ROW_ID column at the start of every run() and strips it before returning. Aggregation nodes use it as the join key, so callers do not need to supply (or care about) any id column on their input data.

Validation runs at construction time so misconfigured DAGs fail loud before any inference round-trips.

MetaEndpointDAG

A typed DAG of endpoints + aggregation nodes.

The DAG joins parallel branches using an internal synthetic row id (:data:DAG_ROW_ID) injected by :meth:run — callers don't need to supply any id column on their input.

Source code in src/workbench/utils/meta_endpoint_dag.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
class MetaEndpointDAG:
    """A typed DAG of endpoints + aggregation nodes.

    The DAG joins parallel branches using an internal synthetic row id
    (:data:`DAG_ROW_ID`) injected by :meth:`run` — callers don't need to
    supply any id column on their input.
    """

    def __init__(self):
        self._endpoints: Dict[str, str] = {}  # node_name → endpoint_name
        self._declared_targets: Dict[str, List[str]] = {}  # node_name → head(s) this node owns
        self._primary: Dict[str, str] = {}  # node_name → panel name for the endpoint's primary head
        self._endpoint_async_flags: Dict[str, bool] = {}  # populated by populate_child_metadata()
        self._endpoint_batch_sizes: Dict[str, int] = {}  # populated by populate_child_metadata()
        self._endpoint_max_instances: Dict[str, int] = {}  # populated by populate_child_metadata()
        self._aggregations: Dict[str, AggregationNode] = {}
        self._edges: List[tuple[str, str]] = []  # (from_node, to_node)
        self._input_nodes: List[str] = []
        self._output_node: Optional[str] = None

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    def add_endpoint(
        self,
        endpoint_name: str,
        targets: Optional[List[str]] = None,
        primary: Optional[str] = None,
        node_name: Optional[str] = None,
    ) -> str:
        """Add an endpoint reference to the DAG.

        Args:
            endpoint_name: Name of a deployed Workbench endpoint.
            targets: The prediction head(s) this node contributes. Only those
                heads' columns (``<target>_pred``, ``<target>_confidence``,
                quantiles, ...) flow downstream; every other head, the
                unprefixed primary aliases (``prediction``, ``confidence``,
                ``q_*``) and any upstream feature-endpoint columns are
                dropped. Omit to take the endpoint's whole output
                — right for feature endpoints and for a single-task child whose
                heads can't collide with anything else in the DAG.
            primary: Panel name for the endpoint's primary head. Only the
                unprefixed ``prediction`` and ``confidence`` flow downstream,
                as ``<primary>_pred`` and ``<primary>_conf`` — so the panel is
                compact and independent of what the model called its target.
                Mutually exclusive with ``targets``.
            node_name: Optional unique node name (defaults to ``endpoint_name``).

        Returns:
            The node name (so callers can chain).
        """
        node = node_name or endpoint_name
        if node in self._endpoints or node in self._aggregations:
            raise ValueError(f"Node '{node}' already exists in this DAG")
        if targets and primary:
            raise ValueError(f"Node '{node}': pass targets= or primary=, not both")
        self._endpoints[node] = endpoint_name
        if targets:
            self._declared_targets[node] = list(targets)
        if primary:
            self._primary[node] = primary
        return node

    def _owned_heads(self, node: str) -> List[str]:
        """Heads this node contributes downstream, by their panel names."""
        if node in self._primary:
            return [self._primary[node]]
        return list(self._declared_targets.get(node, []))

    def add_aggregation(self, node: AggregationNode) -> str:
        """Add an :class:`AggregationNode` to the DAG.

        The node's ``name`` must be unique across the DAG.
        """
        if node.name in self._endpoints or node.name in self._aggregations:
            raise ValueError(f"Node '{node.name}' already exists in this DAG")
        self._aggregations[node.name] = node
        return node.name

    def add_edge(self, from_node: str, to_node: str) -> None:
        """Declare data flow from ``from_node`` to ``to_node``.

        Endpoint nodes accept at most one inbound edge (their input DataFrame
        comes from a single upstream producer). Aggregation nodes can have
        any number of inbound edges.
        """
        if from_node not in self._all_nodes():
            raise ValueError(f"Edge from unknown node '{from_node}'")
        if to_node not in self._all_nodes():
            raise ValueError(f"Edge to unknown node '{to_node}'")
        if to_node in self._endpoints and self._parents_of(to_node):
            raise ValueError(
                f"Endpoint node '{to_node}' already has an upstream parent "
                f"('{self._parents_of(to_node)[0]}'); endpoints take input "
                f"from at most one source."
            )
        self._edges.append((from_node, to_node))

    def set_input_node(self, *nodes: str) -> None:
        """Declare which nodes receive the DAG's input DataFrame directly."""
        for n in nodes:
            if n not in self._endpoints:
                raise ValueError(f"Input nodes must be endpoint nodes; '{n}' is not")
        self._input_nodes = list(nodes)

    def set_output_node(self, node: str) -> None:
        """Declare the terminal node whose output is the DAG's output."""
        if node not in self._all_nodes():
            raise ValueError(f"Unknown output node '{node}'")
        self._output_node = node

    # ------------------------------------------------------------------
    # Inspection
    # ------------------------------------------------------------------

    @property
    def endpoints(self) -> Dict[str, str]:
        """Mapping of node_name → endpoint_name (read-only view)."""
        return self._endpoints

    @property
    def declared_targets(self) -> Dict[str, List[str]]:
        """Mapping of node_name → declared target head(s) (read-only view)."""
        return self._declared_targets

    @property
    def primary_heads(self) -> Dict[str, str]:
        """Mapping of node_name → panel name for its primary head (read-only view)."""
        return self._primary

    @property
    def aggregations(self) -> Dict[str, AggregationNode]:
        """Mapping of node_name → :class:`AggregationNode` (read-only view)."""
        return self._aggregations

    @property
    def input_nodes(self) -> List[str]:
        """Node names that receive the DAG's input DataFrame directly (read-only)."""
        return self._input_nodes

    @property
    def output_node(self) -> Optional[str]:
        """Node name whose output is the DAG's output."""
        return self._output_node

    @property
    def endpoint_async_flags(self) -> Dict[str, bool]:
        """Mapping of endpoint_name → is_async (populated by :meth:`populate_child_metadata`)."""
        return self._endpoint_async_flags

    @property
    def endpoint_batch_sizes(self) -> Dict[str, int]:
        """Mapping of endpoint_name → inference_batch_size (populated by :meth:`populate_child_metadata`)."""
        return self._endpoint_batch_sizes

    @property
    def endpoint_max_instances(self) -> Dict[str, int]:
        """Mapping of endpoint_name → max_instances fleet ceiling (populated by :meth:`populate_child_metadata`)."""
        return self._endpoint_max_instances

    def _all_nodes(self) -> List[str]:
        return list(self._endpoints.keys()) + list(self._aggregations.keys())

    def _parents_of(self, node: str) -> List[str]:
        return [src for src, dst in self._edges if dst == node]

    def topological_order(self) -> List[str]:
        """Return nodes in topological order (parents before children).

        Raises:
            ValueError: If the DAG contains a cycle.
        """
        in_degree = {n: 0 for n in self._all_nodes()}
        for _, dst in self._edges:
            in_degree[dst] += 1

        ready = [n for n, deg in in_degree.items() if deg == 0]
        order: List[str] = []
        while ready:
            node = ready.pop(0)
            order.append(node)
            for src, dst in self._edges:
                if src == node:
                    in_degree[dst] -= 1
                    if in_degree[dst] == 0:
                        ready.append(dst)

        if len(order) != len(in_degree):
            raise ValueError("DAG contains a cycle")
        return order

    # ------------------------------------------------------------------
    # Column contract
    # ------------------------------------------------------------------

    def input_columns(self) -> List[str]:
        """Union of input columns required by every node that receives the
        caller's input directly.

        Used by :class:`MetaEndpoint` as a fallback when deriving the
        feature list during lineage anchoring.
        """
        from workbench.api import Endpoint

        if not self._input_nodes:
            raise ValueError("DAG has no input nodes — call set_input_node() first")

        seen = set()
        cols: List[str] = []
        for node in self._input_nodes:
            for c in Endpoint(self._endpoints[node]).input_columns():
                if c not in seen:
                    seen.add(c)
                    cols.append(c)
        return cols

    def output_columns(self) -> List[str]:
        """Columns the DAG adds, as declared by its nodes — no inference required.

        Walks the topology feeding each endpoint's registered
        :meth:`Endpoint.output_columns` into each aggregation's
        :meth:`AggregationNode.output_columns`. A node with declared
        ``targets`` contributes only those heads' columns.

        The caller's own input columns pass through and are not listed, matching
        :meth:`Endpoint.output_columns` semantics. Size a panel before deploying
        it with ``len(dag.output_columns())``.
        """
        from workbench.api import Endpoint

        if self._output_node is None:
            raise ValueError("DAG has no output node — call set_output_node() first")

        node_columns: Dict[str, List[str]] = {}
        for node in self.topological_order():
            if node in self._endpoints:
                endpoint_name = self._endpoints[node]
                cols = list(Endpoint(endpoint_name).output_columns())
                context = f"Node '{node}' (endpoint '{endpoint_name}')"
                parents = self._parents_of(node)
                upstream = node_columns[parents[0]] if parents else []
                # A node that owns heads contributes only those; upstream feature columns stop here.
                if node in self._primary:
                    node_columns[node] = list(select_primary_columns(cols, self._primary[node], context).values())
                elif node in self._declared_targets:
                    node_columns[node] = select_target_columns(cols, self._declared_targets[node], context)
                else:
                    node_columns[node] = upstream + [c for c in cols if c not in upstream]
            else:
                agg = self._aggregations[node]
                node_columns[node] = agg.output_columns([node_columns[p] for p in self._parents_of(node)])

        return node_columns[self._output_node]

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------

    def validate(self) -> "MetaEndpointDAG":
        """Validate the DAG. Returns self for chaining; raises on failure.

        Checks:
          - At least one input node and exactly one output node declared
          - No cycles
          - Aggregation nodes have at least one parent
          - Endpoint nodes are either input nodes (zero parents) or have
            exactly one upstream parent — never both
          - The output node is reachable from the input nodes
          - No two nodes declare the same target head

        Topology only — no AWS round-trips. A declared head colliding with an
        *undeclared* child's heads needs those children's targets, so it
        surfaces in :meth:`terminal_target` instead.
        """
        if not self._input_nodes:
            raise ValueError("DAG has no input nodes")
        if self._output_node is None:
            raise ValueError("DAG has no output node")

        order = self.topological_order()  # raises on cycle

        for ep_node in self._endpoints:
            parents = self._parents_of(ep_node)
            is_input = ep_node in self._input_nodes
            if is_input and parents:
                raise ValueError(
                    f"Endpoint node '{ep_node}' is declared as an input node but has "
                    f"upstream parents {parents}; pick one or the other."
                )
            if not is_input and not parents:
                raise ValueError(
                    f"Endpoint node '{ep_node}' has no upstream parent and is not "
                    f"declared as an input node — it has no source for its input DataFrame."
                )

        for name in self._aggregations:
            if not self._parents_of(name):
                raise ValueError(f"Aggregation node '{name}' has no upstream parents")

        owner: Dict[str, str] = {}
        for node in list(self._declared_targets) + list(self._primary):
            for target in self._owned_heads(node):
                if target in owner:
                    raise ValueError(
                        f"Nodes '{owner[target]}' and '{node}' both declare target '{target}' — "
                        f"a target can be owned by only one node."
                    )
                owner[target] = node

        reachable = set(self._input_nodes)
        for node in order:
            if node in reachable:
                for src, dst in self._edges:
                    if src == node:
                        reachable.add(dst)
        if self._output_node not in reachable:
            raise ValueError(f"Output node '{self._output_node}' is not reachable from input nodes {self._input_nodes}")

        return self

    # ------------------------------------------------------------------
    # Execution (client-side walker)
    # ------------------------------------------------------------------

    def run(
        self,
        input_df: pd.DataFrame,
        endpoint_invoker: Optional[EndpointInvoker] = None,
    ) -> pd.DataFrame:
        """Execute the DAG against ``input_df`` and return the output node's DataFrame.

        The walker injects a synthetic :data:`DAG_ROW_ID` column at entry
        (used internally to align rows across parallel branches) and
        strips it before returning. Callers don't need to supply any id
        column.

        Walks the DAG in waves: every node whose parents have all finished
        runs concurrently, so parallel branches invoke their endpoints at the
        same time. A cold panel of serverless children then costs one cold
        start rather than one per child. Endpoint nodes call
        :meth:`Endpoint.inference` on either the caller's ``input_df`` (input
        nodes) or their upstream parent's output. Aggregation nodes receive the
        outputs of all their parents and apply their combination logic.

        Failure policy is fail-fast: any exception in any node propagates
        out and the DAG run aborts once its wave finishes.

        Args:
            input_df: DataFrame supplied by the caller. Must contain the
                columns required by every input-node endpoint. No id
                column is required.
            endpoint_invoker: Optional callable ``(endpoint_name, df) -> df``
                used to invoke endpoint nodes. Defaults to using the full
                Workbench ``Endpoint`` API class — appropriate for client-side
                use. Pass a ``fast_inference``-backed invoker when running
                inside a deployed SageMaker container where the full
                Workbench config isn't available.

        Returns:
            The DataFrame at the DAG's output node, with the synthetic
            :data:`DAG_ROW_ID` column removed.
        """
        if self._output_node is None:
            raise ValueError("DAG has no output node — call set_output_node() first")
        if DAG_ROW_ID in input_df.columns:
            raise ValueError(
                f"input_df already contains the reserved column '{DAG_ROW_ID}'. " f"Remove it before calling run()."
            )

        # Inject the synthetic row id. Endpoints will pass this through as an
        # unknown input column; aggregation nodes use it as their join key.
        input_df = input_df.copy()
        input_df[DAG_ROW_ID] = range(len(input_df))

        def run_node(node: str) -> pd.DataFrame:
            if node in self._endpoints:
                return self._run_endpoint(node, input_df, outputs, endpoint_invoker)
            return self._run_aggregation(node, outputs)

        # Workers only read finished parents; `outputs` is written between waves.
        outputs: Dict[str, pd.DataFrame] = {}
        pending = self.topological_order()
        with ThreadPoolExecutor(max_workers=max(1, len(self._endpoints))) as pool:
            while pending:
                wave = [n for n in pending if all(p in outputs for p in self._parents_of(n))]
                futures = {node: pool.submit(run_node, node) for node in wave}
                for node, future in futures.items():
                    outputs[node] = future.result()
                pending = [n for n in pending if n not in outputs]

        result = outputs[self._output_node]
        if DAG_ROW_ID in result.columns:
            result = result.drop(columns=[DAG_ROW_ID])
        return result

    def _run_endpoint(
        self,
        node: str,
        input_df: pd.DataFrame,
        outputs: Dict[str, pd.DataFrame],
        endpoint_invoker: Optional[EndpointInvoker],
    ) -> pd.DataFrame:
        """Execute a single endpoint node.

        Source DataFrame is the caller's input for input nodes, or the
        single upstream parent's output otherwise. The full DataFrame is
        passed to ``endpoint.inference()`` — metadata columns
        (project_id, owner, etc.) flow through alongside the endpoint's
        added columns, matching standard Workbench inference behavior.

        The walker-injected :data:`DAG_ROW_ID` column must survive the
        endpoint round-trip so downstream aggregation nodes can join on
        it. If an endpoint silently strips unknown input columns, this
        will fail loudly — better than misaligned rows.

        A node with declared ``targets`` is trimmed to those heads: every
        column that isn't ``<target>`` or ``<target>_*`` is dropped — the
        endpoint's other heads and any feature columns an upstream endpoint
        added — so a multi-task child contributes one head's worth of
        columns instead of all of them. The caller's own input columns are
        never touched.
        """
        endpoint_name = self._endpoints[node]
        parents = self._parents_of(node)
        source_df = input_df if not parents else outputs[parents[0]]

        if endpoint_invoker is not None:
            result = endpoint_invoker(endpoint_name, source_df)
        else:
            from workbench.api import Endpoint

            result = Endpoint(endpoint_name).inference(source_df)

        if DAG_ROW_ID not in result.columns:
            raise RuntimeError(
                f"Endpoint '{endpoint_name}' dropped the walker-injected '{DAG_ROW_ID}' "
                f"column from its output. The DAG can't align rows across branches "
                f"without it. Endpoints must pass unknown input columns through to "
                f"their output."
            )

        if node in self._primary or node in self._declared_targets:
            passthrough = set(input_df.columns) | {DAG_ROW_ID}
            result = self._trim_to_owned_heads(node, endpoint_name, passthrough, result)
        return result

    def _trim_to_owned_heads(
        self, node: str, endpoint_name: str, passthrough: set, result: pd.DataFrame
    ) -> pd.DataFrame:
        """Keep the caller's ``passthrough`` columns plus the heads this node owns, under their panel names."""
        context = f"Node '{node}' (endpoint '{endpoint_name}')"
        added = [c for c in result.columns if c not in passthrough]
        if node in self._primary:
            mapping = select_primary_columns(added, self._primary[node], context)
        else:
            mapping = {c: c for c in select_target_columns(added, self._declared_targets[node], context)}
        kept = [c for c in result.columns if c in passthrough] + list(mapping)
        return result[kept].rename(columns=mapping)

    def _run_aggregation(self, node: str, outputs: Dict[str, pd.DataFrame]) -> pd.DataFrame:
        """Execute a single aggregation node."""
        agg = self._aggregations[node]
        upstream = [outputs[p] for p in self._parents_of(node)]
        return agg.apply(upstream)

    # ------------------------------------------------------------------
    # Serialization (model artifact + workbench_meta storage)
    # ------------------------------------------------------------------

    def to_dict(self) -> dict:
        """Serialize the DAG topology to a JSON-friendly dict.

        Aggregation nodes are serialized by class name + constructor kwargs;
        deserialization (:meth:`from_dict`) requires the same class to be
        importable.

        Per-endpoint ``is_async`` flags are included only if
        :meth:`populate_child_metadata` has been called. The deployed
        inference container relies on these flags to dispatch invocations
        to ``fast_inference`` or ``async_inference``.
        """
        return {
            "endpoints": dict(self._endpoints),
            "endpoint_targets": {node: list(t) for node, t in self._declared_targets.items()},
            "endpoint_primary": dict(self._primary),
            "endpoint_async": dict(self._endpoint_async_flags),
            "endpoint_batch_size": dict(self._endpoint_batch_sizes),
            "aggregations": [_serialize_aggregation(a) for a in self._aggregations.values()],
            "edges": [list(e) for e in self._edges],
            "input_nodes": list(self._input_nodes),
            "output_node": self._output_node,
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), indent=2)

    def populate_child_metadata(self) -> None:
        """Look up each child's async flag and ``inference_batch_size`` and store both.

        Keyed by endpoint name (not node name) so the deployed invoker can both
        dispatch (``async_inference`` vs ``fast_inference``) and size each child's
        batch to what that child is configured for — a slow GFN2-xTB child wants a
        small batch so one chunk doesn't exceed the per-invocation timeout. A child
        with no/invalid ``inference_batch_size`` is omitted, leaving the invoker on
        ``async_inference``'s default.

        One AWS round-trip per unique endpoint name (reads the meta once for both
        values); called by :meth:`MetaEndpoint.create` before serializing the DAG.
        """
        from workbench.api import Endpoint

        for endpoint_name in set(self._endpoints.values()):
            meta = Endpoint(endpoint_name).workbench_meta() or {}
            self._endpoint_async_flags[endpoint_name] = bool(meta.get("async_endpoint"))
            batch_size = meta.get("inference_batch_size")
            try:
                if batch_size is not None and int(batch_size) > 0:
                    self._endpoint_batch_sizes[endpoint_name] = int(batch_size)
            except (TypeError, ValueError):
                pass  # non-numeric / bogus → leave unset, invoker uses its default
            max_instances = meta.get("max_instances")
            try:
                if max_instances is not None and int(max_instances) > 0:
                    self._endpoint_max_instances[endpoint_name] = int(max_instances)
            except (TypeError, ValueError):
                pass  # non-numeric / bogus → leave unset; meta batch falls back to its default

    def has_async_endpoint(self) -> bool:
        """Return True if any endpoint in the DAG is deployed as async.

        Used by :meth:`MetaEndpoint.create` to decide whether the meta
        endpoint itself must be deployed as async. Lazily calls
        :meth:`populate_child_metadata` if not yet populated.
        """
        if not self._endpoint_async_flags and self._endpoints:
            self.populate_child_metadata()
        return any(self._endpoint_async_flags.values())

    def terminal_target(self) -> Optional[Union[str, List[str]]]:
        """Target column(s) the DAG ultimately predicts, or ``None`` for feature pipelines.

        - Output is an endpoint → its targets.
        - Output is an aggregation → walk back to the closest endpoint(s) and
          union their targets in branch order.

        A node's targets are the ones declared on :meth:`add_endpoint`, else
        every head its model predicts. Returns a bare string for a single
        target, a list for a panel, and ``None`` when nothing upstream predicts
        anything (a feature pipeline).

        Used by :meth:`MetaEndpoint._derive_lineage` to anchor the meta's
        ``target_column`` on what the DAG actually predicts rather than
        inheriting from the (possibly target-less) input endpoint.

        Raises:
            ValueError: Two upstream nodes claim the same target. Declare
                ``targets=`` on :meth:`add_endpoint` to say which one owns it.
        """
        from workbench.api import Endpoint, Model

        def _targets_of(node: str) -> List[str]:
            if node in self._primary or node in self._declared_targets:
                return self._owned_heads(node)
            ep = Endpoint(self._endpoints[node])
            if not ep.exists():
                return []
            target = Model(ep.get_input()).target()
            if not target:
                return []
            return list(target) if isinstance(target, list) else [target]

        if self._output_node in self._endpoints:
            predictors = [self._output_node]
        else:
            # Aggregation output — BFS back until we hit endpoints, in branch order.
            predictors = []
            seen: set = set()
            queue: List[str] = list(self._parents_of(self._output_node))
            while queue:
                node = queue.pop(0)
                if node in seen:
                    continue
                seen.add(node)
                if node in self._endpoints:
                    predictors.append(node)
                else:
                    queue.extend(self._parents_of(node))

        targets: List[str] = []
        owner: Dict[str, str] = {}
        for node in predictors:
            for target in _targets_of(node):
                if target in owner:
                    raise ValueError(
                        f"Nodes '{owner[target]}' and '{node}' both predict '{target}'. Declare "
                        f"targets= on add_endpoint() to say which node owns it."
                    )
                owner[target] = node
                targets.append(target)

        if not targets:
            return None
        return targets[0] if len(targets) == 1 else targets

    @classmethod
    def from_dict(cls, data: dict) -> "MetaEndpointDAG":
        dag = cls()
        declared = data.get("endpoint_targets", {})
        primary = data.get("endpoint_primary", {})
        for node_name, endpoint_name in data.get("endpoints", {}).items():
            dag.add_endpoint(
                endpoint_name, declared.get(node_name), primary=primary.get(node_name), node_name=node_name
            )
        dag._endpoint_async_flags = dict(data.get("endpoint_async", {}))
        dag._endpoint_batch_sizes = {k: int(v) for k, v in data.get("endpoint_batch_size", {}).items()}
        for agg_data in data.get("aggregations", []):
            dag.add_aggregation(_deserialize_aggregation(agg_data))
        for src, dst in data.get("edges", []):
            dag.add_edge(src, dst)
        if data.get("input_nodes"):
            dag.set_input_node(*data["input_nodes"])
        if data.get("output_node"):
            dag.set_output_node(data["output_node"])
        return dag

    @classmethod
    def from_json(cls, payload: str) -> "MetaEndpointDAG":
        return cls.from_dict(json.loads(payload))

aggregations property

Mapping of node_name → :class:AggregationNode (read-only view).

declared_targets property

Mapping of node_name → declared target head(s) (read-only view).

endpoint_async_flags property

Mapping of endpoint_name → is_async (populated by :meth:populate_child_metadata).

endpoint_batch_sizes property

Mapping of endpoint_name → inference_batch_size (populated by :meth:populate_child_metadata).

endpoint_max_instances property

Mapping of endpoint_name → max_instances fleet ceiling (populated by :meth:populate_child_metadata).

endpoints property

Mapping of node_name → endpoint_name (read-only view).

input_nodes property

Node names that receive the DAG's input DataFrame directly (read-only).

output_node property

Node name whose output is the DAG's output.

primary_heads property

Mapping of node_name → panel name for its primary head (read-only view).

add_aggregation(node)

Add an :class:AggregationNode to the DAG.

The node's name must be unique across the DAG.

Source code in src/workbench/utils/meta_endpoint_dag.py 
173
174
175
176
177
178
179
180
181
def add_aggregation(self, node: AggregationNode) -> str:
    """Add an :class:`AggregationNode` to the DAG.

    The node's ``name`` must be unique across the DAG.
    """
    if node.name in self._endpoints or node.name in self._aggregations:
        raise ValueError(f"Node '{node.name}' already exists in this DAG")
    self._aggregations[node.name] = node
    return node.name

add_edge(from_node, to_node)

Declare data flow from from_node to to_node.

Endpoint nodes accept at most one inbound edge (their input DataFrame comes from a single upstream producer). Aggregation nodes can have any number of inbound edges.

Source code in src/workbench/utils/meta_endpoint_dag.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def add_edge(self, from_node: str, to_node: str) -> None:
    """Declare data flow from ``from_node`` to ``to_node``.

    Endpoint nodes accept at most one inbound edge (their input DataFrame
    comes from a single upstream producer). Aggregation nodes can have
    any number of inbound edges.
    """
    if from_node not in self._all_nodes():
        raise ValueError(f"Edge from unknown node '{from_node}'")
    if to_node not in self._all_nodes():
        raise ValueError(f"Edge to unknown node '{to_node}'")
    if to_node in self._endpoints and self._parents_of(to_node):
        raise ValueError(
            f"Endpoint node '{to_node}' already has an upstream parent "
            f"('{self._parents_of(to_node)[0]}'); endpoints take input "
            f"from at most one source."
        )
    self._edges.append((from_node, to_node))

add_endpoint(endpoint_name, targets=None, primary=None, node_name=None)

Add an endpoint reference to the DAG.

Parameters:

Name Type Description Default
endpoint_name str

Name of a deployed Workbench endpoint.

 required
targets Optional[List[str]]

The prediction head(s) this node contributes. Only those heads' columns (<target>_pred, <target>_confidence, quantiles, ...) flow downstream; every other head, the unprefixed primary aliases (prediction, confidence, q_*) and any upstream feature-endpoint columns are dropped. Omit to take the endpoint's whole output — right for feature endpoints and for a single-task child whose heads can't collide with anything else in the DAG.

 None
primary Optional[str]

Panel name for the endpoint's primary head. Only the unprefixed prediction and confidence flow downstream, as <primary>_pred and <primary>_conf — so the panel is compact and independent of what the model called its target. Mutually exclusive with targets.

 None
node_name Optional[str]

Optional unique node name (defaults to endpoint_name).

 None

Returns:

Type Description
str

The node name (so callers can chain).
Source code in src/workbench/utils/meta_endpoint_dag.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def add_endpoint(
    self,
    endpoint_name: str,
    targets: Optional[List[str]] = None,
    primary: Optional[str] = None,
    node_name: Optional[str] = None,
) -> str:
    """Add an endpoint reference to the DAG.

    Args:
        endpoint_name: Name of a deployed Workbench endpoint.
        targets: The prediction head(s) this node contributes. Only those
            heads' columns (``<target>_pred``, ``<target>_confidence``,
            quantiles, ...) flow downstream; every other head, the
            unprefixed primary aliases (``prediction``, ``confidence``,
            ``q_*``) and any upstream feature-endpoint columns are
            dropped. Omit to take the endpoint's whole output
            — right for feature endpoints and for a single-task child whose
            heads can't collide with anything else in the DAG.
        primary: Panel name for the endpoint's primary head. Only the
            unprefixed ``prediction`` and ``confidence`` flow downstream,
            as ``<primary>_pred`` and ``<primary>_conf`` — so the panel is
            compact and independent of what the model called its target.
            Mutually exclusive with ``targets``.
        node_name: Optional unique node name (defaults to ``endpoint_name``).

    Returns:
        The node name (so callers can chain).
    """
    node = node_name or endpoint_name
    if node in self._endpoints or node in self._aggregations:
        raise ValueError(f"Node '{node}' already exists in this DAG")
    if targets and primary:
        raise ValueError(f"Node '{node}': pass targets= or primary=, not both")
    self._endpoints[node] = endpoint_name
    if targets:
        self._declared_targets[node] = list(targets)
    if primary:
        self._primary[node] = primary
    return node

has_async_endpoint()

Return True if any endpoint in the DAG is deployed as async.

Used by :meth:MetaEndpoint.create to decide whether the meta endpoint itself must be deployed as async. Lazily calls :meth:populate_child_metadata if not yet populated.

Source code in src/workbench/utils/meta_endpoint_dag.py 
631
632
633
634
635
636
637
638
639
640
def has_async_endpoint(self) -> bool:
    """Return True if any endpoint in the DAG is deployed as async.

    Used by :meth:`MetaEndpoint.create` to decide whether the meta
    endpoint itself must be deployed as async. Lazily calls
    :meth:`populate_child_metadata` if not yet populated.
    """
    if not self._endpoint_async_flags and self._endpoints:
        self.populate_child_metadata()
    return any(self._endpoint_async_flags.values())

input_columns()

Union of input columns required by every node that receives the caller's input directly.

Used by :class:MetaEndpoint as a fallback when deriving the feature list during lineage anchoring.

Source code in src/workbench/utils/meta_endpoint_dag.py 
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
def input_columns(self) -> List[str]:
    """Union of input columns required by every node that receives the
    caller's input directly.

    Used by :class:`MetaEndpoint` as a fallback when deriving the
    feature list during lineage anchoring.
    """
    from workbench.api import Endpoint

    if not self._input_nodes:
        raise ValueError("DAG has no input nodes — call set_input_node() first")

    seen = set()
    cols: List[str] = []
    for node in self._input_nodes:
        for c in Endpoint(self._endpoints[node]).input_columns():
            if c not in seen:
                seen.add(c)
                cols.append(c)
    return cols

output_columns()

Columns the DAG adds, as declared by its nodes — no inference required.

Walks the topology feeding each endpoint's registered :meth:Endpoint.output_columns into each aggregation's :meth:AggregationNode.output_columns. A node with declared targets contributes only those heads' columns.

The caller's own input columns pass through and are not listed, matching :meth:Endpoint.output_columns semantics. Size a panel before deploying it with len(dag.output_columns()).

Source code in src/workbench/utils/meta_endpoint_dag.py 
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def output_columns(self) -> List[str]:
    """Columns the DAG adds, as declared by its nodes — no inference required.

    Walks the topology feeding each endpoint's registered
    :meth:`Endpoint.output_columns` into each aggregation's
    :meth:`AggregationNode.output_columns`. A node with declared
    ``targets`` contributes only those heads' columns.

    The caller's own input columns pass through and are not listed, matching
    :meth:`Endpoint.output_columns` semantics. Size a panel before deploying
    it with ``len(dag.output_columns())``.
    """
    from workbench.api import Endpoint

    if self._output_node is None:
        raise ValueError("DAG has no output node — call set_output_node() first")

    node_columns: Dict[str, List[str]] = {}
    for node in self.topological_order():
        if node in self._endpoints:
            endpoint_name = self._endpoints[node]
            cols = list(Endpoint(endpoint_name).output_columns())
            context = f"Node '{node}' (endpoint '{endpoint_name}')"
            parents = self._parents_of(node)
            upstream = node_columns[parents[0]] if parents else []
            # A node that owns heads contributes only those; upstream feature columns stop here.
            if node in self._primary:
                node_columns[node] = list(select_primary_columns(cols, self._primary[node], context).values())
            elif node in self._declared_targets:
                node_columns[node] = select_target_columns(cols, self._declared_targets[node], context)
            else:
                node_columns[node] = upstream + [c for c in cols if c not in upstream]
        else:
            agg = self._aggregations[node]
            node_columns[node] = agg.output_columns([node_columns[p] for p in self._parents_of(node)])

    return node_columns[self._output_node]

populate_child_metadata()

Look up each child's async flag and inference_batch_size and store both.

Keyed by endpoint name (not node name) so the deployed invoker can both dispatch (async_inference vs fast_inference) and size each child's batch to what that child is configured for — a slow GFN2-xTB child wants a small batch so one chunk doesn't exceed the per-invocation timeout. A child with no/invalid inference_batch_size is omitted, leaving the invoker on async_inference's default.

One AWS round-trip per unique endpoint name (reads the meta once for both values); called by :meth:MetaEndpoint.create before serializing the DAG.

Source code in src/workbench/utils/meta_endpoint_dag.py 
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
def populate_child_metadata(self) -> None:
    """Look up each child's async flag and ``inference_batch_size`` and store both.

    Keyed by endpoint name (not node name) so the deployed invoker can both
    dispatch (``async_inference`` vs ``fast_inference``) and size each child's
    batch to what that child is configured for — a slow GFN2-xTB child wants a
    small batch so one chunk doesn't exceed the per-invocation timeout. A child
    with no/invalid ``inference_batch_size`` is omitted, leaving the invoker on
    ``async_inference``'s default.

    One AWS round-trip per unique endpoint name (reads the meta once for both
    values); called by :meth:`MetaEndpoint.create` before serializing the DAG.
    """
    from workbench.api import Endpoint

    for endpoint_name in set(self._endpoints.values()):
        meta = Endpoint(endpoint_name).workbench_meta() or {}
        self._endpoint_async_flags[endpoint_name] = bool(meta.get("async_endpoint"))
        batch_size = meta.get("inference_batch_size")
        try:
            if batch_size is not None and int(batch_size) > 0:
                self._endpoint_batch_sizes[endpoint_name] = int(batch_size)
        except (TypeError, ValueError):
            pass  # non-numeric / bogus → leave unset, invoker uses its default
        max_instances = meta.get("max_instances")
        try:
            if max_instances is not None and int(max_instances) > 0:
                self._endpoint_max_instances[endpoint_name] = int(max_instances)
        except (TypeError, ValueError):
            pass  # non-numeric / bogus → leave unset; meta batch falls back to its default

run(input_df, endpoint_invoker=None)

Execute the DAG against input_df and return the output node's DataFrame.

The walker injects a synthetic :data:DAG_ROW_ID column at entry (used internally to align rows across parallel branches) and strips it before returning. Callers don't need to supply any id column.

Walks the DAG in waves: every node whose parents have all finished runs concurrently, so parallel branches invoke their endpoints at the same time. A cold panel of serverless children then costs one cold start rather than one per child. Endpoint nodes call :meth:Endpoint.inference on either the caller's input_df (input nodes) or their upstream parent's output. Aggregation nodes receive the outputs of all their parents and apply their combination logic.

Failure policy is fail-fast: any exception in any node propagates out and the DAG run aborts once its wave finishes.

Parameters:

Name Type Description Default
input_df DataFrame

DataFrame supplied by the caller. Must contain the columns required by every input-node endpoint. No id column is required.

 required
endpoint_invoker Optional[EndpointInvoker]

Optional callable (endpoint_name, df) -> df used to invoke endpoint nodes. Defaults to using the full Workbench Endpoint API class — appropriate for client-side use. Pass a fast_inference-backed invoker when running inside a deployed SageMaker container where the full Workbench config isn't available.

 None

Returns:

Type Description
DataFrame

The DataFrame at the DAG's output node, with the synthetic
DataFrame

data:DAG_ROW_ID column removed.
Source code in src/workbench/utils/meta_endpoint_dag.py 
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
def run(
    self,
    input_df: pd.DataFrame,
    endpoint_invoker: Optional[EndpointInvoker] = None,
) -> pd.DataFrame:
    """Execute the DAG against ``input_df`` and return the output node's DataFrame.

    The walker injects a synthetic :data:`DAG_ROW_ID` column at entry
    (used internally to align rows across parallel branches) and
    strips it before returning. Callers don't need to supply any id
    column.

    Walks the DAG in waves: every node whose parents have all finished
    runs concurrently, so parallel branches invoke their endpoints at the
    same time. A cold panel of serverless children then costs one cold
    start rather than one per child. Endpoint nodes call
    :meth:`Endpoint.inference` on either the caller's ``input_df`` (input
    nodes) or their upstream parent's output. Aggregation nodes receive the
    outputs of all their parents and apply their combination logic.

    Failure policy is fail-fast: any exception in any node propagates
    out and the DAG run aborts once its wave finishes.

    Args:
        input_df: DataFrame supplied by the caller. Must contain the
            columns required by every input-node endpoint. No id
            column is required.
        endpoint_invoker: Optional callable ``(endpoint_name, df) -> df``
            used to invoke endpoint nodes. Defaults to using the full
            Workbench ``Endpoint`` API class — appropriate for client-side
            use. Pass a ``fast_inference``-backed invoker when running
            inside a deployed SageMaker container where the full
            Workbench config isn't available.

    Returns:
        The DataFrame at the DAG's output node, with the synthetic
        :data:`DAG_ROW_ID` column removed.
    """
    if self._output_node is None:
        raise ValueError("DAG has no output node — call set_output_node() first")
    if DAG_ROW_ID in input_df.columns:
        raise ValueError(
            f"input_df already contains the reserved column '{DAG_ROW_ID}'. " f"Remove it before calling run()."
        )

    # Inject the synthetic row id. Endpoints will pass this through as an
    # unknown input column; aggregation nodes use it as their join key.
    input_df = input_df.copy()
    input_df[DAG_ROW_ID] = range(len(input_df))

    def run_node(node: str) -> pd.DataFrame:
        if node in self._endpoints:
            return self._run_endpoint(node, input_df, outputs, endpoint_invoker)
        return self._run_aggregation(node, outputs)

    # Workers only read finished parents; `outputs` is written between waves.
    outputs: Dict[str, pd.DataFrame] = {}
    pending = self.topological_order()
    with ThreadPoolExecutor(max_workers=max(1, len(self._endpoints))) as pool:
        while pending:
            wave = [n for n in pending if all(p in outputs for p in self._parents_of(n))]
            futures = {node: pool.submit(run_node, node) for node in wave}
            for node, future in futures.items():
                outputs[node] = future.result()
            pending = [n for n in pending if n not in outputs]

    result = outputs[self._output_node]
    if DAG_ROW_ID in result.columns:
        result = result.drop(columns=[DAG_ROW_ID])
    return result

set_input_node(*nodes)

Declare which nodes receive the DAG's input DataFrame directly.

Source code in src/workbench/utils/meta_endpoint_dag.py 
202
203
204
205
206
207
def set_input_node(self, *nodes: str) -> None:
    """Declare which nodes receive the DAG's input DataFrame directly."""
    for n in nodes:
        if n not in self._endpoints:
            raise ValueError(f"Input nodes must be endpoint nodes; '{n}' is not")
    self._input_nodes = list(nodes)

set_output_node(node)

Declare the terminal node whose output is the DAG's output.

Source code in src/workbench/utils/meta_endpoint_dag.py 
209
210
211
212
213
def set_output_node(self, node: str) -> None:
    """Declare the terminal node whose output is the DAG's output."""
    if node not in self._all_nodes():
        raise ValueError(f"Unknown output node '{node}'")
    self._output_node = node

terminal_target()

Target column(s) the DAG ultimately predicts, or None for feature pipelines.

  • Output is an endpoint → its targets.
  • Output is an aggregation → walk back to the closest endpoint(s) and union their targets in branch order.

A node's targets are the ones declared on :meth:add_endpoint, else every head its model predicts. Returns a bare string for a single target, a list for a panel, and None when nothing upstream predicts anything (a feature pipeline).

Used by :meth:MetaEndpoint._derive_lineage to anchor the meta's target_column on what the DAG actually predicts rather than inheriting from the (possibly target-less) input endpoint.

Raises:

Type Description
ValueError

Two upstream nodes claim the same target. Declare targets= on :meth:add_endpoint to say which one owns it.
Source code in src/workbench/utils/meta_endpoint_dag.py 
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
def terminal_target(self) -> Optional[Union[str, List[str]]]:
    """Target column(s) the DAG ultimately predicts, or ``None`` for feature pipelines.

    - Output is an endpoint → its targets.
    - Output is an aggregation → walk back to the closest endpoint(s) and
      union their targets in branch order.

    A node's targets are the ones declared on :meth:`add_endpoint`, else
    every head its model predicts. Returns a bare string for a single
    target, a list for a panel, and ``None`` when nothing upstream predicts
    anything (a feature pipeline).

    Used by :meth:`MetaEndpoint._derive_lineage` to anchor the meta's
    ``target_column`` on what the DAG actually predicts rather than
    inheriting from the (possibly target-less) input endpoint.

    Raises:
        ValueError: Two upstream nodes claim the same target. Declare
            ``targets=`` on :meth:`add_endpoint` to say which one owns it.
    """
    from workbench.api import Endpoint, Model

    def _targets_of(node: str) -> List[str]:
        if node in self._primary or node in self._declared_targets:
            return self._owned_heads(node)
        ep = Endpoint(self._endpoints[node])
        if not ep.exists():
            return []
        target = Model(ep.get_input()).target()
        if not target:
            return []
        return list(target) if isinstance(target, list) else [target]

    if self._output_node in self._endpoints:
        predictors = [self._output_node]
    else:
        # Aggregation output — BFS back until we hit endpoints, in branch order.
        predictors = []
        seen: set = set()
        queue: List[str] = list(self._parents_of(self._output_node))
        while queue:
            node = queue.pop(0)
            if node in seen:
                continue
            seen.add(node)
            if node in self._endpoints:
                predictors.append(node)
            else:
                queue.extend(self._parents_of(node))

    targets: List[str] = []
    owner: Dict[str, str] = {}
    for node in predictors:
        for target in _targets_of(node):
            if target in owner:
                raise ValueError(
                    f"Nodes '{owner[target]}' and '{node}' both predict '{target}'. Declare "
                    f"targets= on add_endpoint() to say which node owns it."
                )
            owner[target] = node
            targets.append(target)

    if not targets:
        return None
    return targets[0] if len(targets) == 1 else targets

to_dict()

Serialize the DAG topology to a JSON-friendly dict.

Aggregation nodes are serialized by class name + constructor kwargs; deserialization (:meth:from_dict) requires the same class to be importable.

Per-endpoint is_async flags are included only if :meth:populate_child_metadata has been called. The deployed inference container relies on these flags to dispatch invocations to fast_inference or async_inference.

Source code in src/workbench/utils/meta_endpoint_dag.py 
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
def to_dict(self) -> dict:
    """Serialize the DAG topology to a JSON-friendly dict.

    Aggregation nodes are serialized by class name + constructor kwargs;
    deserialization (:meth:`from_dict`) requires the same class to be
    importable.

    Per-endpoint ``is_async`` flags are included only if
    :meth:`populate_child_metadata` has been called. The deployed
    inference container relies on these flags to dispatch invocations
    to ``fast_inference`` or ``async_inference``.
    """
    return {
        "endpoints": dict(self._endpoints),
        "endpoint_targets": {node: list(t) for node, t in self._declared_targets.items()},
        "endpoint_primary": dict(self._primary),
        "endpoint_async": dict(self._endpoint_async_flags),
        "endpoint_batch_size": dict(self._endpoint_batch_sizes),
        "aggregations": [_serialize_aggregation(a) for a in self._aggregations.values()],
        "edges": [list(e) for e in self._edges],
        "input_nodes": list(self._input_nodes),
        "output_node": self._output_node,
    }

topological_order()

Return nodes in topological order (parents before children).

Raises:

Type Description
ValueError

If the DAG contains a cycle.
Source code in src/workbench/utils/meta_endpoint_dag.py 
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
def topological_order(self) -> List[str]:
    """Return nodes in topological order (parents before children).

    Raises:
        ValueError: If the DAG contains a cycle.
    """
    in_degree = {n: 0 for n in self._all_nodes()}
    for _, dst in self._edges:
        in_degree[dst] += 1

    ready = [n for n, deg in in_degree.items() if deg == 0]
    order: List[str] = []
    while ready:
        node = ready.pop(0)
        order.append(node)
        for src, dst in self._edges:
            if src == node:
                in_degree[dst] -= 1
                if in_degree[dst] == 0:
                    ready.append(dst)

    if len(order) != len(in_degree):
        raise ValueError("DAG contains a cycle")
    return order

validate()

Validate the DAG. Returns self for chaining; raises on failure.

Checks
  • At least one input node and exactly one output node declared
  • No cycles
  • Aggregation nodes have at least one parent
  • Endpoint nodes are either input nodes (zero parents) or have exactly one upstream parent — never both
  • The output node is reachable from the input nodes
  • No two nodes declare the same target head

Topology only — no AWS round-trips. A declared head colliding with an undeclared child's heads needs those children's targets, so it surfaces in :meth:terminal_target instead.

Source code in src/workbench/utils/meta_endpoint_dag.py 
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
def validate(self) -> "MetaEndpointDAG":
    """Validate the DAG. Returns self for chaining; raises on failure.

    Checks:
      - At least one input node and exactly one output node declared
      - No cycles
      - Aggregation nodes have at least one parent
      - Endpoint nodes are either input nodes (zero parents) or have
        exactly one upstream parent — never both
      - The output node is reachable from the input nodes
      - No two nodes declare the same target head

    Topology only — no AWS round-trips. A declared head colliding with an
    *undeclared* child's heads needs those children's targets, so it
    surfaces in :meth:`terminal_target` instead.
    """
    if not self._input_nodes:
        raise ValueError("DAG has no input nodes")
    if self._output_node is None:
        raise ValueError("DAG has no output node")

    order = self.topological_order()  # raises on cycle

    for ep_node in self._endpoints:
        parents = self._parents_of(ep_node)
        is_input = ep_node in self._input_nodes
        if is_input and parents:
            raise ValueError(
                f"Endpoint node '{ep_node}' is declared as an input node but has "
                f"upstream parents {parents}; pick one or the other."
            )
        if not is_input and not parents:
            raise ValueError(
                f"Endpoint node '{ep_node}' has no upstream parent and is not "
                f"declared as an input node — it has no source for its input DataFrame."
            )

    for name in self._aggregations:
        if not self._parents_of(name):
            raise ValueError(f"Aggregation node '{name}' has no upstream parents")

    owner: Dict[str, str] = {}
    for node in list(self._declared_targets) + list(self._primary):
        for target in self._owned_heads(node):
            if target in owner:
                raise ValueError(
                    f"Nodes '{owner[target]}' and '{node}' both declare target '{target}' — "
                    f"a target can be owned by only one node."
                )
            owner[target] = node

    reachable = set(self._input_nodes)
    for node in order:
        if node in reachable:
            for src, dst in self._edges:
                if src == node:
                    reachable.add(dst)
    if self._output_node not in reachable:
        raise ValueError(f"Output node '{self._output_node}' is not reachable from input nodes {self._input_nodes}")

    return self

select_primary_columns(columns, name, context)

Map the unprefixed primary head to its two panel columns.

prediction -> <name>_pred and confidence -> <name>_conf; the rest of the unprefixed family (prediction_std, expected_residual, q_*) is dropped. Returns {source_column: panel_column}, pred first.

Raises:

Type Description
ValueError

columns is non-empty but has no prediction — the endpoint doesn't expose a primary head.
Source code in src/workbench/utils/meta_endpoint_dag.py 
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def select_primary_columns(columns: List[str], name: str, context: str) -> Dict[str, str]:
    """Map the unprefixed primary head to its two panel columns.

    ``prediction`` -> ``<name>_pred`` and ``confidence`` -> ``<name>_conf``; the
    rest of the unprefixed family (``prediction_std``, ``expected_residual``,
    ``q_*``) is dropped. Returns ``{source_column: panel_column}``, pred first.

    Raises:
        ValueError: ``columns`` is non-empty but has no ``prediction`` — the
            endpoint doesn't expose a primary head.
    """
    stems = {"prediction": "pred", "confidence": "conf"}
    mapping = {c: f"{name}_{stem}" for c, stem in stems.items() if c in columns}
    if columns and "prediction" not in mapping:
        raise ValueError(f"{context}: no 'prediction' column among {columns}")
    return mapping

select_target_columns(columns, targets, context)

Keep only the columns belonging to targets (<target> or <target>_*).

The single definition of what a declared head owns — used both for the static column contract (:meth:MetaEndpointDAG.output_columns) and for trimming an endpoint's actual output at run time.

Raises:

Type Description
ValueError

columns is non-empty but nothing matches — a declared target that owns no columns is a typo, not an empty branch.
Source code in src/workbench/utils/meta_endpoint_dag.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def select_target_columns(columns: List[str], targets: List[str], context: str) -> List[str]:
    """Keep only the columns belonging to ``targets`` (``<target>`` or ``<target>_*``).

    The single definition of what a declared head owns — used both for the
    static column contract (:meth:`MetaEndpointDAG.output_columns`) and for
    trimming an endpoint's actual output at run time.

    Raises:
        ValueError: ``columns`` is non-empty but nothing matches — a declared
            target that owns no columns is a typo, not an empty branch.
    """
    kept = [c for c in columns if any(c == t or c.startswith(f"{t}_") for t in targets)]
    if columns and not kept:
        raise ValueError(f"{context}: declared targets {targets} match none of {columns}")
    return kept

Aggregation nodes for MetaEndpointDAG.

An aggregation node combines outputs from one or more upstream nodes (Endpoint or other AggregationNode instances) into a single DataFrame. Two broad categories:

  • Column-union aggregators (Concat): join feature outputs from parallel feature endpoints into a single wide row per id — used for feature-pipeline DAGs (e.g. [2D] + [3D] → Concat).

  • Prediction aggregators (Mean, WeightedMean, Vote, plus the ensemble-strategy ports ConfidenceWeighted, InverseMaeWeighted, ScaledConfidenceWeighted, CalibratedConfidenceWeighted): combine prediction columns from multiple predictor endpoints into a single ensemble prediction with confidence — used for ensemble combination.

Each node declares its input/output column contract so the DAG can be validated statically before any inference runs.

AggregationNode

Base class for DAG aggregation nodes.

Subclasses implement apply() to combine upstream DataFrames and declare output_columns() for static DAG validation.

Aggregation nodes always join across upstream branches using the walker-injected :data:DAG_ROW_ID column, so they don't need to know anything about the caller's id conventions (or whether the caller has any).

Source code in src/workbench/utils/aggregation_nodes.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class AggregationNode:
    """Base class for DAG aggregation nodes.

    Subclasses implement ``apply()`` to combine upstream DataFrames and
    declare ``output_columns()`` for static DAG validation.

    Aggregation nodes always join across upstream branches using the
    walker-injected :data:`DAG_ROW_ID` column, so they don't need to
    know anything about the caller's id conventions (or whether the
    caller has any).
    """

    def __init__(self, name: str):
        self.name = name

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        """Combine upstream DataFrames into one. Subclasses must override."""
        raise NotImplementedError

    def input_columns(self, upstream_outputs: List[List[str]]) -> List[str]:
        """The columns this node expects across all upstream outputs.

        Default: union of all upstream output columns. Subclasses can
        narrow this if they only consume specific columns.
        """
        seen = set()
        cols: List[str] = []
        for upstream in upstream_outputs:
            for c in upstream:
                if c not in seen:
                    seen.add(c)
                    cols.append(c)
        return cols

    def output_columns(self, upstream_outputs: List[List[str]]) -> List[str]:
        """The columns this node emits. Subclasses must override."""
        raise NotImplementedError

apply(upstream)

Combine upstream DataFrames into one. Subclasses must override.

Source code in src/workbench/utils/aggregation_nodes.py 
56
57
58
def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
    """Combine upstream DataFrames into one. Subclasses must override."""
    raise NotImplementedError

input_columns(upstream_outputs)

The columns this node expects across all upstream outputs.

Default: union of all upstream output columns. Subclasses can narrow this if they only consume specific columns.

Source code in src/workbench/utils/aggregation_nodes.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def input_columns(self, upstream_outputs: List[List[str]]) -> List[str]:
    """The columns this node expects across all upstream outputs.

    Default: union of all upstream output columns. Subclasses can
    narrow this if they only consume specific columns.
    """
    seen = set()
    cols: List[str] = []
    for upstream in upstream_outputs:
        for c in upstream:
            if c not in seen:
                seen.add(c)
                cols.append(c)
    return cols

output_columns(upstream_outputs)

The columns this node emits. Subclasses must override.

Source code in src/workbench/utils/aggregation_nodes.py 
75
76
77
def output_columns(self, upstream_outputs: List[List[str]]) -> List[str]:
    """The columns this node emits. Subclasses must override."""
    raise NotImplementedError

CalibratedConfidenceWeighted

Bases: _StrategyAggregator

Per-row weights = confidence × |conf-error correlation| (normalized).

Rewards models whose confidence actually predicts accuracy.

Source code in src/workbench/utils/aggregation_nodes.py 
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
class CalibratedConfidenceWeighted(_StrategyAggregator):
    """Per-row weights = ``confidence × |conf-error correlation|`` (normalized).

    Rewards models whose confidence actually predicts accuracy.
    """

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        self._check_arity(upstream)
        ids, preds, confs = self._stack(upstream)
        calibrated = confs * self.corr_scale
        weights = conf_weights_with_fallback(calibrated, self.model_weights)
        return self._build_output(
            upstream,
            ids,
            prediction=(preds * weights).sum(axis=1),
            prediction_std=preds.std(axis=1),
            confidence=ensemble_confidence(preds, confs, self.corr_scale, self.model_weights, self.optimal_alpha),
        )

Concat

Bases: AggregationNode

Column-union aggregator. Joins upstream DataFrames on the walker's synthetic row id.

Use for feature-pipeline DAGs where parallel feature endpoints contribute disjoint feature column sets that need to be merged into a single wide row.

Source code in src/workbench/utils/aggregation_nodes.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
class Concat(AggregationNode):
    """Column-union aggregator. Joins upstream DataFrames on the walker's
    synthetic row id.

    Use for feature-pipeline DAGs where parallel feature endpoints
    contribute disjoint feature column sets that need to be merged into a
    single wide row.
    """

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        if not upstream:
            raise ValueError(f"Concat[{self.name}]: requires at least one upstream DataFrame")

        out = upstream[0]
        for df in upstream[1:]:
            new_cols = [c for c in df.columns if c == DAG_ROW_ID or c not in out.columns]
            out = out.merge(df[new_cols], on=DAG_ROW_ID, how="inner")
        return out

    def output_columns(self, upstream_outputs: List[List[str]]) -> List[str]:
        seen = set()
        cols: List[str] = []
        for upstream in upstream_outputs:
            for c in upstream:
                if c not in seen:
                    seen.add(c)
                    cols.append(c)
        return cols

ConfidenceWeighted

Bases: _StrategyAggregator

Per-row weights = upstream confidences (normalized).

Falls back to static model_weights when row confidences sum to ~0.

Source code in src/workbench/utils/aggregation_nodes.py 
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
class ConfidenceWeighted(_StrategyAggregator):
    """Per-row weights = upstream confidences (normalized).

    Falls back to static ``model_weights`` when row confidences sum to ~0.
    """

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        self._check_arity(upstream)
        ids, preds, confs = self._stack(upstream)
        weights = conf_weights_with_fallback(confs, self.model_weights)
        return self._build_output(
            upstream,
            ids,
            prediction=(preds * weights).sum(axis=1),
            prediction_std=preds.std(axis=1),
            confidence=ensemble_confidence(preds, confs, self.corr_scale, self.model_weights, self.optimal_alpha),
        )

InverseMaeWeighted

Bases: _StrategyAggregator

Static per-model weights from inverse-MAE.

The caller passes the inverse-MAE-derived weights directly via model_weights. Identical to WeightedMean for the prediction column, but additionally computes calibrated ensemble confidence.

Source code in src/workbench/utils/aggregation_nodes.py 
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
class InverseMaeWeighted(_StrategyAggregator):
    """Static per-model weights from inverse-MAE.

    The caller passes the inverse-MAE-derived weights directly via
    ``model_weights``. Identical to ``WeightedMean`` for the prediction
    column, but additionally computes calibrated ensemble confidence.
    """

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        self._check_arity(upstream)
        ids, preds, confs = self._stack(upstream)
        return self._build_output(
            upstream,
            ids,
            prediction=(preds * self.model_weights).sum(axis=1),
            prediction_std=preds.std(axis=1),
            confidence=ensemble_confidence(preds, confs, self.corr_scale, self.model_weights, self.optimal_alpha),
        )

Mean

Bases: _PredictionAggregator

Simple equal-weight mean of predictions.

Source code in src/workbench/utils/aggregation_nodes.py 
187
188
189
190
191
192
193
194
195
196
197
198
class Mean(_PredictionAggregator):
    """Simple equal-weight mean of predictions."""

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        ids, preds, confs = self._stack(upstream)
        return self._build_output(
            upstream,
            ids,
            prediction=preds.mean(axis=1),
            prediction_std=preds.std(axis=1),
            confidence=confs.mean(axis=1),
        )

ScaledConfidenceWeighted

Bases: _StrategyAggregator

Per-row weights = model_weights × confidence (normalized).

Often the top performer in practice — combines static MAE-derived weighting with per-row confidence scaling.

Source code in src/workbench/utils/aggregation_nodes.py 
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
class ScaledConfidenceWeighted(_StrategyAggregator):
    """Per-row weights = ``model_weights × confidence`` (normalized).

    Often the top performer in practice — combines static MAE-derived
    weighting with per-row confidence scaling.
    """

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        self._check_arity(upstream)
        ids, preds, confs = self._stack(upstream)
        scaled = confs * self.model_weights
        weights = conf_weights_with_fallback(scaled, self.model_weights)
        return self._build_output(
            upstream,
            ids,
            prediction=(preds * weights).sum(axis=1),
            prediction_std=preds.std(axis=1),
            confidence=ensemble_confidence(preds, confs, self.corr_scale, self.model_weights, self.optimal_alpha),
        )

Vote

Bases: _PredictionAggregator

Majority-vote aggregator for classifier predictions.

Expects each upstream's prediction column to hold class labels (string or int). Output prediction is the most common label per row; prediction_std is 0 (placeholder for contract symmetry); confidence is the fraction of upstream models that voted for the winning label.

Source code in src/workbench/utils/aggregation_nodes.py 
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
class Vote(_PredictionAggregator):
    """Majority-vote aggregator for classifier predictions.

    Expects each upstream's ``prediction`` column to hold class labels
    (string or int). Output ``prediction`` is the most common label per
    row; ``prediction_std`` is 0 (placeholder for contract symmetry);
    ``confidence`` is the fraction of upstream models that voted for the
    winning label.
    """

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        if not upstream:
            raise ValueError(f"Vote[{self.name}]: requires at least one upstream DataFrame")

        ids = upstream[0][[DAG_ROW_ID]].copy()
        for df in upstream[1:]:
            ids = ids.merge(df[[DAG_ROW_ID]], on=DAG_ROW_ID, how="inner")

        labels = pd.concat(
            [
                ids.merge(df[[DAG_ROW_ID, "prediction"]], on=DAG_ROW_ID)["prediction"].rename(f"_p{i}")
                for i, df in enumerate(upstream)
            ],
            axis=1,
        )

        modes = labels.mode(axis=1)[0]
        winner_share = (labels.eq(modes, axis=0)).sum(axis=1) / labels.shape[1]

        return self._build_output(
            upstream,
            ids,
            prediction=modes.to_numpy(),
            prediction_std=0.0,
            confidence=winner_share.to_numpy(),
        )

WeightedMean

Bases: _PredictionAggregator

Static-weight mean — caller supplies one weight per upstream.

Source code in src/workbench/utils/aggregation_nodes.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
class WeightedMean(_PredictionAggregator):
    """Static-weight mean — caller supplies one weight per upstream."""

    def __init__(self, name: str, weights: List[float]):
        super().__init__(name)
        if not weights:
            raise ValueError("WeightedMean: weights must be a non-empty list")
        w = np.asarray(weights, dtype=np.float64)
        if (w < 0).any():
            raise ValueError("WeightedMean: weights must be non-negative")
        if w.sum() <= 0:
            raise ValueError("WeightedMean: at least one weight must be positive")
        self.weights = w / w.sum()

    def apply(self, upstream: List[pd.DataFrame]) -> pd.DataFrame:
        if len(upstream) != len(self.weights):
            raise ValueError(
                f"WeightedMean[{self.name}]: got {len(upstream)} upstream frames " f"but {len(self.weights)} weights"
            )
        ids, preds, confs = self._stack(upstream)
        return self._build_output(
            upstream,
            ids,
            prediction=(preds * self.weights).sum(axis=1),
            prediction_std=preds.std(axis=1),
            confidence=(confs * self.weights).sum(axis=1),
        )

Questions?

The SuperCowPowers team is happy to answer any questions you may have about AWS and Workbench.