Skip to content

[Enhancement] Model Serialization & Persistence for Trained Learners via save/load API #892

@aman-coder03

Description

@aman-coder03

Is your feature request related to a problem? Please describe.
causalML has no standard way to persist a trained learner. Once you call .fit(), the only way to reuse that model is to keep the Python process alive or re-train from scratch. This is a significant gap for production workflows where models are trained in one environment(e.g. a batch job) and served or evaluated in another. Currently users are forced to reach for pickle or joblib manually, with no guarantee that the serialized object is compatible with the installed version of causalML, no standardized file format, and no integration with experiment tracking tools like MLflow

Describe the solution you'd like
add SerializableLearner mixin that all base learners (BaseTLearner, BaseSLearner, BaseXLearner, BaseRLearner, BaseDRLearner) inherit from, exposing two methods:

pythonlearner.save("model.causalml")
learner = BaseTRegressor.load("model.causalml")

every saved file would embed __causalml_version__ and __learner_class__ metadata. On load, if the installed version does not match, a CausalMLVersionMismatchWarning is raised so users are never silently running a stale model. MLflow integration would be added as an optional thin wrapper via to_mlflow() and from_mlflow(), raising a helpful ImportError if MLflow is not installed

Describe alternatives you've considered
using raw joblib.dump / joblib.load directly works but puts the burden on the user to handle versioning, class validation, and path conventions themselves. There is no way to know whether a saved file was produced by the same version of CausalML, which makes sharing models across teams or environments risky. A standardized mixin solves this once at the library level rather than in every user's codebase

Additional context

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions