bikes.jobs

 20class EvaluationsJob(base.Job):
 21    """Generate evaluations from a registered model and a dataset.
 22
 23    Parameters:
 24        run_config (services.MlflowService.RunConfig): mlflow run config.
 25        inputs (datasets.ReaderKind): reader for the inputs data.
 26        targets (datasets.ReaderKind): reader for the targets data.
 27        model_type (str): model type (e.g. "regressor", "classifier").
 28        alias_or_version (str | int): alias or version for the  model.
 29        metrics (metrics_.MetricsKind): metric list to compute.
 30        evaluators (list[str]): list of evaluators to use.
 31        evaluator_config (dict[str, dict[str, T.Any]]): options per evaluator name.
 32        thresholds (dict[str, metrics_.Threshold] | None): metric thresholds.
 33    """
 34
 35    KIND: T.Literal["EvaluationsJob"] = "EvaluationsJob"
 36
 37    # Run
 38    run_config: services.MlflowService.RunConfig = services.MlflowService.RunConfig(name="Evaluations")
 39    # Data
 40    inputs: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 41    targets: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 42    # Model
 43    model_type: str = "regressor"
 44    alias_or_version: str | int = "Champion"
 45    # Loader
 46    loader: registries.LoaderKind = pdt.Field(registries.CustomLoader(), discriminator="KIND")
 47    # Metrics
 48    metrics: metrics_.MetricsKind = [metrics_.SklearnMetric()]
 49    # Evaluators
 50    evaluators: list[str] = ["default"]
 51    # The ExplanationsJob owns explanations: the default evaluator's SHAP step cannot read
 52    # the non-numeric inputs anyway and only logs a warning.
 53    evaluator_config: dict[str, dict[str, T.Any]] = {"default": {"log_model_explainability": False}}
 54    # Thresholds
 55    thresholds: dict[str, metrics_.Threshold] = {"r2_score": metrics_.Threshold(threshold=0.5, greater_is_better=True)}
 56
 57    @T.override
 58    def run(self) -> base.Locals:
 59        # services
 60        # - logger
 61        logger = self.logger_service.logger()
 62        logger.info("With logger: {}", logger)
 63        # - mlflow
 64        client = self.mlflow_service.client()
 65        logger.info("With client: {}", services.redact_uri(client.tracking_uri))
 66        with self.mlflow_service.run_context(run_config=self.run_config) as run:
 67            logger.info("With run context: {}", run.info)
 68            # data
 69            # - inputs
 70            logger.info("Read inputs: {}", self.inputs)
 71            inputs_ = self.inputs.read()  # unchecked!
 72            inputs = schemas.InputsSchema.check(inputs_)
 73            logger.debug("- Inputs shape: {}", inputs.shape)
 74            # - targets
 75            logger.info("Read targets: {}", self.targets)
 76            targets_ = self.targets.read()  # unchecked!
 77            targets = schemas.TargetsSchema.check(targets_)
 78            logger.debug("- Targets shape: {}", targets.shape)
 79            # lineage
 80            # - inputs
 81            logger.info("Log lineage: inputs")
 82            inputs_lineage = self.inputs.lineage(data=inputs, name="inputs")
 83            mlflow.log_input(dataset=inputs_lineage, context=self.run_config.name)
 84            logger.debug("- Inputs lineage: {}", inputs_lineage.to_dict())
 85            # - targets
 86            logger.info("Log lineage: targets")
 87            targets_lineage = self.targets.lineage(data=targets, name="targets", targets=schemas.TargetsSchema.cnt)
 88            mlflow.log_input(dataset=targets_lineage, context=self.run_config.name)
 89            logger.debug("- Targets lineage: {}", targets_lineage.to_dict())
 90            # model
 91            logger.info("With model: {}", self.mlflow_service.registry_name)
 92            model_uri = registries.uri_for_model_alias_or_version(
 93                name=self.mlflow_service.registry_name,
 94                alias_or_version=self.alias_or_version,
 95            )
 96            logger.debug("- Model URI: {}", model_uri)
 97            # loader
 98            logger.info("Load model: {}", self.loader)
 99            model = self.loader.load(uri=model_uri)
