sklearn metadata routing: let `treatment` reach a learner inside a Pipeline

[jeongyoonlee]

Part of #980 (v1.0 M1). Closes #854 when it lands.

Reordering fit to (X, y, treatment, …) is necessary but not sufficient for #854. Pipeline.fit(X, y) passes exactly two things; treatment is a third required array with nowhere to ride. After the flip a CausalML learner still can't be a Pipeline step without an adapter — the reordering just removes the first obstacle.

The mechanism sklearn provides is metadata routing: fit(X, y, **fit_params) with set_fit_request(treatment=True), so a caller writes pipe.fit(X, y, treatment=treatment) and the router delivers treatment to the step that asked for it.

Prerequisite

#985, the signature flip. Until then Pipeline's positional y binds to treatment, so an end-to-end test cannot be written. #985 is unblocked and targets v1.0 (Jun 2027).

Scope for the first pass

• Meta-learners only — they are the classes users most want inside a Pipeline.
• Integration tests and documentation before any custom MetadataRouter. BaseLearner subclasses BaseEstimator and the fit signatures declare treatment and p explicitly, so set_fit_request should be generated automatically; confirm that against sklearn.utils.metadata_routing before writing any routing code.
• Cover an S-learner pipeline plus one propensity-consuming learner.
• Cover both paths: routing enabled via sklearn.set_config(enable_metadata_routing=True), and the prefixed routing-disabled form.
• Row-preserving transformers only. Resampling has to move X, y, treatment and p together and belongs in its own design.

Out of scope for the first pass: whether p should route into a learner's wrapped estimator, rather than being consumed by the learner itself.

Why it's separate

The argument order is a breaking change on a deadline (v1.0 API freeze); routing is additive and can land any time after. Coupling them would put a design-heavy feature on the critical path of a mechanical rename.

Acceptance

• A documented, tested example of a CausalML learner as a Pipeline final step with treatment routed, no adapter class.
• CausalML is not compatible with Scikit-Learn Pipeline #854's original reproduction works.

[mohsinm-dev]

This separation from the argument-order flip makes sense. One dependency may be worth making explicit: the end-to-end Pipeline test is blocked by #985. With the current fit(X, treatment, y, ...) signature, Pipeline’s positional y binds to treatment, while routed treatment= supplies it again.

After #985, the implementation may be quite small. BaseLearner already inherits BaseEstimator, and its fit signatures explicitly declare treatment and p. These learners consume that metadata rather than routing it to their underlying models, so scikit-learn should generate set_fit_request(...) automatically. I’d start with integration tests and documentation before introducing a custom MetadataRouter.

A clean initial boundary could cover an S-learner Pipeline, one propensity-consuming learner, both routing-enabled and prefixed routing-disabled usage, and row-preserving transformers only. Row resampling requires coordinated changes to X, y, treatment, and p, so it likely belongs in a separate design.

[jeongyoonlee]

@mohsinm-dev — agreed on all four boundaries: integration tests and docs before any custom MetadataRouter, an S-learner pipeline plus one propensity-consuming learner, both routing-enabled and prefixed routing-disabled usage, and row-preserving transformers only with resampling deferred to its own design.

Sequencing is as you describe. #985 is the prerequisite, and it is currently gated on the deprecation-window length — open question 2 on #938. The window-closing work is done (#981–#984 merged), so #985 is waiting only on that decision.

Two notes for when this starts:

• BaseLearner subclasses BaseEstimator and the fit signatures declare treatment and p explicitly, so set_fit_request should be generated automatically. Worth confirming against sklearn.utils.metadata_routing before assuming it.
• The learners consume this metadata rather than forwarding it to their inner models. Whether p should also route into a wrapped estimator is a separate question and probably out of scope for the first pass.