sklearn metadata routing: let `treatment` reach a learner inside a Pipeline
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 6k
- Forks
- 877
- PR merge metrics
- No merged PRs in 30d
Description
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.BaseLearnersubclassesBaseEstimatorand thefitsignatures declaretreatmentandpexplicitly, soset_fit_requestshould be generated automatically; confirm that againstsklearn.utils.metadata_routingbefore 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,treatmentandptogether 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
Pipelinefinal step withtreatmentrouted, no adapter class. - #854's original reproduction works.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with sklearn.utils.metadata_routing and Pipeline, then inspect BaseLearner and the meta-learner fit signatures after #985 lands. Add integration coverage for an S-learner and a propensity-consuming learner with enabled and prefixed routing, plus the documented Pipeline example. Done means treatment reaches the final learner without an adapter and #854's reproduction works.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, scikit-learn
- Domain
- documentation, machine-learning, testing-qa
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100