100            logger.debug("- Model: {}", model)
101            # outputs
102            logger.info("Predict outputs: {}", len(inputs))
103            outputs = model.predict(inputs=inputs)  # checked
104            logger.debug("- Outputs shape: {}", outputs.shape)
105            # dataset
106            logger.info("Create dataset: inputs & targets & outputs")
107            # Preserve domain counts while preventing unsigned residuals in MLflow metrics.
108            dataset_ = pd.concat([inputs, targets, outputs], axis="columns").astype(
109                {schemas.TargetsSchema.cnt: "float64", schemas.OutputsSchema.prediction: "float64"},
110            )
111            dataset = mlflow.data.from_pandas(  # type: ignore[attr-defined]
112                df=dataset_,
113                name="evaluation",
114                targets=schemas.TargetsSchema.cnt,
115                predictions=schemas.OutputsSchema.prediction,
116            )
117            logger.debug("- Dataset: {}", dataset.to_dict())
118            # metrics
119            logger.debug("Convert metrics: {}", self.metrics)
120            extra_metrics = [metric.to_mlflow() for metric in self.metrics]
121            logger.debug("- Extra metrics: {}", extra_metrics)
122            # thresholds
123            logger.info("Convert thresholds: {}", self.thresholds)
124            validation_thresholds = {name: threshold.to_mlflow() for name, threshold in self.thresholds.items()}
125            logger.debug("- Validation thresholds: {}", validation_thresholds)
126            # evaluations
127            logger.info("Compute evaluations: {}", self.model_type)
128            evaluations = mlflow.models.evaluate(
129                data=dataset,
130                model_type=self.model_type,
131                evaluators=self.evaluators,
132                evaluator_config=self.evaluator_config,
133                extra_metrics=extra_metrics,
134            )
135            logger.debug("- Evaluations metrics: {}", evaluations.metrics)
136            # validation
137            # - MLflow 3 split threshold validation out of `evaluate` into a dedicated call
138            logger.info("Validate evaluations: {}", validation_thresholds)
139            mlflow.validate_evaluation_results(
140                candidate_result=evaluations,
141                baseline_result=None,
142                validation_thresholds=validation_thresholds,
143            )
144            # notify
145            self.alerts_service.notify(
146                title="Evaluations Job Finished",
147                message=f"Evaluation metrics: {evaluations.metrics}",
148            )
149        return locals()

17class ExplanationsJob(base.Job):
18    """Generate explanations from the model and a data sample.
19
20    Parameters:
21        inputs_samples (datasets.ReaderKind): reader for the samples data.
22        models_explanations (datasets.WriterKind): writer for models explanation.
23        samples_explanations (datasets.WriterKind): writer for samples explanation.
24        alias_or_version (str | int): alias or version for the  model.
25        loader (registries.CustomLoader): registry loader for the model.
26    """
27
28    KIND: T.Literal["ExplanationsJob"] = "ExplanationsJob"
29
30    # Samples
31    inputs_samples: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
32    # Explanations
33    models_explanations: datasets.WriterKind = pdt.Field(..., discriminator="KIND")
34    samples_explanations: datasets.WriterKind = pdt.Field(..., discriminator="KIND")
35    # Model
36    alias_or_version: str | int = "Champion"
37    # Loader
38    # only the project's custom models implement the explanation methods
39    loader: registries.CustomLoader = registries.CustomLoader()
40
41    @T.override
42    def run(self) -> base.Locals:
43        # services
44        logger = self.logger_service.logger()
45        logger.info("With logger: {}", logger)
46        # inputs
47        logger.info("Read samples: {}", self.inputs_samples)
48        inputs_samples = self.inputs_samples.read()  # unchecked!
49        inputs_samples = schemas.InputsSchema.check(inputs_samples)
50        logger.debug("- Inputs samples shape: {}", inputs_samples.shape)
51        # model
52        logger.info("With model: {}", self.mlflow_service.registry_name)
53        model_uri = registries.uri_for_model_alias_or_version(
54            name=self.mlflow_service.registry_name,
55            alias_or_version=self.alias_or_version,
56        )
57        logger.debug("- Model URI: {}", model_uri)
58        # loader
59        logger.info("Load model: {}", self.loader)
60        model = self.loader.load(uri=model_uri).model.unwrap_python_model().model
61        logger.debug("- Model: {}", model)
62        # explanations
63        # - models
64        logger.info("Explain model: {}", model)
65        models_explanations = model.explain_model()
66        logger.debug("- Models explanations shape: {}", models_explanations.shape)
67        # # - samples
68        logger.info("Explain samples: {}", len(inputs_samples))
69        samples_explanations = model.explain_samples(inputs=inputs_samples)
70        logger.debug("- Samples explanations shape: {}", samples_explanations.shape)
71        # write
72        # - model
73        logger.info("Write models explanations: {}", self.models_explanations)
74        self.models_explanations.write(data=models_explanations)
75        # - samples
76        logger.info("Write samples explanations: {}", self.samples_explanations)
77        self.samples_explanations.write(data=samples_explanations)
78        # notify
79        self.alerts_service.notify(
80            title="Explanations Job Finished",
81            message=f"Features Count: {len(models_explanations)}",
82        )
83        return locals()

