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            dataset_ = pd.concat([inputs, targets, outputs], axis="columns")
108            dataset = mlflow.data.from_pandas(  # type: ignore[attr-defined]
109                df=dataset_,
110                name="evaluation",
111                targets=schemas.TargetsSchema.cnt,
112                predictions=schemas.OutputsSchema.prediction,
113            )
114            logger.debug("- Dataset: {}", dataset.to_dict())
115            # metrics
116            logger.debug("Convert metrics: {}", self.metrics)
117            extra_metrics = [metric.to_mlflow() for metric in self.metrics]
118            logger.debug("- Extra metrics: {}", extra_metrics)
119            # thresholds
120            logger.info("Convert thresholds: {}", self.thresholds)
121            validation_thresholds = {name: threshold.to_mlflow() for name, threshold in self.thresholds.items()}
122            logger.debug("- Validation thresholds: {}", validation_thresholds)
123            # evaluations
124            logger.info("Compute evaluations: {}", self.model_type)
125            evaluations = mlflow.models.evaluate(
126                data=dataset,
127                model_type=self.model_type,
128                evaluators=self.evaluators,
129                evaluator_config=self.evaluator_config,
130                extra_metrics=extra_metrics,
131            )
132            logger.debug("- Evaluations metrics: {}", evaluations.metrics)
133            # validation
134            # - MLflow 3 split threshold validation out of `evaluate` into a dedicated call
135            logger.info("Validate evaluations: {}", validation_thresholds)
136            mlflow.validate_evaluation_results(
137                candidate_result=evaluations,
138                baseline_result=None,
139                validation_thresholds=validation_thresholds,
140            )
141            # notify
142            self.alerts_service.notify(
143                title="Evaluations Job Finished",
144                message=f"Evaluation metrics: {evaluations.metrics}",
145            )
146        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()

13class PromotionJob(base.Job):
14    """Define a job for promoting a registered model version with an alias.
15
16    https://mlflow.org/docs/latest/model-registry.html#concepts
17
18    Parameters:
19        alias (str): the mlflow alias to transition the registered model version.
20        version (int | None): the model version to transition (use None for latest).
21    """
22
23    KIND: T.Literal["PromotionJob"] = "PromotionJob"
24
25    alias: str = "Champion"
26    version: int | None = None
27
28    @T.override
29    def run(self) -> base.Locals:
30        # services
31        # - logger
32        logger = self.logger_service.logger()
33        logger.info("With logger: {}", logger)
34        # - mlflow
35        client = self.mlflow_service.client()
36        logger.info("With client: {}", client)
37        name = self.mlflow_service.registry_name
38        # version
39        if self.version is None:  # use the latest model version
40            try:
41                version = client.search_model_versions(
42                    f"name='{name}'", max_results=1, order_by=["version_number DESC"]
43                )[0].version
44            except IndexError as error:
45                raise ValueError(f"No version registered for model: {name}") from error
46        else:
47            version = self.version
48        logger.info("From version: {}", version)
49        # alias
50        logger.info("To alias: {}", self.alias)
51        # promote
52        logger.info("Promote model: {}", name)
53        client.set_registered_model_alias(name=name, alias=self.alias, version=version)
54        model_version = client.get_model_version_by_alias(name=name, alias=self.alias)
55        logger.debug("- Model version: {}", model_version)
56        # notify
57        self.alerts_service.notify(
58            title="Promotion Job Finished",
59            message=f"Version: {model_version.version} @ {self.alias}",
60        )
61        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()

 19class TuningJob(base.Job):
 20    """Find the best hyperparameters for a model.
 21
 22    Parameters:
 23        run_config (services.MlflowService.RunConfig): mlflow run config.
 24        inputs (datasets.ReaderKind): reader for the inputs data.
 25        targets (datasets.ReaderKind): reader for the targets data.
 26        model (models.ModelKind): machine learning model to tune.
 27        metric (metrics.MetricKind): tuning metric to optimize.
 28        splitter (splitters.SplitterKind): data sets splitter.
 29        searcher: (searchers.SearcherKind): hparams searcher.
 30    """
 31
 32    KIND: T.Literal["TuningJob"] = "TuningJob"
 33
 34    # Run
 35    run_config: services.MlflowService.RunConfig = services.MlflowService.RunConfig(name="Tuning")
 36    # Data
 37    inputs: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 38    targets: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 39    # Model
 40    model: models.ModelKind = pdt.Field(models.BaselineSklearnModel(), discriminator="KIND")
 41    # Metric
 42    metric: metrics.MetricKind = pdt.Field(metrics.SklearnMetric(), discriminator="KIND")
 43    # splitter
 44    splitter: splitters.SplitterKind = pdt.Field(splitters.TimeSeriesSplitter(), discriminator="KIND")
 45    # Searcher
 46    searcher: searchers.SearcherKind = pdt.Field(
 47        searchers.GridCVSearcher(
 48            param_grid={
 49                "max_depth": [3, 5, 7],
 50            }
 51        ),
 52        discriminator="KIND",
 53    )
 54
 55    @T.override
 56    def run(self) -> base.Locals:
 57        """Run the tuning job in context."""
 58        # services
 59        # - logger
 60        logger = self.logger_service.logger()
 61        logger.info("With logger: {}", logger)
 62        with self.mlflow_service.run_context(run_config=self.run_config) as run:
 63            logger.info("With run context: {}", run.info)
 64            # data
 65            # - inputs
 66            logger.info("Read inputs: {}", self.inputs)
 67            inputs_ = self.inputs.read()  # unchecked!
 68            inputs = schemas.InputsSchema.check(inputs_)
 69            logger.debug("- Inputs shape: {}", inputs.shape)
 70            # - targets
 71            logger.info("Read targets: {}", self.targets)
 72            targets_ = self.targets.read()  # unchecked!
 73            targets = schemas.TargetsSchema.check(targets_)
 74            logger.debug("- Targets shape: {}", targets.shape)
 75            # lineage
 76            # - inputs
 77            logger.info("Log lineage: inputs")
 78            inputs_lineage = self.inputs.lineage(data=inputs, name="inputs")
 79            mlflow.log_input(dataset=inputs_lineage, context=self.run_config.name)
 80            logger.debug("- Inputs lineage: {}", inputs_lineage.to_dict())
 81            # - targets
 82            logger.info("Log lineage: targets")
 83            targets_lineage = self.targets.lineage(data=targets, name="targets", targets=schemas.TargetsSchema.cnt)
 84            mlflow.log_input(dataset=targets_lineage, context=self.run_config.name)
 85            logger.debug("- Targets lineage: {}", targets_lineage.to_dict())
 86            # model
 87            logger.info("With model: {}", self.model)
 88            # metric
 89            logger.info("With metric: {}", self.metric)
 90            # splitter
 91            logger.info("With splitter: {}", self.splitter)
 92            # searcher
 93            logger.info("Run searcher: {}", self.searcher)
 94            results, best_score, best_params = self.searcher.search(
 95                model=self.model,
 96                metric=self.metric,
 97                inputs=inputs,
 98                targets=targets,
 99                cv=self.splitter,
100            )
101            logger.debug("- Results: {}", results.shape)
102            logger.debug("- Best Score: {}", best_score)
103            logger.debug("- Best Params: {}", best_params)
104            # notify
105            self.alerts_service.notify(title="Tuning Job Finished", message=f"Best score: {best_score}")
106        return locals()