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.

Scope

• Implement metadata routing for treatment (and p, where learners accept propensity scores) on the meta-learners first — they are the classes users most want inside a Pipeline.
• Decide the interaction with sklearn.set_config(enable_metadata_routing=True), which is opt-in and off by default: does CausalML require it, or support both paths?
• Establish which transformers in a Pipeline are expected to pass treatment through untouched, and what happens when an intermediate step resamples rows.
• End-to-end test: a real Pipeline([...]) with a CausalML learner as the final step, fit and predicted without an adapter.

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.