17class InferenceJob(base.Job):
18    """Generate batch predictions from a registered model.
19
20    Parameters:
21        inputs (datasets.ReaderKind): reader for the inputs data.
22        outputs (datasets.WriterKind): writer for the outputs data.
23        alias_or_version (str | int): alias or version for the  model.
24        loader (registries.LoaderKind): registry loader for the model.
25    """
26
27    KIND: T.Literal["InferenceJob"] = "InferenceJob"
28
29    # Inputs
30    inputs: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
31    # Outputs
32    outputs: datasets.WriterKind = pdt.Field(..., discriminator="KIND")
33    # Model
34    alias_or_version: str | int = "Champion"
35    # Loader
36    loader: registries.LoaderKind = pdt.Field(registries.CustomLoader(), discriminator="KIND")
37
38    @T.override
39    def run(self) -> base.Locals:
40        # services
41        logger = self.logger_service.logger()
42        logger.info("With logger: {}", logger)
43        # inputs
44        logger.info("Read inputs: {}", self.inputs)
45        inputs_ = self.inputs.read()  # unchecked!
46        inputs = schemas.InputsSchema.check(inputs_)
47        logger.debug("- Inputs shape: {}", inputs.shape)
48        # model
49        logger.info("With model: {}", self.mlflow_service.registry_name)
50        model_uri = registries.uri_for_model_alias_or_version(
51            name=self.mlflow_service.registry_name,
52            alias_or_version=self.alias_or_version,
53        )
54        logger.debug("- Model URI: {}", model_uri)
55        # loader
56        logger.info("Load model: {}", self.loader)
57        model = self.loader.load(uri=model_uri)
58        logger.debug("- Model: {}", model)
59        # outputs
60        logger.info("Predict outputs: {}", len(inputs))
61        outputs = model.predict(inputs=inputs)  # checked
62        logger.debug("- Outputs shape: {}", outputs.shape)
63        # write
64        logger.info("Write outputs: {}", self.outputs)
65        self.outputs.write(data=outputs)
66        # notify
67        self.alerts_service.notify(title="Inference Job Finished", message=f"Outputs Shape: {outputs.shape}")
68        return locals()

15class PipelineJob(base.Job):
16    """Tune, train, evaluate and promote one candidate, then run optional consumers.
17
18    Child jobs share this job's services. Evaluation failures propagate before the
19    alias changes. All consumers use the immutable version returned by training.
20    """
21
22    KIND: T.Literal["PipelineJob"] = "PipelineJob"
23    tuning: TuningJob
24    training: TrainingJob
25    evaluations: EvaluationsJob
26    alias: str = "Champion"
27    inference: InferenceJob | None = None
28    explanations: ExplanationsJob | None = None
29
30    def _configured(self, job: base.Job, **updates: T.Any) -> base.Job:
31        """Validate a stage with shared services and explicit candidate overrides."""
32        config = (
33            job.model_dump()
34            | {
35                "logger_service": self.logger_service,
36                "alerts_service": self.alerts_service,
37                "mlflow_service": self.mlflow_service,
38            }
39            | updates
40        )
41        return type(job).model_validate(config)
42
43    @T.override
44    def run(self) -> base.Locals:
45        """Run every stage; failed validation leaves the current alias unchanged."""
46        tuning = self._configured(self.tuning).run()
47        selected = type(self.tuning.model).model_validate(self.tuning.model.model_dump() | tuning["best_params"])
48        run_config = self.training.run_config.model_dump()
49        run_config["tags"] = (run_config.get("tags") or {}) | {
50            "tuning_run_id": tuning["run"].info.run_id,
51        }
52        training = self._configured(self.training, model=selected, run_config=run_config).run()
53        version = int(training["model_version"].version)
54        evaluations = self._configured(self.evaluations, alias_or_version=version).run()
55        promotion = PromotionJob(
56            logger_service=self.logger_service,
57            alerts_service=self.alerts_service,
58            mlflow_service=self.mlflow_service,
59            version=version,
60            alias=self.alias,
61        ).run()
62        inference = self._configured(self.inference, alias_or_version=version).run() if self.inference else None
63        explanations = (
64            self._configured(self.explanations, alias_or_version=version).run() if self.explanations else None
65        )
66        return locals()

