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        thresholds (dict[str, metrics_.Threshold] | None): metric thresholds.
 32    """
 33
 34    KIND: T.Literal["EvaluationsJob"] = "EvaluationsJob"
 35
 36    # Run
 37    run_config: services.MlflowService.RunConfig = services.MlflowService.RunConfig(name="Evaluations")
 38    # Data
 39    inputs: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 40    targets: datasets.ReaderKind = pdt.Field(..., discriminator="KIND")
 41    # Model
 42    model_type: str = "regressor"
 43    alias_or_version: str | int = "Champion"
 44    # Loader
 45    loader: registries.LoaderKind = pdt.Field(registries.CustomLoader(), discriminator="KIND")
 46    # Metrics
 47    metrics: metrics_.MetricsKind = [metrics_.SklearnMetric()]
 48    # Evaluators
 49    evaluators: list[str] = ["default"]
 50    # Thresholds
 51    thresholds: dict[str, metrics_.Threshold] = {"r2_score": metrics_.Threshold(threshold=0.5, greater_is_better=True)}
 52
 53    @T.override
 54    def run(self) -> base.Locals:
 55        # services
 56        # - logger
 57        logger = self.logger_service.logger()
 58        logger.info("With logger: {}", logger)
 59        # - mlflow
 60        client = self.mlflow_service.client()
 61        logger.info("With client: {}", client.tracking_uri)
 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.mlflow_service.registry_name)
 88            model_uri = registries.uri_for_model_alias_or_version(
 89                name=self.mlflow_service.registry_name,
 90                alias_or_version=self.alias_or_version,
 91            )
 92            logger.debug("- Model URI: {}", model_uri)
 93            # loader
 94            logger.info("Load model: {}", self.loader)
 95            model = self.loader.load(uri=model_uri)
 96            logger.debug("- Model: {}", model)
 97            # outputs
 98            logger.info("Predict outputs: {}", len(inputs))
 99            outputs = model.predict(inputs=inputs)  # checked
100            logger.debug("- Outputs shape: {}", outputs.shape)
101            # dataset
102            logger.info("Create dataset: inputs & targets & outputs")
103            dataset_ = pd.concat([inputs, targets, outputs], axis="columns")
104            dataset = mlflow.data.from_pandas(  # type: ignore[attr-defined]
105                df=dataset_,
106                name="evaluation",
107                targets=schemas.TargetsSchema.cnt,
108                predictions=schemas.OutputsSchema.prediction,
109            )
110            logger.debug("- Dataset: {}", dataset.to_dict())
111            # metrics
112            logger.debug("Convert metrics: {}", self.metrics)
113            extra_metrics = [metric.to_mlflow() for metric in self.metrics]
114            logger.debug("- Extra metrics: {}", extra_metrics)
115            # thresholds
116            logger.info("Convert thresholds: {}", self.thresholds)
117            validation_thresholds = {name: threshold.to_mlflow() for name, threshold in self.thresholds.items()}
118            logger.debug("- Validation thresholds: {}", validation_thresholds)
119            # evaluations
120            logger.info("Compute evaluations: {}", self.model_type)
121            evaluations = mlflow.evaluate(
122                data=dataset,
123                model_type=self.model_type,
124                evaluators=self.evaluators,
125                extra_metrics=extra_metrics,
126            )
127            logger.debug("- Evaluations metrics: {}", evaluations.metrics)
128            # validation
129            # - MLflow 3 split threshold validation out of `evaluate` into a dedicated call
130            logger.info("Validate evaluations: {}", validation_thresholds)
131            mlflow.validate_evaluation_results(
132                candidate_result=evaluations,
133                baseline_result=None,
134                validation_thresholds=validation_thresholds,
135            )
136            # notify
137            self.alerts_service.notify(
138                title="Evaluations Job Finished",
139                message=f"Evaluation metrics: {evaluations.metrics}",
140            )
141        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.LoaderKind): 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    loader: registries.LoaderKind = pdt.Field(registries.CustomLoader(), discriminator="KIND")
39
40    @T.override
41    def run(self) -> base.Locals:
42        # services
43        logger = self.logger_service.logger()
44        logger.info("With logger: {}", logger)
45        # inputs
46        logger.info("Read samples: {}", self.inputs_samples)
47        inputs_samples = self.inputs_samples.read()  # unchecked!
48        inputs_samples = schemas.InputsSchema.check(inputs_samples)
49        logger.debug("- Inputs samples shape: {}", inputs_samples.shape)
50        # model
51        logger.info("With model: {}", self.mlflow_service.registry_name)
52        model_uri = registries.uri_for_model_alias_or_version(
53            name=self.mlflow_service.registry_name,
54            alias_or_version=self.alias_or_version,
55        )
56        logger.debug("- Model URI: {}", model_uri)
57        # loader
58        logger.info("Load model: {}", self.loader)
59        model = self.loader.load(uri=model_uri).model.unwrap_python_model().model
60        logger.debug("- Model: {}", model)
61        # explanations
62        # - models
63        logger.info("Explain model: {}", model)
64        models_explanations = model.explain_model()
65        logger.debug("- Models explanations shape: {}", models_explanations.shape)
66        # # - samples
67        logger.info("Explain samples: {}", len(inputs_samples))
68        samples_explanations = model.explain_samples(inputs=inputs_samples)
69        logger.debug("- Samples explanations shape: {}", samples_explanations.shape)
70        # write
71        # - model
72        logger.info("Write models explanations: {}", self.models_explanations)
73        self.models_explanations.write(data=models_explanations)
74        # - samples
75        logger.info("Write samples explanations: {}", self.samples_explanations)
76        self.samples_explanations.write(data=samples_explanations)
77        # notify
78        self.alerts_service.notify(
79            title="Explanations Job Finished",
80            message=f"Features Count: {len(models_explanations)}",
81        )
82        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            version = client.search_model_versions(f"name='{name}'", max_results=1, order_by=["version_number DESC"])[
41                0
42            ].version
43        else:
44            version = self.version
45        logger.info("From version: {}", version)
46        # alias
47        logger.info("To alias: {}", self.alias)
48        # promote
49        logger.info("Promote model: {}", name)
50        client.set_registered_model_alias(name=name, alias=self.alias, version=version)
51        model_version = client.get_model_version_by_alias(name=name, alias=self.alias)
52        logger.debug("- Model version: {}", model_version)
53        # notify
54        self.alerts_service.notify(
55            title="Promotion Job Finished",
56            message=f"Version: {model_version.version} @ {self.alias}",
57        )
58        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: {}", 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)
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()