15class PromotionJob(base.Job):
16    """Define a job for promoting a registered model version with an alias.
17
18    https://mlflow.org/docs/latest/model-registry.html#concepts
19
20    Parameters:
21        alias (str): the mlflow alias to transition the registered model version.
22        version (int): positive, explicit model version to transition.
23    """
24
25    KIND: T.Literal["PromotionJob"] = "PromotionJob"
26
27    alias: str = "Champion"
28    version: int = pdt.Field(..., gt=0)
29
30    @T.override
31    def run(self) -> base.Locals:
32        # services
33        # - logger
34        logger = self.logger_service.logger()
35        logger.info("With logger: {}", logger)
36        # - mlflow
37        client = self.mlflow_service.client()
38        logger.info("With client: {}", client)
39        name = self.mlflow_service.registry_name
40        # version
41        version = self.version
42        logger.info("From version: {}", version)
43        # alias
44        logger.info("To alias: {}", self.alias)
45        # promote
46        logger.info("Promote model: {}", name)
47        client.set_registered_model_alias(name=name, alias=self.alias, version=version)
48        model_version = client.get_model_version_by_alias(name=name, alias=self.alias)
49        logger.debug("- Model version: {}", model_version)
50        # notify
51        self.alerts_service.notify(
52            title="Promotion Job Finished",
53            message=f"Version: {model_version.version} @ {self.alias}",
54        )
55        return locals()

 20class TrainingJob(base.Job):
 21    """Train and register a single AI/ML model.
 22
 23    Parameters:
 24        run_config (services.MlflowService.RunConfig): mlflow run config.
 25        inputs (datasets.ReaderKind): reader for the inputs data.
 26        targets (datasets.ReaderKind): reader for the targets data.
 27        model (models.ModelKind): machine learning model to train.
 28        metrics (metrics_.MetricsKind): metric list to compute.
 29        splitter (splitters.SplitterKind): data sets splitter.
 30        saver (registries.SaverKind): model saver.
 31        signer (signers.SignerKind): model signer.
 32        registry (registries.RegisterKind): model register.
 33    """
 34
 35    KIND: T.Literal["TrainingJob"] = "TrainingJob"
 36
 37    # Run
 38    run_config: services.MlflowService.RunConfig = services.MlflowService.RunConfig(name="Training")
 39    # Data
 40    inputs: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 41    targets: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 42    # Model
 43    model: models.ModelKind = pdt.Field(models.BaselineSklearnModel(), discriminator="KIND")
 44    # Metrics
 45    metrics: metrics_.MetricsKind = [metrics_.SklearnMetric()]
 46    # Splitter
 47    splitter: splitters.SplitterKind = pdt.Field(splitters.TrainTestSplitter(), discriminator="KIND")
 48    # Saver
 49    saver: registries.SaverKind = pdt.Field(registries.CustomSaver(), discriminator="KIND")
 50    # Signer
 51    signer: signers.SignerKind = pdt.Field(signers.InferSigner(), discriminator="KIND")
 52    # Registrer
 53    # - avoid shadowing pydantic `register` pydantic function
 54    registry: registries.RegisterKind = pdt.Field(registries.MlflowRegister(), discriminator="KIND")
 55
 56    @T.override
 57    def run(self) -> base.Locals:
 58        # services
 59        # - logger
 60        logger = self.logger_service.logger()
 61        logger.info("With logger: {}", logger)
 62        # - mlflow
 63        client = self.mlflow_service.client()
 64        logger.info("With client: {}", services.redact_uri(client.tracking_uri))
 65        with self.mlflow_service.run_context(run_config=self.run_config) as run:
 66            logger.info("With run context: {}", run.info)
 67            # data
 68            # - inputs
 69            logger.info("Read inputs: {}", self.inputs)
 70            inputs_ = self.inputs.read()  # unchecked!
 71            inputs = schemas.InputsSchema.check(inputs_)
 72            logger.debug("- Inputs shape: {}", inputs.shape)
 73            # - targets
 74            logger.info("Read targets: {}", self.targets)
 75            targets_ = self.targets.read()  # unchecked!
 76            targets = schemas.TargetsSchema.check(targets_)
 77            logger.debug("- Targets shape: {}", targets.shape)
 78            # lineage
 79            # - inputs
 80            logger.info("Log lineage: inputs")
 81            inputs_lineage = self.inputs.lineage(data=inputs, name="inputs")
 82            mlflow.log_input(dataset=inputs_lineage, context=self.run_config.name)
 83            logger.debug("- Inputs lineage: {}", inputs_lineage.to_dict())
 84            # - targets
 85            logger.info("Log lineage: targets")
 86            targets_lineage = self.targets.lineage(data=targets, name="targets", targets=schemas.TargetsSchema.cnt)
 87            mlflow.log_input(dataset=targets_lineage, context=self.run_config.name)
 88            logger.debug("- Targets lineage: {}", targets_lineage.to_dict())
 89            # splitter
 90            logger.info("With splitter: {}", self.splitter)
 91            # - index
 92            train_index, test_index = next(self.splitter.split(inputs=inputs, targets=targets))
 93            # - inputs
 94            inputs_train = T.cast(schemas.Inputs, inputs.iloc[train_index])
 95            inputs_test = T.cast(schemas.Inputs, inputs.iloc[test_index])
 96            logger.debug("- Inputs train shape: {}", inputs_train.shape)
 97            logger.debug("- Inputs test shape: {}", inputs_test.shape)
 98            # - targets
 99            targets_train = T.cast(schemas.Targets, targets.iloc[train_index])
100            targets_test = T.cast(schemas.Targets, targets.iloc[test_index])
101            logger.debug("- Targets train shape: {}", targets_train.shape)
102            logger.debug("- Targets test shape: {}", targets_test.shape)
103            # model
104            logger.info("Fit model: {}", self.model)
105            self.model.fit(inputs=inputs_train, targets=targets_train)
106            # outputs
107            logger.info("Predict outputs: {}", len(inputs_test))
108            outputs_test = self.model.predict(inputs=inputs_test)
109            logger.debug("- Outputs test shape: {}", outputs_test.shape)
110            # metrics
111            for i, metric in enumerate(self.metrics, start=1):
112                logger.info("{}. Compute metric: {}", i, metric)
113                score = metric.score(targets=targets_test, outputs=outputs_test)
114                client.log_metric(run_id=run.info.run_id, key=metric.name, value=score)
115                logger.debug("- Metric score: {}", score)
116            # signer
117            logger.info("Sign model: {}", self.signer)
118            model_signature = self.signer.sign(inputs=inputs, outputs=outputs_test)
119            logger.debug("- Model signature: {}", model_signature.to_dict())
120            # saver
121            logger.info("Save model: {}", self.saver)
122            model_info = self.saver.save(model=self.model, signature=model_signature, input_example=inputs.head())
123            logger.debug("- Model URI: {}", model_info.model_uri)
124            # register
125            logger.info("Register model: {}", self.registry)
126            model_version = self.registry.register(
127                name=self.mlflow_service.registry_name, model_uri=model_info.model_uri
128            )
129            logger.debug("- Model version: {}", model_version)
130            # notify
131            self.alerts_service.notify(
132                title="Training Job Finished",
133                message=f"Model version: {model_version.version}",
134            )
135        return locals()