star

venv/lib/python3.9/site-packages/mlflow/recipes/steps/automl/flaml.py

@@ -0,0 +1,177 @@
+import importlib
+import logging
+from typing import TYPE_CHECKING, Any, Callable
+
+import pandas as pd
+
+if TYPE_CHECKING:
+    from sklearn.base import BaseEstimator
+
+import mlflow
+from mlflow import MlflowException
+from mlflow.models import EvaluationMetric
+from mlflow.models.evaluation.evaluators.classifier import _get_binary_classifier_metrics
+from mlflow.models.evaluation.evaluators.regressor import _get_regressor_metrics
+from mlflow.recipes.utils.metrics import RecipeMetric, _load_custom_metrics
+
+_logger = logging.getLogger(__name__)
+
+_AUTOML_DEFAULT_TIME_BUDGET = 600
+_MLFLOW_TO_FLAML_METRICS = {
+    "mean_absolute_error": "mae",
+    "mean_squared_error": "mse",
+    "root_mean_squared_error": "rmse",
+    "r2_score": "r2",
+    "mean_absolute_percentage_error": "mape",
+    "f1_score": "f1",
+    "f1_score_micro": "micro_f1",
+    "f1_score_macro": "macro_f1",
+    "accuracy_score": "accuracy",
+    "roc_auc": "roc_auc",
+    "roc_auc_ovr": "roc_auc_ovr",
+    "roc_auc_ovo": "roc_auc_ovo",
+    "log_loss": "log_loss",
+}
+
+# metrics that are not supported natively in FLAML
+_SKLEARN_METRICS = ["recall_score", "precision_score"]
+
+
+def get_estimator_and_best_params(
+    X,
+    y,
+    task: str,
+    extended_task: str,
+    step_config: dict[str, Any],
+    recipe_root: str,
+    evaluation_metrics: dict[str, RecipeMetric],
+    primary_metric: str,
+) -> tuple["BaseEstimator", dict[str, Any]]:
+    return _create_model_automl(
+        X, y, task, extended_task, step_config, recipe_root, evaluation_metrics, primary_metric
+    )
+
+
+def _create_custom_metric_flaml(
+    task: str, metric_name: str, coeff: int, eval_metric: EvaluationMetric
+) -> Callable:
+    def calc_metric(X, y, estimator) -> dict[str, float]:
+        y_pred = estimator.predict(X)
+        builtin_metrics = (
+            _get_regressor_metrics(y, y_pred, sample_weights=None)
+            if task == "regression"
+            else _get_binary_classifier_metrics(y_true=y, y_pred=y_pred)
+        )
+        res_df = pd.DataFrame()
+        res_df["prediction"] = y_pred
+        res_df["target"] = y if task == "classification" else y.values
+        return eval_metric.eval_fn(res_df, builtin_metrics)
+
+    def custom_metric(
+        X_val,
+        y_val,
+        estimator,
+        labels,
+        X_train,
+        y_train,
+        weight_val=None,
+        weight_train=None,
+        *args,
+    ):
+        val_metric = coeff * calc_metric(X_val, y_val, estimator)
+        train_metric = calc_metric(X_train, y_train, estimator)
+        main_metric = coeff * val_metric
+        return main_metric, {
+            f"{metric_name}_train": train_metric,
+            f"{metric_name}_val": val_metric,
+        }
+
+    return custom_metric
+
+
+def _create_sklearn_metric_flaml(metric_name: str, coeff: int, avg: str = "binary") -> Callable:
+    def sklearn_metric(
+        X_val,
+        y_val,
+        estimator,
+        labels,
+        X_train,
+        y_train,
+        weight_val=None,
+        weight_train=None,
+        *args,
+    ):
+        custom_metrics_mod = importlib.import_module("sklearn.metrics")
+        eval_fn = getattr(custom_metrics_mod, metric_name)
+        val_metric = coeff * eval_fn(y_val, estimator.predict(X_val), average=avg)
+        train_metric = coeff * eval_fn(y_train, estimator.predict(X_train), average=avg)
+        return val_metric, {
+            f"{metric_name}_train": train_metric,
+            f"{metric_name}_val": val_metric,
+        }
+
+    return sklearn_metric
+
+
+def _create_model_automl(
+    X,
+    y,
+    task: str,
+    extended_task: str,
+    step_config: dict[str, Any],
+    recipe_root: str,
+    evaluation_metrics: dict[str, RecipeMetric],
+    primary_metric: str,
+) -> tuple["BaseEstimator", dict[str, Any]]:
+    try:
+        from flaml import AutoML
+    except ImportError:
+        raise MlflowException("Please install FLAML to use AutoML!")
+
+    try:
+        if primary_metric in _MLFLOW_TO_FLAML_METRICS and primary_metric in evaluation_metrics:
+            metric = _MLFLOW_TO_FLAML_METRICS[primary_metric]
+            if primary_metric == "roc_auc" and extended_task == "classification/multiclass":
+                metric = "roc_auc_ovr"
+        elif primary_metric in _SKLEARN_METRICS and primary_metric in evaluation_metrics:
+            metric = _create_sklearn_metric_flaml(
+                primary_metric,
+                -1 if evaluation_metrics[primary_metric].greater_is_better else 1,
+                "macro" if extended_task in ["classification/multiclass"] else "binary",
+            )
+        elif primary_metric in evaluation_metrics:
+            metric = _create_custom_metric_flaml(
+                task,
+                primary_metric,
+                -1 if evaluation_metrics[primary_metric].greater_is_better else 1,
+                _load_custom_metrics(recipe_root, [evaluation_metrics[primary_metric]])[0],
+            )
+        else:
+            raise MlflowException(
+                f"There is no FLAML alternative or custom metric for {primary_metric} metric."
+            )
+
+        automl_settings = step_config.get("flaml_params", {})
+        automl_settings["time_budget"] = step_config.get(
+            "time_budget_secs", _AUTOML_DEFAULT_TIME_BUDGET
+        )
+        automl_settings["metric"] = metric
+        automl_settings["task"] = task
+        # Disabled Autologging, because during the hyperparameter search
+        # it tries to log the same parameters multiple times.
+        mlflow.autolog(disable=True)
+        automl = AutoML()
+        automl.fit(X, y, **automl_settings)
+        mlflow.autolog(disable=False, log_models=False)
+        if automl.model is None:
+            raise MlflowException(
+                "AutoML (FLAML) could not train a suitable algorithm. "
+                "Maybe you should increase `time_budget_secs`parameter "
+                "to give AutoML process more time."
+            )
+        return automl.model.estimator, automl.best_config
+    except Exception as e:
+        _logger.warning(e, exc_info=e, stack_info=True)
+        raise MlflowException(
+            f"Error has occurred during training of AutoML model using FLAML: {e!r}"
+        )

venv/lib/python3.9/site-packages/mlflow/recipes/steps/evaluate.py

@@ -0,0 +1,502 @@
+import datetime
+import logging
+import operator
+import os
+import sys
+import warnings
+from collections import namedtuple
+from pathlib import Path
+from typing import Any
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.recipes.cards import BaseCard
+from mlflow.recipes.step import BaseStep, StepClass
+from mlflow.recipes.steps.train import TrainStep
+from mlflow.recipes.utils.execution import get_step_output_path
+from mlflow.recipes.utils.metrics import (
+    _get_builtin_metrics,
+    _get_custom_metrics,
+    _get_extended_task,
+    _get_model_type_from_template,
+    _get_primary_metric,
+    _load_custom_metrics,
+    transform_multiclass_metric,
+)
+from mlflow.recipes.utils.step import get_merged_eval_metrics, validate_classification_config
+from mlflow.recipes.utils.tracking import (
+    TrackingConfig,
+    apply_recipe_tracking_config,
+    get_recipe_tracking_config,
+    get_run_tags_env_vars,
+)
+from mlflow.tracking.fluent import _get_experiment_id, _set_experiment_primary_metric
+from mlflow.utils.databricks_utils import get_databricks_env_vars, get_databricks_run_url
+from mlflow.utils.string_utils import strip_prefix
+
+_logger = logging.getLogger(__name__)
+
+
+_FEATURE_IMPORTANCE_PLOT_FILE = "feature_importance.png"
+_VALIDATION_METRIC_PREFIX = "val_"
+
+MetricValidationResult = namedtuple(
+    "MetricValidationResult", ["metric", "greater_is_better", "value", "threshold", "validated"]
+)
+
+
+class EvaluateStep(BaseStep):
+    def __init__(self, step_config: dict[str, Any], recipe_root: str) -> None:
+        super().__init__(step_config, recipe_root)
+        self.tracking_config = TrackingConfig.from_dict(self.step_config)
+
+    def _validate_and_apply_step_config(self):
+        self.target_col = self.step_config.get("target_col")
+        if self.target_col is None:
+            raise MlflowException(
+                "Missing target_col config in recipe config.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        self.recipe = self.step_config.get("recipe")
+        if self.recipe is None:
+            raise MlflowException(
+                "Missing recipe config in recipe config.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        self.positive_class = self.step_config.get("positive_class")
+        self.extended_task = _get_extended_task(self.recipe, self.positive_class)
+        self.model_validation_status = "UNKNOWN"
+        self.primary_metric = _get_primary_metric(
+            self.step_config.get("primary_metric"), self.extended_task
+        )
+        self.user_defined_custom_metrics = {
+            metric.name: metric
+            for metric in _get_custom_metrics(self.step_config, self.extended_task)
+        }
+        self.evaluation_metrics = {
+            metric.name: metric for metric in _get_builtin_metrics(self.extended_task)
+        }
+        self.evaluation_metrics.update(self.user_defined_custom_metrics)
+        if self.primary_metric is not None and self.primary_metric not in self.evaluation_metrics:
+            raise MlflowException(
+                f"The primary metric '{self.primary_metric}' is a custom metric, but its"
+                " corresponding custom metric configuration is missing from `recipe.yaml`.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+    def _validate_validation_criteria(self):
+        """
+        Validates validation criteria don't contain undefined metrics
+        """
+        val_metrics = {vc["metric"] for vc in self.step_config.get("validation_criteria", [])}
+        if not val_metrics:
+            return
+        undefined_metrics = val_metrics.difference(self.evaluation_metrics.keys())
+        if undefined_metrics:
+            raise MlflowException(
+                f"Validation criteria contain undefined metrics: {sorted(undefined_metrics)}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+    def _check_validation_criteria(self, metrics, validation_criteria):
+        """
+        return a list of `MetricValidationResult` tuple instances.
+        """
+        summary = []
+        for val_criterion in validation_criteria:
+            metric_name = val_criterion["metric"]
+            metric_val = metrics.get(metric_name)
+            if metric_val is None:
+                raise MlflowException(
+                    f"The metric {metric_name} is defined in the recipe's validation criteria"
+                    " but was not returned from mlflow evaluation.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+            greater_is_better = self.evaluation_metrics[metric_name].greater_is_better
+            comp_func = operator.ge if greater_is_better else operator.le
+            threshold = val_criterion["threshold"]
+            validated = comp_func(metric_val, threshold)
+            summary.append(
+                MetricValidationResult(
+                    metric=metric_name,
+                    greater_is_better=greater_is_better,
+                    value=metric_val,
+                    threshold=threshold,
+                    validated=validated,
+                )
+            )
+        return summary
+
+    def _run(self, output_directory):
+        def my_warn(*args, **kwargs):
+            timestamp = datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S")
+            stacklevel = 1 if "stacklevel" not in kwargs else kwargs["stacklevel"]
+            frame = sys._getframe(stacklevel)
+            filename = frame.f_code.co_filename
+            lineno = frame.f_lineno
+            message = f"{timestamp} {filename}:{lineno}: {args[0]}\n"
+            with open(os.path.join(output_directory, "warning_logs.txt"), "a") as f:
+                f.write(message)
+
+        original_warn = warnings.warn
+        warnings.warn = my_warn
+        try:
+            import pandas as pd
+
+            with open(os.path.join(output_directory, "warning_logs.txt"), "w"):
+                pass
+
+            self._validate_validation_criteria()
+
+            test_df_path = get_step_output_path(
+                recipe_root_path=self.recipe_root,
+                step_name="split",
+                relative_path="test.parquet",
+            )
+            test_df = pd.read_parquet(test_df_path)
+            validate_classification_config(self.task, self.positive_class, test_df, self.target_col)
+
+            validation_df_path = get_step_output_path(
+                recipe_root_path=self.recipe_root,
+                step_name="split",
+                relative_path="validation.parquet",
+            )
+            validation_df = pd.read_parquet(validation_df_path)
+
+            run_id_path = get_step_output_path(
+                recipe_root_path=self.recipe_root,
+                step_name="train",
+                relative_path="run_id",
+            )
+            run_id = Path(run_id_path).read_text()
+
+            model_uri = get_step_output_path(
+                recipe_root_path=self.recipe_root,
+                step_name="train",
+                relative_path=TrainStep.SKLEARN_MODEL_ARTIFACT_RELATIVE_PATH,
+            )
+
+            apply_recipe_tracking_config(self.tracking_config)
+            exp_id = _get_experiment_id()
+
+            primary_metric_greater_is_better = self.evaluation_metrics[
+                self.primary_metric
+            ].greater_is_better
+
+            _set_experiment_primary_metric(
+                exp_id, f"test_{self.primary_metric}", primary_metric_greater_is_better
+            )
+
+            with mlflow.start_run(run_id=run_id):
+                eval_metrics = {}
+                for dataset_name, dataset, evaluator_config in (
+                    (
+                        "validation",
+                        validation_df,
+                        {
+                            "explainability_algorithm": "kernel",
+                            "explainability_nsamples": 10,
+                            "metric_prefix": _VALIDATION_METRIC_PREFIX,
+                        },
+                    ),
+                    (
+                        "test",
+                        test_df,
+                        {
+                            "log_model_explainability": False,
+                            "metric_prefix": "test_",
+                        },
+                    ),
+                ):
+                    if self.extended_task == "classification/binary":
+                        evaluator_config["pos_label"] = self.positive_class
+                    eval_result = mlflow.evaluate(
+                        model=model_uri,
+                        data=dataset,
+                        targets=self.target_col,
+                        model_type=_get_model_type_from_template(self.recipe),
+                        evaluators="default",
+                        extra_metrics=_load_custom_metrics(
+                            self.recipe_root,
+                            self.evaluation_metrics.values(),
+                        ),
+                        evaluator_config=evaluator_config,
+                    )
+                    eval_result.save(os.path.join(output_directory, f"eval_{dataset_name}"))
+                    eval_metrics[dataset_name] = {
+                        transform_multiclass_metric(
+                            strip_prefix(k, evaluator_config["metric_prefix"]), self.extended_task
+                        ): v
+                        for k, v in eval_result.metrics.items()
+                    }
+
+                validation_results = self._validate_model(eval_metrics, output_directory)
+
+            card = self._build_profiles_and_card(
+                run_id, model_uri, eval_metrics, validation_results, output_directory
+            )
+            card.save_as_html(output_directory)
+            self._log_step_card(run_id, self.name)
+            return card
+        finally:
+            warnings.warn = original_warn
+
+    def _validate_model(self, eval_metrics, output_directory):
+        validation_criteria = self.step_config.get("validation_criteria")
+        validation_results = None
+        if validation_criteria:
+            validation_results = self._check_validation_criteria(
+                eval_metrics["test"], validation_criteria
+            )
+            self.model_validation_status = (
+                "VALIDATED" if all(cr.validated for cr in validation_results) else "REJECTED"
+            )
+        else:
+            self.model_validation_status = "UNKNOWN"
+        Path(output_directory, "model_validation_status").write_text(self.model_validation_status)
+        return validation_results
+
+    def _build_profiles_and_card(
+        self, run_id, model_uri, eval_metrics, validation_results, output_directory
+    ):
+        """
+        Constructs data profiles of predictions and errors and a step card instance corresponding
+        to the current evaluate step state.
+
+        Args:
+            run_id: The ID of the MLflow Run to which to log model evaluation results.
+            model_uri: The URI of the model being evaluated.
+            eval_metrics: The evaluation result keyed by dataset name from `mlflow.evaluate`.
+            validation_results: A list of `MetricValidationResult` instances.
+            output_directory: Output directory used by the evaluate step.
+        """
+        import pandas as pd
+
+        # Build card
+        card = BaseCard(self.recipe_name, self.name)
+        # Tab 0: model performance summary.
+        metric_df = (
+            get_merged_eval_metrics(
+                eval_metrics,
+                ordered_metric_names=[self.primary_metric, *self.user_defined_custom_metrics],
+            )
+            .reset_index()
+            .rename(columns={"index": "Metric"})
+        )
+
+        def row_style(row):
+            if row.Metric == self.primary_metric or row.Metric in self.user_defined_custom_metrics:
+                return pd.Series("font-weight: bold", row.index)
+            else:
+                return pd.Series("", row.index)
+
+        metric_table_html = BaseCard.render_table(
+            metric_df.style.format({"training": "{:.6g}", "validation": "{:.6g}"}).apply(
+                row_style, axis=1
+            )
+        )
+
+        card.add_tab(
+            "Model Performance (Test)",
+            "<h3 class='section-title'>Summary Metrics</h3>"
+            "<b>NOTE</b>: Use evaluation metrics over test dataset with care. "
+            "Fine-tuning model over the test dataset is not advised."
+            "{{ METRICS }} ",
+        ).add_html("METRICS", metric_table_html)
+
+        # Tab 1: model validation results, if exists.
+        if validation_results is not None:
+
+            def get_icon(validated):
+                return (
+                    # check mark button emoji
+                    "\u2705"
+                    if validated
+                    # cross mark emoji
+                    else "\u274c"
+                )
+
+            result_df = pd.DataFrame(validation_results).assign(
+                validated=lambda df: df["validated"].map(get_icon)
+            )
+
+            criteria_html = BaseCard.render_table(
+                result_df.style.format({"value": "{:.6g}", "threshold": "{:.6g}"})
+            )
+            card.add_tab("Model Validation", "{{ METRIC_VALIDATION_RESULTS }}").add_html(
+                "METRIC_VALIDATION_RESULTS",
+                "<h3 class='section-title'>Model Validation Results (Test Dataset)</h3> "
+                + criteria_html,
+            )
+
+        # Tab 2: Classifier plots.
+        if self.recipe == "classification/v1":
+            classifiers_plot_tab = card.add_tab(
+                "Model Performance Plots",
+                "{{ CONFUSION_MATRIX }} {{CONFUSION_MATRIX_PLOT}}"
+                + "{{ LIFT_CURVE }} {{LIFT_CURVE_PLOT}}"
+                + "{{ PR_CURVE }} {{PR_CURVE_PLOT}}"
+                + "{{ ROC_CURVE }} {{ROC_CURVE_PLOT}}",
+            )
+            confusion_matrix_path = os.path.join(
+                output_directory,
+                "eval_validation/artifacts",
+                f"{_VALIDATION_METRIC_PREFIX}confusion_matrix.png",
+            )
+            if os.path.exists(confusion_matrix_path):
+                classifiers_plot_tab.add_html(
+                    "CONFUSION_MATRIX",
+                    '<h3 class="section-title">Confusion Matrix Plot</h3>',
+                )
+                classifiers_plot_tab.add_image(
+                    "CONFUSION_MATRIX_PLOT", confusion_matrix_path, width=400
+                )
+
+            lift_curve_path = os.path.join(
+                output_directory,
+                "eval_validation/artifacts",
+                f"{_VALIDATION_METRIC_PREFIX}lift_curve_plot.png",
+            )
+            if os.path.exists(lift_curve_path):
+                classifiers_plot_tab.add_html(
+                    "LIFT_CURVE",
+                    '<h3 class="section-title">Lift Curve Plot</h3>',
+                )
+                classifiers_plot_tab.add_image("LIFT_CURVE_PLOT", lift_curve_path, width=400)
+
+            pr_curve_path = os.path.join(
+                output_directory,
+                "eval_validation/artifacts",
+                f"{_VALIDATION_METRIC_PREFIX}precision_recall_curve_plot.png",
+            )
+            if os.path.exists(pr_curve_path):
+                classifiers_plot_tab.add_html(
+                    "PR_CURVE",
+                    '<h3 class="section-title">Precision Recall Curve Plot</h3>',
+                )
+                classifiers_plot_tab.add_image("PR_CURVE_PLOT", pr_curve_path, width=400)
+
+            roc_curve_path = os.path.join(
+                output_directory,
+                "eval_validation/artifacts",
+                f"{_VALIDATION_METRIC_PREFIX}roc_curve_plot.png",
+            )
+            if os.path.exists(roc_curve_path):
+                classifiers_plot_tab.add_html(
+                    "ROC_CURVE",
+                    '<h3 class="section-title">ROC Curve Plot</h3>',
+                )
+                classifiers_plot_tab.add_image("ROC_CURVE_PLOT", roc_curve_path, width=400)
+
+        # Tab 3: SHAP plots.
+        def _add_shap_plots(card):
+            """Contingent on shap being installed."""
+            shap_plot_tab = card.add_tab(
+                "Feature Importance",
+                '<h3 class="section-title">Feature Importance on Validation Dataset</h3>'
+                '<h3 class="section-title">SHAP Bar Plot</h3>{{SHAP_BAR_PLOT}}'
+                '<h3 class="section-title">SHAP Beeswarm Plot</h3>{{SHAP_BEESWARM_PLOT}}',
+            )
+
+            shap_bar_plot_path = os.path.join(
+                output_directory,
+                "eval_validation/artifacts",
+                f"{_VALIDATION_METRIC_PREFIX}shap_feature_importance_plot.png",
+            )
+            shap_beeswarm_plot_path = os.path.join(
+                output_directory,
+                "eval_validation/artifacts",
+                f"{_VALIDATION_METRIC_PREFIX}shap_beeswarm_plot.png",
+            )
+            shap_plot_tab.add_image("SHAP_BAR_PLOT", shap_bar_plot_path, width=800)
+            shap_plot_tab.add_image("SHAP_BEESWARM_PLOT", shap_beeswarm_plot_path, width=800)
+
+        try:
+            import shap  # noqa: F401
+            from matplotlib import pyplot  # noqa: F401
+
+            _add_shap_plots(card)
+        except ImportError:
+            _logger.warning(
+                "SHAP or matplotlib package is not installed, so shap plots will not be added."
+            )
+
+        # Tab 3: Warning log outputs.
+        warning_output_path = os.path.join(output_directory, "warning_logs.txt")
+        if os.path.exists(warning_output_path):
+            warnings_output_tab = card.add_tab("Warning Logs", "{{ STEP_WARNINGS }}")
+            with open(warning_output_path) as f:
+                warnings_output_tab.add_html("STEP_WARNINGS", f"<pre>{f.read()}</pre>")
+
+        # Tab 4: Run summary.
+        run_summary_card_tab = card.add_tab(
+            "Run Summary",
+            "{{ RUN_ID }} "
+            + "{{ MODEL_URI }}"
+            + "{{ VALIDATION_STATUS }}"
+            + "{{ EXE_DURATION }}"
+            + "{{ LAST_UPDATE_TIME }}",
+        ).add_markdown(
+            "VALIDATION_STATUS", f"**Validation status:** `{self.model_validation_status}`"
+        )
+        run_url = get_databricks_run_url(
+            tracking_uri=mlflow.get_tracking_uri(),
+            run_id=run_id,
+        )
+        model_uri = f"runs:/{run_id}/train/{TrainStep.MODEL_ARTIFACT_RELATIVE_PATH}"
+        model_url = get_databricks_run_url(
+            tracking_uri=mlflow.get_tracking_uri(),
+            run_id=run_id,
+            artifact_path=f"train/{TrainStep.MODEL_ARTIFACT_RELATIVE_PATH}",
+        )
+
+        if run_url is not None:
+            run_summary_card_tab.add_html(
+                "RUN_ID", f"<b>MLflow Run ID:</b> <a href={run_url}>{run_id}</a><br><br>"
+            )
+        else:
+            run_summary_card_tab.add_markdown("RUN_ID", f"**MLflow Run ID:** `{run_id}`")
+
+        if model_url is not None:
+            run_summary_card_tab.add_html(
+                "MODEL_URI", f"<b>MLflow Model URI:</b> <a href={model_url}>{model_uri}</a>"
+            )
+        else:
+            run_summary_card_tab.add_markdown("MODEL_URI", f"**MLflow Model URI:** `{model_uri}`")
+
+        return card
+
+    @classmethod
+    def from_recipe_config(cls, recipe_config, recipe_root):
+        step_config = {}
+        if recipe_config.get("steps", {}).get("evaluate", {}) is not None:
+            step_config.update(recipe_config.get("steps", {}).get("evaluate", {}))
+        step_config["target_col"] = recipe_config.get("target_col")
+        if "positive_class" in recipe_config:
+            step_config["positive_class"] = recipe_config.get("positive_class")
+        if recipe_config.get("custom_metrics") is not None:
+            step_config["custom_metrics"] = recipe_config["custom_metrics"]
+        if recipe_config.get("primary_metric") is not None:
+            step_config["primary_metric"] = recipe_config["primary_metric"]
+        step_config["recipe"] = recipe_config.get("recipe")
+        step_config.update(
+            get_recipe_tracking_config(
+                recipe_root_path=recipe_root,
+                recipe_config=recipe_config,
+            ).to_dict()
+        )
+        return cls(step_config, recipe_root)
+
+    @property
+    def name(self):
+        return "evaluate"
+
+    @property
+    def environment(self):
+        environ = get_databricks_env_vars(tracking_uri=self.tracking_config.tracking_uri)
+        environ.update(get_run_tags_env_vars(recipe_root_path=self.recipe_root))
+        return environ
+
+    def step_class(self):
+        return StepClass.TRAINING

venv/lib/python3.9/site-packages/mlflow/recipes/steps/ingest/__init__.py

@@ -0,0 +1,275 @@
+import abc
+import logging
+import os
+from pathlib import Path
+from typing import Any, Optional
+
+import pandas as pd
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.recipes.artifacts import DataframeArtifact
+from mlflow.recipes.cards import BaseCard
+from mlflow.recipes.step import BaseStep, StepClass
+from mlflow.recipes.steps.ingest.datasets import (
+    CustomDataset,
+    DeltaTableDataset,
+    ParquetDataset,
+    SparkSqlDataset,
+)
+from mlflow.recipes.utils.step import get_pandas_data_profiles, validate_classification_config
+from mlflow.utils.file_utils import read_parquet_as_pandas_df
+
+_logger = logging.getLogger(__name__)
+
+
+class BaseIngestStep(BaseStep, metaclass=abc.ABCMeta):
+    _DATASET_FORMAT_SPARK_TABLE = "spark_table"
+    _DATASET_FORMAT_DELTA = "delta"
+    _DATASET_FORMAT_PARQUET = "parquet"
+    _DATASET_PROFILE_OUTPUT_NAME = "dataset_profile.html"
+    _STEP_CARD_OUTPUT_NAME = "card.pkl"
+    _SUPPORTED_DATASETS = [
+        ParquetDataset,
+        DeltaTableDataset,
+        SparkSqlDataset,
+        # NB: The custom dataset is deliberately listed last as a catch-all for any
+        # format not matched by the datasets above. When mapping a format to a dataset,
+        # datasets are explored in the listed order
+        CustomDataset,
+    ]
+
+    def _validate_and_apply_step_config(self):
+        dataset_format = self.step_config.get("using")
+        if not dataset_format:
+            raise MlflowException(
+                message=(
+                    "Dataset format must be specified via the `using` key within the `ingest`"
+                    " section of recipe.yaml"
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        if self.step_class() == StepClass.TRAINING:
+            self.target_col = self.step_config.get("target_col")
+            if self.target_col is None:
+                raise MlflowException(
+                    "Missing target_col config in recipe config.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+            self.positive_class = self.step_config.get("positive_class")
+        for dataset_class in BaseIngestStep._SUPPORTED_DATASETS:
+            if dataset_class.handles_format(dataset_format):
+                self.dataset = dataset_class.from_config(
+                    dataset_config=self.step_config,
+                    recipe_root=self.recipe_root,
+                )
+                break
+        else:
+            raise MlflowException(
+                message=f"Unrecognized dataset format: {dataset_format}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
+
+    def _run(self, output_directory: str) -> BaseCard:
+        dataset_dst_path = os.path.abspath(os.path.join(output_directory, self.dataset_output_name))
+        self.dataset.resolve_to_parquet(
+            dst_path=dataset_dst_path,
+        )
+        _logger.debug("Successfully stored data in parquet format at '%s'", dataset_dst_path)
+
+        ingested_df = read_parquet_as_pandas_df(data_parquet_path=dataset_dst_path)
+        if self.step_class() == StepClass.TRAINING:
+            if self.target_col not in ingested_df.columns:
+                raise MlflowException(
+                    f"Target column '{self.target_col}' not found in ingested dataset.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+            if self.task == "classification":
+                validate_classification_config(
+                    self.task, self.positive_class, ingested_df, self.target_col
+                )
+                cardinality = ingested_df[self.target_col].nunique()
+                if cardinality > 2 and self.positive_class is not None:
+                    raise MlflowException(
+                        f"Target column '{self.target_col}' must have a cardinality of 2,"
+                        f"found '{cardinality}'.",
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+                if self.positive_class is not None and cardinality != 2:
+                    raise MlflowException(
+                        f"`For binary classification problems, "
+                        f"target column '{self.target_col}' must have a cardinality of 2,"
+                        f"found '{cardinality}'. `positive_class` was set, "
+                        f"so we treat this problem as a binary classification problem. ",
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+
+        ingested_dataset_profile = None
+        if not self.skip_data_profiling:
+            _logger.debug("Profiling ingested dataset")
+            ingested_dataset_profile = get_pandas_data_profiles(
+                [["Profile of Ingested Dataset", ingested_df]]
+            )
+            dataset_profile_path = Path(
+                str(os.path.join(output_directory, BaseIngestStep._DATASET_PROFILE_OUTPUT_NAME))
+            )
+            dataset_profile_path.write_text(ingested_dataset_profile, encoding="utf-8")
+            _logger.debug(f"Wrote dataset profile to '{dataset_profile_path}'")
+
+        schema = pd.io.json.build_table_schema(ingested_df, index=False)
+
+        return self._build_step_card(
+            ingested_dataset_profile=ingested_dataset_profile,
+            ingested_rows=len(ingested_df),
+            schema=schema,
+            data_preview=ingested_df.head(),
+            dataset_src_location=getattr(self.dataset, "location", None),
+            dataset_sql=getattr(self.dataset, "sql", None),
+        )
+
+    def _build_step_card(  # noqa: D417
+        self,
+        ingested_dataset_profile: str,
+        ingested_rows: int,
+        schema: dict,
+        data_preview: pd.DataFrame = None,
+        dataset_src_location: Optional[str] = None,
+        dataset_sql: Optional[str] = None,
+    ) -> BaseCard:
+        """
+        Constructs a step card instance corresponding to the current ingest step state.
+
+        Args:
+            ingested_dataset_path: The local filesystem path to the ingested parquet dataset file.
+            dataset_src_location: The source location of the dataset (e.g. '/tmp/myfile.parquet',
+                's3://mybucket/mypath', ...), if the dataset is a location-based dataset. Either
+                ``dataset_src_location`` or ``dataset_sql`` must be specified.
+            dataset_sql: The Spark SQL query string that defines the dataset
+                (e.g. 'SELECT * FROM my_spark_table'), if the dataset is a Spark SQL dataset. Either
+                ``dataset_src_location`` or ``dataset_sql`` must be specified.
+
+        Returns:
+            An BaseCard instance corresponding to the current ingest step state.
+
+        """
+        if dataset_src_location is None and dataset_sql is None:
+            raise MlflowException(
+                message=(
+                    "Failed to build step card because neither a dataset location nor a"
+                    " dataset Spark SQL query were specified"
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        card = BaseCard(self.recipe_name, self.name)
+        if not self.skip_data_profiling:
+            (  # Tab #1 -- Ingested dataset profile.
+                card.add_tab("Data Profile", "{{PROFILE}}").add_pandas_profile(
+                    "PROFILE", ingested_dataset_profile
+                )
+            )
+        # Tab #2 -- Ingested dataset schema.
+        schema_html = BaseCard.render_table(schema["fields"])
+        card.add_tab("Data Schema", "{{SCHEMA}}").add_html("SCHEMA", schema_html)
+
+        if data_preview is not None:
+            # Tab #3 -- Ingested dataset preview.
+            card.add_tab("Data Preview", "{{DATA_PREVIEW}}").add_html(
+                "DATA_PREVIEW", BaseCard.render_table(data_preview)
+            )
+
+        (  # Tab #4 -- Step run summary.
+            card.add_tab(
+                "Run Summary",
+                "{{ INGESTED_ROWS }}"
+                + "{{ DATA_SOURCE }}"
+                + "{{ EXE_DURATION }}"
+                + "{{ LAST_UPDATE_TIME }}",
+            )
+            .add_markdown(
+                name="INGESTED_ROWS",
+                markdown=f"**Number of rows ingested:** `{ingested_rows}`",
+            )
+            .add_markdown(
+                name="DATA_SOURCE",
+                markdown=(
+                    f"**Dataset source location:** `{dataset_src_location}`"
+                    if dataset_src_location is not None
+                    else f"**Dataset SQL:** `{dataset_sql}`"
+                ),
+            )
+        )
+        return card
+
+
+class IngestStep(BaseIngestStep):
+    _DATASET_OUTPUT_NAME = "dataset.parquet"
+
+    def __init__(self, step_config: dict[str, Any], recipe_root: str):
+        super().__init__(step_config, recipe_root)
+        self.dataset_output_name = IngestStep._DATASET_OUTPUT_NAME
+
+    @classmethod
+    def from_recipe_config(cls, recipe_config: dict[str, Any], recipe_root: str):
+        ingest_config = recipe_config.get("steps", {}).get("ingest", {})
+        target_config = {"target_col": recipe_config.get("target_col")}
+        if "positive_class" in recipe_config:
+            target_config["positive_class"] = recipe_config.get("positive_class")
+        return cls(
+            step_config={
+                **ingest_config,
+                **target_config,
+                **{"recipe": recipe_config.get("recipe", "regression/v1")},
+            },
+            recipe_root=recipe_root,
+        )
+
+    @property
+    def name(self) -> str:
+        return "ingest"
+
+    def get_artifacts(self):
+        return [
+            DataframeArtifact(
+                "ingested_data", self.recipe_root, self.name, IngestStep._DATASET_OUTPUT_NAME
+            )
+        ]
+
+    def step_class(self):
+        return StepClass.TRAINING
+
+
+class IngestScoringStep(BaseIngestStep):
+    _DATASET_OUTPUT_NAME = "scoring-dataset.parquet"
+
+    def __init__(self, step_config: dict[str, Any], recipe_root: str):
+        super().__init__(step_config, recipe_root)
+        self.dataset_output_name = IngestScoringStep._DATASET_OUTPUT_NAME
+
+    @classmethod
+    def from_recipe_config(cls, recipe_config: dict[str, Any], recipe_root: str):
+        step_config = recipe_config.get("steps", {}).get("ingest_scoring", {})
+        step_config["recipe"] = recipe_config.get("recipe")
+        return cls(
+            step_config=step_config,
+            recipe_root=recipe_root,
+        )
+
+    @property
+    def name(self) -> str:
+        return "ingest_scoring"
+
+    def get_artifacts(self):
+        return [
+            DataframeArtifact(
+                "ingested_scoring_data",
+                self.recipe_root,
+                self.name,
+                IngestScoringStep._DATASET_OUTPUT_NAME,
+            )
+        ]
+
+    def step_class(self):
+        return StepClass.PREDICTION

venv/lib/python3.9/site-packages/mlflow/recipes/steps/ingest/datasets.py

@@ -0,0 +1,675 @@
+import importlib
+import logging
+import os
+import pathlib
+import posixpath
+import sys
+from abc import abstractmethod
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Optional, Union
+from urllib.parse import urlparse
+
+from mlflow.artifacts import download_artifacts
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import BAD_REQUEST, INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repo import (
+    _NUM_DEFAULT_CPUS,
+    _NUM_MAX_THREADS,
+    _NUM_MAX_THREADS_PER_CPU,
+)
+from mlflow.utils._spark_utils import (
+    _create_local_spark_session_for_recipes,
+    _get_active_spark_session,
+)
+from mlflow.utils.file_utils import (
+    TempDir,
+    download_file_using_http_uri,
+    get_local_path_or_none,
+    local_file_uri_to_path,
+    read_parquet_as_pandas_df,
+    write_pandas_df_as_parquet,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+_USER_DEFINED_INGEST_STEP_MODULE = "steps.ingest"
+
+
+class _Dataset:
+    """
+    Base class representing an ingestable dataset.
+    """
+
+    def __init__(self, dataset_format: str):
+        """
+        Args:
+            dataset_format: The format of the dataset (e.g. 'csv', 'parquet', ...).
+        """
+        self.dataset_format = dataset_format
+
+    @abstractmethod
+    def resolve_to_parquet(self, dst_path: str):
+        """
+        Fetches the dataset, converts it to parquet, and stores it at the specified `dst_path`.
+
+        Args:
+            dst_path: The local filesystem path at which to store the resolved parquet dataset
+                (e.g. `<execution_directory_path>/steps/ingest/outputs/dataset.parquet`).
+        """
+
+    @classmethod
+    def from_config(cls, dataset_config: dict[str, Any], recipe_root: str) -> "_Dataset":
+        """
+        Constructs a dataset instance from the specified dataset configuration
+        and recipe root path.
+
+        Args:
+            dataset_config: Dictionary representation of the recipe dataset configuration
+                (i.e. the `data` section of recipe.yaml).
+            recipe_root: The absolute path of the associated recipe root directory on the
+                local filesystem.
+
+        Returns:
+            A `_Dataset` instance representing the configured dataset.
+        """
+        if not cls.handles_format(dataset_config.get("using")):
+            raise MlflowException(
+                f"Invalid format {dataset_config.get('using')} for dataset {cls}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        return cls._from_config(dataset_config, recipe_root)
+
+    @classmethod
+    @abstractmethod
+    def _from_config(cls, dataset_config, recipe_root) -> "_Dataset":
+        """
+        Constructs a dataset instance from the specified dataset configuration
+        and recipe root path.
+
+        Args:
+            dataset_config: Dictionary representation of the recipe dataset configuration
+                (i.e. the `data` section of recipe.yaml).
+            recipe_root: The absolute path of the associated recipe root directory on the
+                local filesystem.
+
+        Returns:
+            A `_Dataset` instance representing the configured dataset.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def handles_format(dataset_format: str) -> bool:
+        """
+        Determines whether or not the dataset class is a compatible representation of the
+        specified dataset format.
+
+        Args:
+            dataset_format: The format of the dataset (e.g. 'csv', 'parquet', ...).
+
+        Returns:
+            `True` if the dataset class is a compatible representation of the specified
+            dataset format, `False` otherwise.
+        """
+
+    @classmethod
+    def _get_required_config(cls, dataset_config: dict[str, Any], key: str) -> Any:
+        """
+        Obtains the value associated with the specified dataset configuration key, first verifying
+        that the key is present in the config and throwing if it is not.
+
+        Args:
+            dataset_config: Dictionary representation of the recipe dataset configuration
+                (i.e. the `data` section of recipe.yaml).
+            key: The key within the dataset configuration for which to fetch the associated
+                value.
+
+        Returns:
+            The value associated with the specified configuration key.
+        """
+        try:
+            return dataset_config[key]
+        except KeyError:
+            raise MlflowException(
+                f"The `{key}` configuration key must be specified for dataset with"
+                f" using '{dataset_config.get('using')}' format"
+            ) from None
+
+
+class _LocationBasedDataset(_Dataset):
+    """
+    Base class representing an ingestable dataset with a configurable `location` attribute.
+    """
+
+    def __init__(
+        self,
+        location: Union[str, list[str]],
+        dataset_format: str,
+        recipe_root: str,
+    ):
+        """
+        Args:
+            location: The location of the dataset (one dataset as a string or list of multiple
+                datasets)
+                (e.g. '/tmp/myfile.parquet', './mypath', 's3://mybucket/mypath', or YAML list:
+                    location:
+                        - http://www.myserver.com/dataset/df1.csv
+                        - http://www.myserver.com/dataset/df1.csv
+                )
+
+            dataset_format: The format of the dataset (e.g. 'csv', 'parquet', ...).
+
+            recipe_root: The absolute path of the associated recipe root directory on the local
+                filesystem.
+        """
+        super().__init__(dataset_format=dataset_format)
+        self.location = (
+            _LocationBasedDataset._sanitize_local_dataset_multiple_locations_if_necessary(
+                dataset_location=location,
+                recipe_root=recipe_root,
+            )
+        )
+
+    @abstractmethod
+    def resolve_to_parquet(self, dst_path: str):
+        pass
+
+    @classmethod
+    def _from_config(cls, dataset_config: dict[str, Any], recipe_root: str) -> "_Dataset":
+        return cls(
+            location=cls._get_required_config(dataset_config=dataset_config, key="location"),
+            recipe_root=recipe_root,
+            dataset_format=cls._get_required_config(dataset_config=dataset_config, key="using"),
+        )
+
+    @staticmethod
+    def _sanitize_local_dataset_multiple_locations_if_necessary(
+        dataset_location: Union[str, list[str]], recipe_root: str
+    ) -> list[str]:
+        if isinstance(dataset_location, str):
+            return [
+                _LocationBasedDataset._sanitize_local_dataset_location_if_necessary(
+                    dataset_location, recipe_root
+                )
+            ]
+        elif isinstance(dataset_location, list):
+            return [
+                _LocationBasedDataset._sanitize_local_dataset_location_if_necessary(
+                    locaton, recipe_root
+                )
+                for locaton in dataset_location
+            ]
+        else:
+            raise MlflowException(f"Unsupported location type: {type(dataset_location)}")
+
+    @staticmethod
+    def _sanitize_local_dataset_location_if_necessary(
+        dataset_location: str, recipe_root: str
+    ) -> str:
+        """
+        Checks whether or not the specified `dataset_location` is a local filesystem location and,
+        if it is, converts it to an absolute path if it is not already absolute.
+
+        Args:
+            dataset_location: The dataset location from the recipe dataset configuration.
+            recipe_root: The absolute path of the recipe root directory on the local
+                filesystem.
+
+        Returns:
+            The sanitized dataset location.
+        """
+        local_dataset_path_or_none = get_local_path_or_none(path_or_uri=dataset_location)
+        if local_dataset_path_or_none is None:
+            return dataset_location
+
+        # If the local dataset path is a file: URI, convert it to a filesystem path
+        local_dataset_path = local_file_uri_to_path(uri=local_dataset_path_or_none)
+        local_dataset_path = pathlib.Path(local_dataset_path)
+        if local_dataset_path.is_absolute():
+            return str(local_dataset_path)
+        else:
+            # Use pathlib to join the local dataset relative path with the recipe root
+            # directory to correctly handle the case where the root path is Windows-formatted
+            # and the local dataset relative path is POSIX-formatted
+            return str(pathlib.Path(recipe_root) / local_dataset_path)
+
+    @staticmethod
+    @abstractmethod
+    def handles_format(dataset_format: str) -> bool:
+        pass
+
+
+class _DownloadThenConvertDataset(_LocationBasedDataset):
+    """
+    Base class representing a location-based ingestible dataset that is resolved in two distinct
+    phases: 1. Download the dataset files to the local filesystem. 2. Convert the dataset files to
+    parquet format, aggregating them together as a single parquet file.
+    `_DownloadThenConvertDataset` implements phase (1) and provides an abstract method
+    for phase (2).
+    """
+
+    _FILE_DOWNLOAD_CHUNK_SIZE_BYTES = 10**7  # 10MB
+
+    def resolve_to_parquet(self, dst_path: str):
+        with TempDir(chdr=True) as tmpdir:
+            _logger.debug("Resolving input data from '%s'", self.location)
+            local_dataset_path = _DownloadThenConvertDataset._download_dataset(
+                dataset_location=self.location,
+                dst_path=tmpdir.path(),
+            )
+
+            if os.path.isdir(local_dataset_path):
+                # NB: Sort the file names alphanumerically to ensure a consistent
+                # ordering across invocations
+                if self.dataset_format == "custom":
+                    dataset_file_paths = sorted(pathlib.Path(local_dataset_path).glob("*"))
+                else:
+                    dataset_file_paths = sorted(
+                        pathlib.Path(local_dataset_path).glob(f"*.{self.dataset_format}")
+                    )
+                if len(dataset_file_paths) == 0:
+                    raise MlflowException(
+                        message=(
+                            "Did not find any data files with the specified format"
+                            f" '{self.dataset_format}' in the resolved data directory with path"
+                            f" '{local_dataset_path}'. Directory contents:"
+                            f" {os.listdir(local_dataset_path)}."
+                        ),
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+            else:
+                if self.dataset_format != "custom" and not local_dataset_path.endswith(
+                    f".{self.dataset_format}"
+                ):
+                    raise MlflowException(
+                        message=(
+                            f"Resolved data file with path '{local_dataset_path}' does not have the"
+                            f" expected format '{self.dataset_format}'."
+                        ),
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+                dataset_file_paths = [local_dataset_path]
+
+            _logger.debug("Resolved input data to '%s'", local_dataset_path)
+            _logger.debug("Converting dataset to parquet format, if necessary")
+            return self._convert_to_parquet(
+                dataset_file_paths=dataset_file_paths,
+                dst_path=dst_path,
+            )
+
+    @staticmethod
+    def _download_dataset(dataset_location: list[str], dst_path: str):
+        dest_locations = _DownloadThenConvertDataset._download_all_datasets_in_parallel(
+            dataset_location, dst_path
+        )
+        if len(dest_locations) == 1:
+            return dest_locations[0]
+        else:
+            res_path = pathlib.Path(dest_locations[0])
+            if res_path.is_dir():
+                return str(res_path)
+            else:
+                return str(res_path.parent)
+
+    @staticmethod
+    def _download_all_datasets_in_parallel(dataset_location, dst_path):
+        num_cpus = os.cpu_count() or _NUM_DEFAULT_CPUS
+        with ThreadPoolExecutor(
+            max_workers=min(num_cpus * _NUM_MAX_THREADS_PER_CPU, _NUM_MAX_THREADS)
+        ) as executor:
+            futures = []
+            for location in dataset_location:
+                future = executor.submit(
+                    _DownloadThenConvertDataset._download_one_dataset,
+                    dataset_location=location,
+                    dst_path=dst_path,
+                )
+                futures.append(future)
+
+            dest_locations = []
+            failed_downloads = []
+            for future in as_completed(futures):
+                try:
+                    dest_locations.append(future.result())
+                except Exception as e:
+                    failed_downloads.append(repr(e))
+            if len(failed_downloads) > 0:
+                raise MlflowException(
+                    "During downloading of the datasets a number "
+                    + f"of errors have occurred: {failed_downloads}"
+                )
+            return dest_locations
+
+    @staticmethod
+    def _download_one_dataset(dataset_location: str, dst_path: str):
+        parsed_location_uri = urlparse(dataset_location)
+        if parsed_location_uri.scheme in ["http", "https"]:
+            dst_file_name = posixpath.basename(parsed_location_uri.path)
+            dst_file_path = os.path.join(dst_path, dst_file_name)
+            download_file_using_http_uri(
+                http_uri=dataset_location,
+                download_path=dst_file_path,
+                chunk_size=_DownloadThenConvertDataset._FILE_DOWNLOAD_CHUNK_SIZE_BYTES,
+            )
+            return dst_file_path
+        else:
+            return download_artifacts(artifact_uri=dataset_location, dst_path=dst_path)
+
+    @abstractmethod
+    def _convert_to_parquet(self, dataset_file_paths: list[str], dst_path: str):
+        """
+        Converts the specified dataset files to parquet format and aggregates them together,
+        writing the consolidated parquet file to the specified destination path.
+
+        Args:
+            dataset_file_paths: A list of local filesystem of dataset files to convert to
+                parquet format.
+            dst_path: The local filesystem path at which to store the resolved parquet dataset
+                (e.g. `<execution_directory_path>/steps/ingest/outputs/dataset.parquet`).
+
+        """
+
+
+class _PandasConvertibleDataset(_DownloadThenConvertDataset):
+    """
+    Base class representing a location-based ingestable dataset that can be parsed and converted to
+    parquet using a series of Pandas DataFrame ``read_*`` and ``concat`` operations.
+    """
+
+    def _convert_to_parquet(self, dataset_file_paths: list[str], dst_path: str):
+        import pandas as pd
+
+        aggregated_dataframe = None
+        for data_file_path in dataset_file_paths:
+            _path = pathlib.Path(data_file_path)
+            data_file_as_dataframe = self._load_file_as_pandas_dataframe(
+                local_data_file_path=data_file_path,
+            )
+            aggregated_dataframe = (
+                pd.concat([aggregated_dataframe, data_file_as_dataframe])
+                if aggregated_dataframe is not None
+                else data_file_as_dataframe
+            )
+
+        write_pandas_df_as_parquet(df=aggregated_dataframe, data_parquet_path=dst_path)
+
+    @abstractmethod
+    def _load_file_as_pandas_dataframe(self, local_data_file_path: str):
+        """
+        Loads the specified file as a Pandas DataFrame.
+
+        Args:
+            local_data_file_path: The local filesystem path of the file to load.
+
+        Returns:
+            A Pandas DataFrame representation of the specified file.
+
+        """
+
+    @staticmethod
+    @abstractmethod
+    def handles_format(dataset_format: str) -> bool:
+        pass
+
+
+class ParquetDataset(_PandasConvertibleDataset):
+    """
+    Representation of a dataset in parquet format with files having the `.parquet` extension.
+    """
+
+    def _load_file_as_pandas_dataframe(self, local_data_file_path: str):
+        return read_parquet_as_pandas_df(data_parquet_path=local_data_file_path)
+
+    @staticmethod
+    def handles_format(dataset_format: str) -> bool:
+        return dataset_format == "parquet"
+
+
+class CustomDataset(_PandasConvertibleDataset):
+    """
+    Representation of a location-based dataset with files containing a consistent, custom
+    extension (e.g. 'csv', 'csv.gz', 'json', ...), as well as a custom function used to load
+    and convert the dataset to parquet format.
+    """
+
+    def __init__(
+        self,
+        location: str,
+        dataset_format: str,
+        loader_method: str,
+        recipe_root: str,
+    ):
+        """
+        Args:
+            location: The location of the dataset
+                (e.g. '/tmp/myfile.parquet', './mypath', 's3://mybucket/mypath', ...).
+            dataset_format: The format of the dataset (e.g. 'csv', 'parquet', ...).
+            loader_method: The custom loader method used to load and convert the dataset
+                to parquet format, e.g.`load_file_as_dataframe`.
+            recipe_root: The absolute path of the associated recipe root directory on the
+                local filesystem.
+        """
+        super().__init__(
+            location=location,
+            dataset_format=dataset_format,
+            recipe_root=recipe_root,
+        )
+        self.recipe_root = recipe_root
+        self.loader_method = loader_method
+
+    def _validate_user_code_output(self, func, *args):
+        import pandas as pd
+
+        ingested_df = func(*args)
+        if not isinstance(ingested_df, pd.DataFrame):
+            raise MlflowException(
+                message=(
+                    "The `ingested_data` is not a DataFrame, please make sure "
+                    f"'{_USER_DEFINED_INGEST_STEP_MODULE}.{self.loader_method}' "
+                    "returns a Pandas DataFrame object."
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            ) from None
+        return ingested_df
+
+    def _load_file_as_pandas_dataframe(self, local_data_file_path: str):
+        try:
+            sys.path.append(self.recipe_root)
+            loader_method = getattr(
+                importlib.import_module(_USER_DEFINED_INGEST_STEP_MODULE),
+                self.loader_method,
+            )
+        except Exception as e:
+            raise MlflowException(
+                message=(
+                    "Failed to import custom dataset loader function"
+                    f" '{_USER_DEFINED_INGEST_STEP_MODULE}.{self.loader_method}' for"
+                    f" ingesting dataset with format '{self.dataset_format}'.",
+                ),
+                error_code=BAD_REQUEST,
+            ) from e
+
+        try:
+            return self._validate_user_code_output(
+                loader_method, local_data_file_path, self.dataset_format
+            )
+        except MlflowException as e:
+            raise e
+        except NotImplementedError:
+            raise MlflowException(
+                message=(
+                    f"Unable to load data file at path '{local_data_file_path}' with format"
+                    f" '{self.dataset_format}' using custom loader method"
+                    f" '{loader_method.__name__}' because it is not"
+                    " supported. Please update the custom loader method to support this"
+                    " format."
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            ) from None
+        except Exception as e:
+            raise MlflowException(
+                message=(
+                    f"Unable to load data file at path '{local_data_file_path}' with format"
+                    f" '{self.dataset_format}' using custom loader method"
+                    f" '{loader_method.__name__}'."
+                ),
+                error_code=BAD_REQUEST,
+            ) from e
+
+    @classmethod
+    def _from_config(cls, dataset_config: dict[str, Any], recipe_root: str) -> "_Dataset":
+        return cls(
+            location=cls._get_required_config(dataset_config=dataset_config, key="location"),
+            dataset_format=cls._get_required_config(dataset_config=dataset_config, key="using"),
+            loader_method=cls._get_required_config(
+                dataset_config=dataset_config, key="loader_method"
+            ),
+            recipe_root=recipe_root,
+        )
+
+    @staticmethod
+    def handles_format(dataset_format: str) -> bool:
+        return dataset_format is not None
+
+
+class _SparkDatasetMixin:
+    """
+    Mixin class providing Spark-related utilities for Datasets that use Spark for resolution
+    and conversion to parquet format.
+    """
+
+    def _convert_spark_df_to_pandas(self, spark_df):
+        import pandas as pd
+
+        datetime_cols = [
+            field.name for field in spark_df.schema.fields if str(field.dataType) == "DateType"
+        ]
+        pandas_df = spark_df.toPandas()
+        pandas_df[datetime_cols] = pandas_df[datetime_cols].apply(pd.to_datetime, errors="coerce")
+
+        return pandas_df
+
+    def _get_or_create_spark_session(self):
+        """
+        Obtains the active Spark session, throwing if a session does not exist.
+
+        Returns:
+            The active Spark session.
+        """
+        try:
+            spark_session = _get_active_spark_session()
+            if spark_session:
+                _logger.debug("Found active spark session")
+            else:
+                spark_session = _create_local_spark_session_for_recipes()
+                _logger.debug("Creating new spark session")
+            return spark_session
+        except Exception as e:
+            raise MlflowException(
+                message=(
+                    f"Encountered an error while searching for an active Spark session to"
+                    f" load the dataset with format '{self.dataset_format}'. Please create a"
+                    f" Spark session and try again."
+                ),
+                error_code=BAD_REQUEST,
+            ) from e
+
+
+class DeltaTableDataset(_SparkDatasetMixin, _LocationBasedDataset):
+    """
+    Representation of a dataset in delta format with files having the `.delta` extension.
+    """
+
+    def __init__(
+        self,
+        location: str,
+        dataset_format: str,
+        recipe_root: str,
+        version: Optional[int] = None,
+        timestamp: Optional[str] = None,
+    ):
+        """
+        Args:
+            location: The location of the dataset (e.g. '/tmp/myfile.parquet', './mypath',
+                's3://mybucket/mypath', ...).
+            dataset_format: The format of the dataset (e.g. 'csv', 'parquet', ...).
+            recipe_root: The absolute path of the associated recipe root directory on the
+                local filesystem.
+            version: The version of the Delta table to read.
+            timestamp: The timestamp at which to read the Delta table.
+        """
+        super().__init__(location=location, dataset_format=dataset_format, recipe_root=recipe_root)
+        self.version = version
+        self.timestamp = timestamp
+
+    def resolve_to_parquet(self, dst_path: str):
+        spark_session = self._get_or_create_spark_session()
+        spark_read_op = spark_session.read.format("delta")
+        if self.version is not None:
+            spark_read_op = spark_read_op.option("versionAsOf", self.version)
+        if self.timestamp is not None:
+            spark_read_op = spark_read_op.option("timestampAsOf", self.timestamp)
+        spark_df = spark_read_op.load(self.location)
+        pandas_df = self._convert_spark_df_to_pandas(spark_df)
+        write_pandas_df_as_parquet(df=pandas_df, data_parquet_path=dst_path)
+
+    @staticmethod
+    def handles_format(dataset_format: str) -> bool:
+        return dataset_format == "delta"
+
+    @classmethod
+    def _from_config(cls, dataset_config: dict[str, Any], recipe_root: str) -> "_Dataset":
+        return cls(
+            location=cls._get_required_config(dataset_config=dataset_config, key="location"),
+            recipe_root=recipe_root,
+            dataset_format=cls._get_required_config(dataset_config=dataset_config, key="using"),
+            version=dataset_config.get("version"),
+            timestamp=dataset_config.get("timestamp"),
+        )
+
+
+class SparkSqlDataset(_SparkDatasetMixin, _Dataset):
+    """
+    Representation of a Spark SQL dataset defined by a Spark SQL query string
+    (e.g. `SELECT * FROM my_spark_table`).
+    """
+
+    def __init__(self, sql: str, location: str, dataset_format: str):
+        """
+        Args:
+            sql: The Spark SQL query string that defines the dataset
+                (e.g. 'SELECT * FROM my_spark_table').
+            location: The location of the dataset
+                (e.g. 'catalog.schema.table', 'schema.table', 'table').
+            dataset_format: The format of the dataset (e.g. 'csv', 'parquet', ...).
+        """
+        super().__init__(dataset_format=dataset_format)
+        self.sql = sql
+        self.location = location
+
+    def resolve_to_parquet(self, dst_path: str):
+        if self.location is None and self.sql is None:
+            raise MlflowException(
+                "Either location or sql configuration key must be specified for "
+                "dataset with format spark_sql"
+            ) from None
+        spark_session = self._get_or_create_spark_session()
+        spark_df = None
+        if self.sql is not None:
+            spark_df = spark_session.sql(self.sql)
+        elif self.location is not None:
+            spark_df = spark_session.table(self.location)
+        pandas_df = self._convert_spark_df_to_pandas(spark_df)
+        write_pandas_df_as_parquet(df=pandas_df, data_parquet_path=dst_path)
+
+    @classmethod
+    def _from_config(cls, dataset_config: dict[str, Any], recipe_root: str) -> "_Dataset":
+        return cls(
+            sql=dataset_config.get("sql"),
+            location=dataset_config.get("location"),
+            dataset_format=cls._get_required_config(dataset_config=dataset_config, key="using"),
+        )
+
+    @staticmethod
+    def handles_format(dataset_format: str) -> bool:
+        return dataset_format == "spark_sql"

venv/lib/python3.9/site-packages/mlflow/recipes/steps/predict.py

@@ -0,0 +1,288 @@
+import logging
+import os
+import time
+from typing import Any
+
+import mlflow
+from mlflow.exceptions import BAD_REQUEST, INVALID_PARAMETER_VALUE, MlflowException
+from mlflow.recipes.artifacts import DataframeArtifact, RegisteredModelVersionInfo
+from mlflow.recipes.cards import BaseCard
+from mlflow.recipes.step import BaseStep, StepClass
+from mlflow.recipes.steps.register import _REGISTERED_MV_INFO_FILE
+from mlflow.recipes.utils.execution import get_step_output_path
+from mlflow.recipes.utils.step import get_pandas_data_profiles
+from mlflow.recipes.utils.tracking import (
+    TrackingConfig,
+    apply_recipe_tracking_config,
+    get_recipe_tracking_config,
+)
+from mlflow.utils._spark_utils import (
+    _create_local_spark_session_for_recipes,
+    _get_active_spark_session,
+)
+from mlflow.utils.databricks_utils import get_databricks_env_vars
+from mlflow.utils.file_utils import write_spark_dataframe_to_parquet_on_local_disk
+
+_logger = logging.getLogger(__name__)
+
+
+# This should maybe imported from the ingest scoring step for consistency
+_INPUT_FILE_NAME = "scoring-dataset.parquet"
+_SCORED_OUTPUT_FILE_NAME = "scored.parquet"
+_PREDICTION_COLUMN_NAME = "prediction"
+
+# Max dataframe size for profiling after scoring
+_MAX_PROFILE_SIZE = 10000
+# Environment manager for Spark UDF model restoration
+_ENV_MANAGER = "virtualenv"
+
+
+class PredictStep(BaseStep):
+    def __init__(self, step_config: dict[str, Any], recipe_root: str) -> None:
+        super().__init__(step_config, recipe_root)
+        self.tracking_config = TrackingConfig.from_dict(self.step_config)
+
+    def _validate_and_apply_step_config(self):
+        required_configuration_keys = ["using", "location"]
+        for key in required_configuration_keys:
+            if key not in self.step_config:
+                raise MlflowException(
+                    f"The `{key}` configuration key must be specified for the predict step.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+        if self.step_config["using"] not in {"parquet", "delta", "table"}:
+            raise MlflowException(
+                "Invalid `using` in predict step configuration.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        if "model_uri" not in self.step_config:
+            try:
+                register_config = self.step_config["model_registry"]
+                model_name = register_config["model_name"]
+            except KeyError:
+                raise MlflowException(
+                    "No model specified for batch scoring: model_registry does not have "
+                    "`model_uri` and does not have `model_name` configuration key.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+            else:
+                self.step_config["model_uri"] = f"models:/{model_name}/latest"
+        self.registry_uri = self.step_config.get("registry_uri", None)
+        self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
+        self.save_mode = self.step_config.get("save_mode", "overwrite")
+        self.run_end_time = None
+        self.execution_duration = None
+
+    def _build_profiles_and_card(self, scored_sdf) -> BaseCard:
+        # Build profiles for scored dataset
+        card = BaseCard(self.recipe_name, self.name)
+
+        scored_size = scored_sdf.count()
+
+        if not self.skip_data_profiling:
+            _logger.info("Profiling scored dataset")
+            if scored_size > _MAX_PROFILE_SIZE:
+                _logger.info("Sampling scored dataset for profiling because dataset size is large.")
+                sample_percentage = _MAX_PROFILE_SIZE / scored_size
+                scored_sdf = scored_sdf.sample(sample_percentage)
+            scored_df = scored_sdf.toPandas()
+            scored_dataset_profile = get_pandas_data_profiles(
+                [["Profile of Scored Dataset", scored_df]]
+            )
+
+            # Optional tab : data profile for scored data:
+            card.add_tab("Scored Data Profile", "{{PROFILE}}").add_pandas_profile(
+                "PROFILE", scored_dataset_profile
+            )
+
+        # Tab #1/2: run summary.
+        (
+            card.add_tab(
+                "Run Summary",
+                """
+                {{ SCORED_DATA_NUM_ROWS }}
+                {{ EXE_DURATION }}
+                {{ LAST_UPDATE_TIME }}
+                """,
+            ).add_markdown(
+                "SCORED_DATA_NUM_ROWS",
+                f"**Number of scored dataset rows:** `{scored_size}`",
+            )
+        )
+
+        return card
+
+    def _run(self, output_directory):
+        import pandas as pd
+        from pyspark.sql.functions import struct
+
+        run_start_time = time.time()
+
+        apply_recipe_tracking_config(self.tracking_config)
+        if self.registry_uri:
+            mlflow.set_registry_uri(self.registry_uri)
+
+        # Get or create spark session
+        try:
+            spark = _get_active_spark_session()
+            if spark:
+                _logger.info("Found active spark session")
+            else:
+                _logger.info("Creating new spark session")
+                spark = _create_local_spark_session_for_recipes()
+        except Exception as e:
+            raise MlflowException(
+                message=(
+                    "Encountered an error while getting or creating an active Spark session to"
+                    " score dataset with spark UDF."
+                ),
+                error_code=BAD_REQUEST,
+            ) from e
+
+        # read cleaned dataset
+        ingested_data_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name="ingest_scoring",
+            relative_path=_INPUT_FILE_NAME,
+        )
+        # Because the cached parquet file is not on DBFS, we have to first load it as a pandas df
+        input_pdf = pd.read_parquet(ingested_data_path)
+        input_sdf = spark.createDataFrame(input_pdf)
+        if _PREDICTION_COLUMN_NAME in input_sdf.columns:
+            _logger.warning(
+                f"Input scoring dataframe already contains a column '{_PREDICTION_COLUMN_NAME}'. "
+                f"This column will be dropped in favor of the predict output column name."
+            )
+
+        # get model uri
+        model_uri = self.step_config["model_uri"]
+        registered_model_file_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name="register",
+            relative_path=_REGISTERED_MV_INFO_FILE,
+        )
+        if os.path.exists(registered_model_file_path):
+            rmi = RegisteredModelVersionInfo.from_json(path=registered_model_file_path)
+            model_uri = f"models:/{rmi.name}/{rmi.version}"
+
+        # scored dataset
+        result_type = self.step_config.get("result_type", "double")
+        predict = mlflow.pyfunc.spark_udf(
+            spark, model_uri, result_type=result_type, env_manager=_ENV_MANAGER
+        )
+        scored_sdf = input_sdf.withColumn(
+            _PREDICTION_COLUMN_NAME, predict(struct(*input_sdf.columns))
+        )
+
+        # check if output location is already populated for non-delta output formats
+        output_format = self.step_config["using"]
+        output_location = self.step_config["location"]
+        output_populated = False
+        if self.save_mode in ["default", "error", "errorifexists"]:
+            if output_format == "parquet" or output_format == "delta":
+                output_populated = os.path.exists(output_location)
+            else:
+                try:
+                    output_populated = spark._jsparkSession.catalog().tableExists(output_location)
+                except Exception:
+                    # swallow spark failures
+                    pass
+        if output_populated:
+            raise MlflowException(
+                message=(
+                    f"Output location `{output_location}` using format `{output_format}` is "
+                    "already populated. To overwrite, please change the spark `save_mode` in "
+                    "the predict step configuration."
+                ),
+                error_code=BAD_REQUEST,
+            )
+
+        if output_format == "table":
+            try:
+                from delta.tables import DeltaTable
+
+                output_populated = DeltaTable.forName(spark, output_location)
+            except Exception:
+                # swallow spark failures
+                pass
+
+            if output_populated:
+                _logger.info(f"Table already exists at {output_location}")
+                # If the table already exists, we are just setting up the table properties to
+                # ensure that the table can be written with column names with spaces.
+                spark.sql(
+                    f"ALTER TABLE {output_location} SET TBLPROPERTIES "
+                    "('delta.columnMapping.mode'='name','delta.minReaderVersion'='2',"
+                    "'delta.minWriterVersion'='5')"
+                )
+            else:
+                _logger.info(f"Creating a new table at {output_location}")
+                from delta.tables import DeltaTable
+
+                # If the table location specified doesn't exist, we are creating a new table
+                # with properties required to ensure that column names can have spaces.
+                DeltaTable.create().addColumns(scored_sdf.schema).property(
+                    "delta.minReaderVersion", "2"
+                ).property("delta.minWriterVersion", "5").property(
+                    "delta.columnMapping.mode", "name"
+                ).tableName(output_location).execute()
+                # We are overriding the save_mode to append for the create case, since the table
+                # is already created above, so adding any record to the empty table can be
+                # appended to the table
+                self.save_mode = "append"
+
+        # save predictions
+        if output_format in ["parquet", "delta"]:
+            scored_sdf.coalesce(1).write.format(output_format).mode(self.save_mode).save(
+                output_location
+            )
+        else:
+            scored_sdf.write.format("delta").mode(self.save_mode).saveAsTable(output_location)
+
+        # predict step artifacts
+        write_spark_dataframe_to_parquet_on_local_disk(
+            scored_sdf, os.path.join(output_directory, _SCORED_OUTPUT_FILE_NAME)
+        )
+
+        self.run_end_time = time.time()
+        self.execution_duration = self.run_end_time - run_start_time
+        return self._build_profiles_and_card(scored_sdf)
+
+    @classmethod
+    def from_recipe_config(cls, recipe_config, recipe_root):
+        step_config = {}
+        if recipe_config.get("steps", {}).get("predict", {}) is not None:
+            step_config.update(recipe_config.get("steps", {}).get("predict", {}))
+        if recipe_config.get("steps", {}).get("predict", {}).get("output", {}) is not None:
+            step_config.update(recipe_config.get("steps", {}).get("predict", {}).get("output", {}))
+        step_config["register"] = recipe_config.get("steps", {}).get("register", {})
+        step_config["model_registry"] = recipe_config.get("model_registry", {})
+        step_config["recipe"] = recipe_config.get("recipe", "regression/v1")
+        if recipe_config.get("model_registry", {}).get("registry_uri") is not None:
+            step_config["registry_uri"] = recipe_config.get("model_registry", {}).get(
+                "registry_uri"
+            )
+
+        step_config.update(
+            get_recipe_tracking_config(
+                recipe_root_path=recipe_root,
+                recipe_config=recipe_config,
+            ).to_dict()
+        )
+        return cls(step_config, recipe_root)
+
+    @property
+    def name(self):
+        return "predict"
+
+    @property
+    def environment(self):
+        return get_databricks_env_vars(tracking_uri=self.tracking_config.tracking_uri)
+
+    def get_artifacts(self):
+        return [
+            DataframeArtifact("scored_data", self.recipe_root, self.name, _SCORED_OUTPUT_FILE_NAME)
+        ]
+
+    def step_class(self):
+        return StepClass.PREDICTION

venv/lib/python3.9/site-packages/mlflow/recipes/steps/register.py

@@ -0,0 +1,205 @@
+import logging
+from pathlib import Path
+from typing import Any
+
+import mlflow
+from mlflow.entities import SourceType
+from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
+from mlflow.recipes.artifacts import ModelVersionArtifact, RegisteredModelVersionInfo
+from mlflow.recipes.cards import BaseCard
+from mlflow.recipes.step import BaseStep, StepClass
+from mlflow.recipes.steps.train import TrainStep
+from mlflow.recipes.utils.execution import get_step_output_path
+from mlflow.recipes.utils.tracking import (
+    TrackingConfig,
+    apply_recipe_tracking_config,
+    get_recipe_tracking_config,
+)
+from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+from mlflow.utils.databricks_utils import (
+    get_databricks_env_vars,
+    get_databricks_model_version_url,
+    get_databricks_run_url,
+)
+from mlflow.utils.mlflow_tags import MLFLOW_RECIPE_TEMPLATE_NAME, MLFLOW_SOURCE_TYPE
+
+_logger = logging.getLogger(__name__)
+
+_REGISTERED_MV_INFO_FILE = "registered_model_version.json"
+
+
+class RegisterStep(BaseStep):
+    def __init__(self, step_config: dict[str, Any], recipe_root: str):
+        super().__init__(step_config, recipe_root)
+        self.tracking_config = TrackingConfig.from_dict(self.step_config)
+
+    def _validate_and_apply_step_config(self):
+        self.num_dropped_rows = None
+        self.model_uri = None
+        self.model_details = None
+        self.version = None
+
+        self.register_model_name = self.step_config.get("model_name")
+        if self.register_model_name is None:
+            raise MlflowException(
+                "Missing 'model_name' config in register step config.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        self.allow_non_validated_model = self.step_config.get("allow_non_validated_model", False)
+        self.registry_uri = self.step_config.get("registry_uri", None)
+
+    def _run(self, output_directory):
+        apply_recipe_tracking_config(self.tracking_config)
+
+        run_id_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name="train",
+            relative_path="run_id",
+        )
+        run_id = Path(run_id_path).read_text()
+
+        model_validation_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name="evaluate",
+            relative_path="model_validation_status",
+        )
+        model_validation = Path(model_validation_path).read_text()
+        artifact_path = "train/model"
+        tags = {
+            MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.RECIPE),
+            MLFLOW_RECIPE_TEMPLATE_NAME: self.step_config["recipe"],
+        }
+        self.model_uri = f"runs:/{run_id}/{artifact_path}"
+        if model_validation == "VALIDATED" or (
+            model_validation == "UNKNOWN" and self.allow_non_validated_model
+        ):
+            if self.registry_uri:
+                mlflow.set_registry_uri(self.registry_uri)
+            self.model_details = mlflow.register_model(
+                model_uri=self.model_uri,
+                name=self.register_model_name,
+                tags=tags,
+                await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
+            )
+            self.version = self.model_details.version
+            registered_model_info = RegisteredModelVersionInfo(
+                name=self.register_model_name, version=self.version
+            )
+            registered_model_info.to_json(
+                path=str(Path(output_directory) / _REGISTERED_MV_INFO_FILE)
+            )
+        else:
+            raise MlflowException(
+                f"Model registration on {self.model_uri} failed because it "
+                "is not validated. Bypass by setting allow_non_validated_model to True. "
+            )
+
+        card = self._build_card(run_id)
+        card.save_as_html(output_directory)
+        self._log_step_card(run_id, self.name)
+        return card
+
+    def _build_card(self, run_id: str) -> BaseCard:
+        card = BaseCard(self.recipe_name, self.name)
+        card_tab = card.add_tab(
+            "Run Summary",
+            "{{ MODEL_NAME }}"
+            + "{{ MODEL_VERSION }}"
+            + "{{ MODEL_SOURCE_URI }}"
+            + "{{ ALERTS }}"
+            + "{{ EXE_DURATION }}"
+            + "{{ LAST_UPDATE_TIME }}",
+        )
+
+        if self.version is not None:
+            model_version_url = get_databricks_model_version_url(
+                registry_uri=mlflow.get_registry_uri(),
+                name=self.register_model_name,
+                version=self.version,
+            )
+
+            if model_version_url is not None:
+                card_tab.add_html(
+                    "MODEL_NAME",
+                    (
+                        f"<b>Model Name:</b> <a href={model_version_url}>"
+                        f"{self.register_model_name}</a><br><br>"
+                    ),
+                )
+                card_tab.add_html(
+                    "MODEL_VERSION",
+                    (
+                        f"<b>Model Version</b> <a href={model_version_url}>"
+                        f"{self.version}</a><br><br>"
+                    ),
+                )
+            else:
+                card_tab.add_markdown(
+                    "MODEL_NAME",
+                    f"**Model Name:** `{self.register_model_name}`",
+                )
+                card_tab.add_markdown(
+                    "MODEL_VERSION",
+                    f"**Model Version:** `{self.version}`",
+                )
+
+        model_source_url = get_databricks_run_url(
+            tracking_uri=mlflow.get_tracking_uri(),
+            run_id=run_id,
+            artifact_path=f"train/{TrainStep.MODEL_ARTIFACT_RELATIVE_PATH}",
+        )
+
+        if self.model_uri is not None and model_source_url is not None:
+            card_tab.add_html(
+                "MODEL_SOURCE_URI",
+                f"<b>Model Source URI</b> <a href={model_source_url}>{self.model_uri}</a>",
+            )
+        elif self.model_uri is not None:
+            card_tab.add_markdown(
+                "MODEL_SOURCE_URI",
+                f"**Model Source URI:** `{self.model_uri}`",
+            )
+
+        return card
+
+    @classmethod
+    def from_recipe_config(cls, recipe_config, recipe_root):
+        step_config = {}
+        if recipe_config.get("steps", {}).get("register") is not None:
+            step_config.update(recipe_config.get("steps", {}).get("register"))
+        step_config["recipe"] = recipe_config.get("recipe")
+        if recipe_config.get("model_registry", {}).get("registry_uri") is not None:
+            step_config["registry_uri"] = recipe_config.get("model_registry", {}).get(
+                "registry_uri"
+            )
+        if recipe_config.get("model_registry", {}).get("model_name") is not None:
+            step_config["model_name"] = recipe_config.get("model_registry", {}).get("model_name")
+        step_config.update(
+            get_recipe_tracking_config(
+                recipe_root_path=recipe_root,
+                recipe_config=recipe_config,
+            ).to_dict()
+        )
+        return cls(step_config, recipe_root)
+
+    @property
+    def name(self):
+        return "register"
+
+    @property
+    def environment(self):
+        return get_databricks_env_vars(tracking_uri=self.tracking_config.tracking_uri)
+
+    def get_artifacts(self):
+        return [
+            ModelVersionArtifact(
+                "registered_model_version",
+                self.recipe_root,
+                self.name,
+                self.tracking_config.tracking_uri,
+            )
+        ]
+
+    def step_class(self):
+        return StepClass.TRAINING

venv/lib/python3.9/site-packages/mlflow/recipes/steps/split.py

@@ -0,0 +1,482 @@
+import importlib
+import logging
+import os
+import sys
+import time
+from enum import Enum
+from functools import partial
+from multiprocessing.pool import Pool, ThreadPool
+
+import numpy as np
+import pandas as pd
+from packaging.version import Version
+
+from mlflow.exceptions import BAD_REQUEST, INVALID_PARAMETER_VALUE, MlflowException
+from mlflow.recipes.artifacts import DataframeArtifact
+from mlflow.recipes.cards import BaseCard
+from mlflow.recipes.step import BaseStep, StepClass
+from mlflow.recipes.utils.execution import get_step_output_path
+from mlflow.recipes.utils.step import get_pandas_data_profiles, validate_classification_config
+from mlflow.store.artifact.artifact_repo import _NUM_DEFAULT_CPUS
+from mlflow.utils.time import Timer
+
+_logger = logging.getLogger(__name__)
+
+
+_SPLIT_HASH_BUCKET_NUM = 1000
+_INPUT_FILE_NAME = "dataset.parquet"
+_OUTPUT_TRAIN_FILE_NAME = "train.parquet"
+_OUTPUT_VALIDATION_FILE_NAME = "validation.parquet"
+_OUTPUT_TEST_FILE_NAME = "test.parquet"
+_USER_DEFINED_SPLIT_STEP_MODULE = "steps.split"
+_MAX_CLASSES_TO_PROFILE = 5
+
+
+class SplitValues(Enum):
+    """
+    Represents the custom split return values.
+    """
+
+    # Indicates that the row is part of test split
+    TEST = "TEST"
+    # Indicates that the row is part of train split
+    TRAINING = "TRAINING"
+    # Indicates that the row is part of validation split
+    VALIDATION = "VALIDATION"
+
+
+def _make_elem_hashable(elem):
+    if isinstance(elem, list):
+        return tuple(_make_elem_hashable(e) for e in elem)
+    elif isinstance(elem, dict):
+        return tuple((_make_elem_hashable(k), _make_elem_hashable(v)) for k, v in elem.items())
+    elif isinstance(elem, np.ndarray):
+        return elem.shape, tuple(elem.flatten(order="C"))
+    else:
+        return elem
+
+
+def _run_split(task, input_df, split_ratios, target_col):
+    if task == "classification":
+        return _perform_stratified_split_per_class(input_df, split_ratios, target_col)
+    elif task == "regression":
+        return _perform_split(input_df, split_ratios)
+
+
+def _perform_stratified_split_per_class(input_df, split_ratios, target_col):
+    classes = np.unique(input_df[target_col])
+    partial_func = partial(
+        _perform_split_for_one_class,
+        input_df=input_df,
+        split_ratios=split_ratios,
+        target_col=target_col,
+    )
+
+    with ThreadPool(os.cpu_count() or _NUM_DEFAULT_CPUS) as p:
+        zipped_dfs = p.map(partial_func, classes)
+        train_df, validation_df, test_df = [pd.concat(x) for x in list(zip(*zipped_dfs))]
+        return train_df, validation_df, test_df
+
+
+def _perform_split_for_one_class(
+    class_value,
+    input_df,
+    split_ratios,
+    target_col,
+):
+    filtered_df = input_df[input_df[target_col] == class_value]
+    return _perform_split(filtered_df, split_ratios, n_jobs=2)
+
+
+def _perform_split(input_df, split_ratios, n_jobs=-1):
+    hash_buckets = _create_hash_buckets(input_df, n_jobs=n_jobs)
+    train_df, validation_df, test_df = _get_split_df(input_df, hash_buckets, split_ratios)
+    return train_df, validation_df, test_df
+
+
+def _get_split_df(input_df, hash_buckets, split_ratios):
+    # split dataset into train / validation / test splits
+    train_ratio, validation_ratio, test_ratio = split_ratios
+    ratio_sum = train_ratio + validation_ratio + test_ratio
+    train_bucket_end = train_ratio / ratio_sum
+    validation_bucket_end = (train_ratio + validation_ratio) / ratio_sum
+    train_df = input_df[hash_buckets.map(lambda x: x < train_bucket_end)]
+    validation_df = input_df[
+        hash_buckets.map(lambda x: train_bucket_end <= x < validation_bucket_end)
+    ]
+    test_df = input_df[hash_buckets.map(lambda x: x >= validation_bucket_end)]
+
+    empty_splits = [
+        split_name
+        for split_name, split_df in [
+            ("train split", train_df),
+            ("validation split", validation_df),
+            ("test split", test_df),
+        ]
+        if len(split_df) == 0
+    ]
+    if len(empty_splits) > 0:
+        _logger.warning(f"The following input dataset splits are empty: {','.join(empty_splits)}.")
+    return train_df, validation_df, test_df
+
+
+def _parallelize(data, func, n_jobs=-1):
+    n_jobs = n_jobs if 0 < n_jobs <= _NUM_DEFAULT_CPUS else _NUM_DEFAULT_CPUS
+    data_split = np.array_split(data, n_jobs)
+    pool = Pool(n_jobs)
+    data = pd.concat(pool.map(func, data_split))
+    pool.close()
+    pool.join()
+    return data
+
+
+def _run_on_subset(func, data_subset):
+    if Version(pd.__version__) >= Version("2.1.0"):
+        return data_subset.map(func)
+    return data_subset.applymap(func)
+
+
+def _parallelize_on_rows(data, func, n_jobs=-1):
+    return _parallelize(data, partial(_run_on_subset, func), n_jobs=n_jobs)
+
+
+def _hash_pandas_dataframe(input_df, n_jobs=-1):
+    from pandas.util import hash_pandas_object
+
+    hashed_input_df = _parallelize_on_rows(input_df, _make_elem_hashable, n_jobs=n_jobs)
+    return hash_pandas_object(hashed_input_df)
+
+
+def _create_hash_buckets(input_df, n_jobs=-1):
+    # Create hash bucket used for splitting dataset
+    # Note: use `hash_pandas_object` instead of python builtin hash because it is stable
+    # across different process runs / different python versions
+    with Timer() as t:
+        hash_buckets = _hash_pandas_dataframe(input_df, n_jobs=n_jobs).map(
+            lambda x: (x % _SPLIT_HASH_BUCKET_NUM) / _SPLIT_HASH_BUCKET_NUM
+        )
+    _logger.debug(
+        f"Creating hash buckets on input dataset containing {len(input_df)} "
+        f"rows consumes {t:.3f} seconds."
+    )
+    return hash_buckets
+
+
+def _validate_user_code_output(post_split, train_df, validation_df, test_df):
+    try:
+        (
+            post_filter_train_df,
+            post_filter_validation_df,
+            post_filter_test_df,
+        ) = post_split(train_df, validation_df, test_df)
+    except Exception:
+        raise MlflowException(
+            message="Error in cleaning up the data frame post split step."
+            " Expected output is a tuple with (train_df, validation_df, test_df)"
+        ) from None
+
+    for post_split_df, pre_split_df, split_type in [
+        [post_filter_train_df, train_df, "train"],
+        [post_filter_validation_df, validation_df, "validation"],
+        [post_filter_test_df, test_df, "test"],
+    ]:
+        if not isinstance(post_split_df, pd.DataFrame):
+            raise MlflowException(
+                message="The split data is not a DataFrame, please return the correct data."
+            ) from None
+        if list(pre_split_df.columns) != list(post_split_df.columns):
+            raise MlflowException(
+                message="The number of columns post split step are different."
+                f" Column list for {split_type} dataset pre-slit is {list(pre_split_df.columns)}"
+                f" and post split is {list(post_split_df.columns)}. "
+                "Split filter function should be used to filter rows rather than filtering columns."
+            ) from None
+
+    return (
+        post_filter_train_df,
+        post_filter_validation_df,
+        post_filter_test_df,
+    )
+
+
+class SplitStep(BaseStep):
+    def _validate_and_apply_step_config(self):
+        self.run_end_time = None
+        self.execution_duration = None
+        self.num_dropped_rows = None
+
+        self.target_col = self.step_config.get("target_col")
+        self.positive_class = self.step_config.get("positive_class")
+        self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
+        if self.target_col is None:
+            raise MlflowException(
+                "Missing target_col config in recipe config.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
+
+        if "using" in self.step_config:
+            if self.step_config["using"] not in ["custom", "split_ratios"]:
+                raise MlflowException(
+                    f"Invalid split step configuration value {self.step_config['using']} for "
+                    f"key 'using'. Supported values are: ['custom', 'split_ratios']",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+        else:
+            self.step_config["using"] = "split_ratios"
+
+        if self.step_config["using"] == "split_ratios":
+            self.split_ratios = self.step_config.get("split_ratios", [0.75, 0.125, 0.125])
+            if not (
+                isinstance(self.split_ratios, list)
+                and len(self.split_ratios) == 3
+                and all(isinstance(x, (int, float)) and x > 0 for x in self.split_ratios)
+            ):
+                raise MlflowException(
+                    "Config split_ratios must be a list containing 3 positive numbers."
+                )
+
+        if "split_method" not in self.step_config and self.step_config["using"] == "custom":
+            raise MlflowException(
+                "Missing 'split_method' configuration in the split step, which is using 'custom'.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+    def _build_profiles_and_card(self, train_df, validation_df, test_df) -> BaseCard:
+        from sklearn.utils import compute_class_weight
+
+        def _set_target_col_as_first(df, target_col):
+            columns = list(df.columns)
+            col = columns.pop(columns.index(target_col))
+            return df[[col] + columns]
+
+        # Build card
+        card = BaseCard(self.recipe_name, self.name)
+
+        if not self.skip_data_profiling:
+            # Build profiles for input dataset, and train / validation / test splits
+            train_df = _set_target_col_as_first(train_df, self.target_col)
+            validation_df = _set_target_col_as_first(validation_df, self.target_col)
+            test_df = _set_target_col_as_first(test_df, self.target_col)
+            data_profile = get_pandas_data_profiles(
+                [
+                    ["Train", train_df.reset_index(drop=True)],
+                    ["Validation", validation_df.reset_index(drop=True)],
+                    ["Test", test_df.reset_index(drop=True)],
+                ]
+            )
+
+            # Tab #1 - #3: data profiles for train/validation and test.
+            card.add_tab("Compare Splits", "{{PROFILE}}").add_pandas_profile(
+                "PROFILE", data_profile
+            )
+
+            if self.task == "classification":
+                if self.positive_class is not None:
+                    mask = train_df[self.target_col] == self.positive_class
+                    dfs_for_profiles = [
+                        ("Positive", train_df[mask]),
+                        ("Negative", train_df[~mask]),
+                    ]
+                    sub_title = "Positive vs Negative"
+                else:
+                    classes = np.unique(train_df[self.target_col])
+                    class_weights = compute_class_weight(
+                        class_weight="balanced",
+                        classes=classes,
+                        y=train_df[self.target_col],
+                    )
+                    class_weights = list(zip(classes, class_weights))
+                    class_weights = sorted(class_weights, key=lambda x: x[1], reverse=True)
+                    if len(class_weights) > _MAX_CLASSES_TO_PROFILE:
+                        class_weights = class_weights[:_MAX_CLASSES_TO_PROFILE]
+                    dfs_for_profiles = [
+                        (name, train_df[(train_df[self.target_col] == name)])
+                        for name, _ in class_weights
+                    ]
+                    sub_title = f"Top {min(5, len(class_weights))} Classes"
+
+                profiles = [
+                    [
+                        str(p[0]),
+                        p[1].drop(columns=[self.target_col]).reset_index(drop=True),
+                    ]
+                    for p in dfs_for_profiles
+                ]
+                generated_profile = get_pandas_data_profiles(profiles)
+                # Tab #4: data profiles positive negative training split.
+                card.add_tab(
+                    f"Compare Training Data ({sub_title})", "{{PROFILE}}"
+                ).add_pandas_profile("PROFILE", generated_profile)
+
+        # Tab #5: run summary.
+        (
+            card.add_tab(
+                "Run Summary",
+                """
+                {{ SCHEMA_LOCATION }}
+                {{ TRAIN_SPLIT_NUM_ROWS }}
+                {{ VALIDATION_SPLIT_NUM_ROWS }}
+                {{ TEST_SPLIT_NUM_ROWS }}
+                {{ NUM_DROPPED_ROWS }}
+                {{ EXE_DURATION}}
+                {{ LAST_UPDATE_TIME }}
+                """,
+            )
+            .add_markdown(
+                "NUM_DROPPED_ROWS", f"**Number of dropped rows:** `{self.num_dropped_rows}`"
+            )
+            .add_markdown(
+                "TRAIN_SPLIT_NUM_ROWS", f"**Number of train dataset rows:** `{len(train_df)}`"
+            )
+            .add_markdown(
+                "VALIDATION_SPLIT_NUM_ROWS",
+                f"**Number of validation dataset rows:** `{len(validation_df)}`",
+            )
+            .add_markdown(
+                "TEST_SPLIT_NUM_ROWS", f"**Number of test dataset rows:** `{len(test_df)}`"
+            )
+        )
+
+        return card
+
+    def _validate_and_execute_custom_split(self, split_fn, input_df):
+        custom_split_mapping_series = split_fn(input_df)
+        if not isinstance(custom_split_mapping_series, pd.Series):
+            raise MlflowException(
+                "Return type of the custom split function should be a pandas series",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        copy_df = input_df.copy()
+        copy_df["split"] = custom_split_mapping_series
+        train_df = input_df[copy_df["split"] == SplitValues.TRAINING.value].reset_index(drop=True)
+        validation_df = input_df[copy_df["split"] == SplitValues.VALIDATION.value].reset_index(
+            drop=True
+        )
+        test_df = input_df[copy_df["split"] == SplitValues.TEST.value].reset_index(drop=True)
+
+        if train_df.size + validation_df.size + test_df.size != input_df.size:
+            incorrect_args = custom_split_mapping_series[
+                ~custom_split_mapping_series.isin(
+                    [
+                        SplitValues.TRAINING.value,
+                        SplitValues.VALIDATION.value,
+                        SplitValues.TEST.value,
+                    ]
+                )
+            ].unique()
+            raise MlflowException(
+                f"Returned pandas series from custom split step should only contain "
+                f"{SplitValues.TRAINING.value}, {SplitValues.VALIDATION.value} or "
+                f"{SplitValues.TEST.value} as values. Value returned back: {incorrect_args}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        return train_df, validation_df, test_df
+
+    def _run_custom_split(self, input_df):
+        split_fn = getattr(
+            importlib.import_module(_USER_DEFINED_SPLIT_STEP_MODULE),
+            self.step_config["split_method"],
+        )
+        return self._validate_and_execute_custom_split(split_fn, input_df)
+
+    def _run(self, output_directory):
+        run_start_time = time.time()
+
+        # read ingested dataset
+        ingested_data_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name="ingest",
+            relative_path=_INPUT_FILE_NAME,
+        )
+        input_df = pd.read_parquet(ingested_data_path)
+        validate_classification_config(self.task, self.positive_class, input_df, self.target_col)
+
+        # drop rows which target value is missing
+        raw_input_num_rows = len(input_df)
+        # Make sure the target column is actually present in the input DF.
+        if self.target_col not in input_df.columns:
+            raise MlflowException(
+                f"Target column '{self.target_col}' not found in ingested dataset.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        input_df = input_df.dropna(how="any", subset=[self.target_col])
+        self.num_dropped_rows = raw_input_num_rows - len(input_df)
+
+        # split dataset
+        if self.step_config["using"] == "custom":
+            train_df, validation_df, test_df = self._run_custom_split(input_df)
+        else:
+            train_df, validation_df, test_df = _run_split(
+                self.task, input_df, self.split_ratios, self.target_col
+            )
+        # Import from user function module to process dataframes
+        post_split_config = self.step_config.get("post_split_method", None)
+        post_split_filter_config = self.step_config.get("post_split_filter_method", None)
+        if post_split_config is not None:
+            sys.path.append(self.recipe_root)
+            post_split = getattr(
+                importlib.import_module(_USER_DEFINED_SPLIT_STEP_MODULE), post_split_config
+            )
+            _logger.debug(f"Running {post_split_config} on train, validation and test datasets.")
+            (
+                train_df,
+                validation_df,
+                test_df,
+            ) = _validate_user_code_output(post_split, train_df, validation_df, test_df)
+
+        elif post_split_filter_config is not None:
+            sys.path.append(self.recipe_root)
+            post_split_filter = getattr(
+                importlib.import_module(_USER_DEFINED_SPLIT_STEP_MODULE), post_split_filter_config
+            )
+            _logger.debug(
+                f"Running {post_split_filter_config} on train, validation and test datasets."
+            )
+            train_df = train_df[post_split_filter(train_df)]
+
+        if min(len(train_df), len(validation_df), len(test_df)) < 4:
+            raise MlflowException(
+                f"Train, validation, and testing datasets cannot be less than 4 rows. Train has "
+                f"{len(train_df)} rows, validation has {len(validation_df)} rows, and test has "
+                f"{len(test_df)} rows.",
+                error_code=BAD_REQUEST,
+            )
+        # Output train / validation / test splits
+        train_df.to_parquet(os.path.join(output_directory, _OUTPUT_TRAIN_FILE_NAME))
+        validation_df.to_parquet(os.path.join(output_directory, _OUTPUT_VALIDATION_FILE_NAME))
+        test_df.to_parquet(os.path.join(output_directory, _OUTPUT_TEST_FILE_NAME))
+
+        self.run_end_time = time.time()
+        self.execution_duration = self.run_end_time - run_start_time
+        return self._build_profiles_and_card(train_df, validation_df, test_df)
+
+    @classmethod
+    def from_recipe_config(cls, recipe_config, recipe_root):
+        step_config = {}
+        if recipe_config.get("steps", {}).get("split", {}) is not None:
+            step_config.update(recipe_config.get("steps", {}).get("split", {}))
+        step_config["target_col"] = recipe_config.get("target_col")
+        step_config["positive_class"] = recipe_config.get("positive_class")
+        step_config["recipe"] = recipe_config.get("recipe", "regression/v1")
+        return cls(step_config, recipe_root)
+
+    @property
+    def name(self):
+        return "split"
+
+    def get_artifacts(self):
+        return [
+            DataframeArtifact(
+                "training_data", self.recipe_root, self.name, _OUTPUT_TRAIN_FILE_NAME
+            ),
+            DataframeArtifact(
+                "validation_data", self.recipe_root, self.name, _OUTPUT_VALIDATION_FILE_NAME
+            ),
+            DataframeArtifact("test_data", self.recipe_root, self.name, _OUTPUT_TEST_FILE_NAME),
+        ]
+
+    def step_class(self):
+        return StepClass.TRAINING

venv/lib/python3.9/site-packages/mlflow/recipes/steps/transform.py

@@ -0,0 +1,264 @@
+import importlib
+import logging
+import os
+import sys
+import time
+
+import cloudpickle
+from packaging.version import Version
+
+from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
+from mlflow.recipes.artifacts import DataframeArtifact, TransformerArtifact
+from mlflow.recipes.cards import BaseCard
+from mlflow.recipes.step import BaseStep, StepClass
+from mlflow.recipes.utils.execution import get_step_output_path
+from mlflow.recipes.utils.step import get_pandas_data_profiles, validate_classification_config
+from mlflow.recipes.utils.tracking import TrackingConfig, get_recipe_tracking_config
+
+_logger = logging.getLogger(__name__)
+
+_USER_DEFINED_TRANSFORM_STEP_MODULE = "steps.transform"
+
+
+def _generate_feature_names(num_features):
+    max_length = len(str(num_features))
+    return ["f_" + str(i).zfill(max_length) for i in range(num_features)]
+
+
+def _get_output_feature_names(transformer, num_features, input_features):
+    import sklearn
+
+    # `get_feature_names_out` was introduced in scikit-learn 1.0.0.
+    if Version(sklearn.__version__) < Version("1.0.0"):
+        return _generate_feature_names(num_features)
+
+    try:
+        # `get_feature_names_out` fails if `transformer` contains a transformer that doesn't
+        # implement `get_feature_names_out`. For example, `FunctionTransformer` only implements
+        # `get_feature_names_out` when it's instantiated with `feature_names_out`.
+        # In scikit-learn >= 1.1.0, all transformers implement `get_feature_names_out`.
+        # In scikit-learn == 1.0.*, some transformers implement `get_feature_names_out`.
+        return transformer.get_feature_names_out(input_features)
+    except Exception as e:
+        _logger.warning(
+            f"Failed to get output feature names with `get_feature_names_out`: {e}. "
+            "Falling back to using auto-generated feature names."
+        )
+        return _generate_feature_names(num_features)
+
+
+def _validate_user_code_output(transformer_fn):
+    transformer = transformer_fn()
+    if transformer is not None and not (hasattr(transformer, "fit") and callable(transformer.fit)):
+        raise MlflowException(
+            message="The transformer provided doesn't have a fit method."
+        ) from None
+
+    if transformer is not None and not (
+        hasattr(transformer, "transform") and callable(transformer.transform)
+    ):
+        raise MlflowException(
+            message="The transformer provided doesn't have a transform method."
+        ) from None
+
+    return transformer
+
+
+class TransformStep(BaseStep):
+    def __init__(self, step_config, recipe_root):
+        super().__init__(step_config, recipe_root)
+        self.tracking_config = TrackingConfig.from_dict(self.step_config)
+
+    def _validate_and_apply_step_config(self):
+        self.target_col = self.step_config.get("target_col")
+        self.positive_class = self.step_config.get("positive_class")
+        if self.target_col is None:
+            raise MlflowException(
+                "Missing target_col config in recipe config.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        if "using" in self.step_config:
+            if self.step_config["using"] not in ["custom"]:
+                raise MlflowException(
+                    f"Invalid transform step configuration value {self.step_config['using']} for "
+                    f"key 'using'. Supported values are: ['custom']",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+        else:
+            self.step_config["using"] = "custom"
+        self.run_end_time = None
+        self.execution_duration = None
+        self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
+
+    def _run(self, output_directory):
+        import pandas as pd
+
+        run_start_time = time.time()
+
+        train_data_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name="split",
+            relative_path="train.parquet",
+        )
+        train_df = pd.read_parquet(train_data_path)
+        validate_classification_config(self.task, self.positive_class, train_df, self.target_col)
+
+        validation_data_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name="split",
+            relative_path="validation.parquet",
+        )
+        validation_df = pd.read_parquet(validation_data_path)
+
+        sys.path.append(self.recipe_root)
+
+        def get_identity_transformer():
+            from sklearn.pipeline import Pipeline
+            from sklearn.preprocessing import FunctionTransformer
+
+            return Pipeline(steps=[("identity", FunctionTransformer())])
+
+        if "transformer_method" not in self.step_config and self.step_config["using"] == "custom":
+            raise MlflowException(
+                "Missing 'transformer_method' configuration in the transform step, "
+                "which is using 'custom'.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        method_config = self.step_config.get("transformer_method")
+        transformer = None
+        if method_config and self.step_config["using"] == "custom":
+            transformer_fn = getattr(
+                importlib.import_module(_USER_DEFINED_TRANSFORM_STEP_MODULE), method_config
+            )
+            transformer = _validate_user_code_output(transformer_fn)
+        transformer = transformer if transformer else get_identity_transformer()
+        transformer.fit(train_df.drop(columns=[self.target_col]), train_df[self.target_col])
+
+        def transform_dataset(dataset):
+            features = dataset.drop(columns=[self.target_col])
+            transformed_features = transformer.transform(features)
+            if not isinstance(transformed_features, pd.DataFrame):
+                num_features = transformed_features.shape[1]
+                columns = _get_output_feature_names(transformer, num_features, features.columns)
+                transformed_features = pd.DataFrame(transformed_features, columns=columns)
+            transformed_features[self.target_col] = dataset[self.target_col].values
+            return transformed_features
+
+        train_transformed = transform_dataset(train_df)
+        validation_transformed = transform_dataset(validation_df)
+
+        with open(os.path.join(output_directory, "transformer.pkl"), "wb") as f:
+            cloudpickle.dump(transformer, f)
+
+        train_transformed.to_parquet(
+            os.path.join(output_directory, "transformed_training_data.parquet")
+        )
+        validation_transformed.to_parquet(
+            os.path.join(output_directory, "transformed_validation_data.parquet")
+        )
+
+        self.run_end_time = time.time()
+        self.execution_duration = self.run_end_time - run_start_time
+
+        return self._build_profiles_and_card(train_df, train_transformed, transformer)
+
+    def _build_profiles_and_card(self, train_df, train_transformed, transformer) -> BaseCard:
+        # Build card
+        card = BaseCard(self.recipe_name, self.name)
+
+        if not self.skip_data_profiling:
+            # Tab 1: build profiles for train_transformed
+            train_transformed_profile = get_pandas_data_profiles(
+                [["Profile of Train Transformed Dataset", train_transformed]]
+            )
+            card.add_tab("Data Profile (Train Transformed)", "{{PROFILE}}").add_pandas_profile(
+                "PROFILE", train_transformed_profile
+            )
+
+        # Tab 3: transformer diagram
+        from sklearn import set_config
+        from sklearn.utils import estimator_html_repr
+
+        set_config(display="diagram")
+        transformer_repr = estimator_html_repr(transformer)
+        card.add_tab("Transformer", "{{TRANSFORMER}}").add_html("TRANSFORMER", transformer_repr)
+
+        # Tab 4: transformer input schema
+        card.add_tab("Input Schema", "{{INPUT_SCHEMA}}").add_html(
+            "INPUT_SCHEMA",
+            BaseCard.render_table({"Name": n, "Type": t} for n, t in train_df.dtypes.items()),
+        )
+
+        # Tab 5: transformer output schema
+        try:
+            card.add_tab("Output Schema", "{{OUTPUT_SCHEMA}}").add_html(
+                "OUTPUT_SCHEMA",
+                BaseCard.render_table(
+                    {"Name": n, "Type": t} for n, t in train_transformed.dtypes.items()
+                ),
+            )
+        except Exception as e:
+            card.add_tab("Output Schema", "{{OUTPUT_SCHEMA}}").add_html(
+                "OUTPUT_SCHEMA", f"Failed to extract transformer schema. Error: {e}"
+            )
+
+        # Tab 6: transformer output data preview
+        card.add_tab("Data Preview", "{{DATA_PREVIEW}}").add_html(
+            "DATA_PREVIEW", BaseCard.render_table(train_transformed.head())
+        )
+
+        # Tab 7: run summary
+        (
+            card.add_tab(
+                "Run Summary",
+                """
+                {{ EXE_DURATION }}
+                {{ LAST_UPDATE_TIME }}
+                """,
+            )
+        )
+
+        return card
+
+    @classmethod
+    def from_recipe_config(cls, recipe_config, recipe_root):
+        step_config = {}
+        if recipe_config.get("steps", {}).get("transform", {}) is not None:
+            step_config.update(recipe_config.get("steps", {}).get("transform", {}))
+        step_config["target_col"] = recipe_config.get("target_col")
+        step_config["recipe"] = recipe_config.get("recipe", "regression/v1")
+        if "positive_class" in recipe_config:
+            step_config["positive_class"] = recipe_config.get("positive_class")
+        step_config.update(
+            get_recipe_tracking_config(
+                recipe_root_path=recipe_root,
+                recipe_config=recipe_config,
+            ).to_dict()
+        )
+        return cls(step_config, recipe_root)
+
+    @property
+    def name(self):
+        return "transform"
+
+    def get_artifacts(self):
+        return [
+            DataframeArtifact(
+                "transformed_training_data",
+                self.recipe_root,
+                self.name,
+                "transformed_training_data.parquet",
+            ),
+            DataframeArtifact(
+                "transformed_validation_data",
+                self.recipe_root,
+                self.name,
+                "transformed_validation_data.parquet",
+            ),
+            TransformerArtifact(
+                "transformer", self.recipe_root, self.name, self.tracking_config.tracking_uri
+            ),
+        ]
+
+    def step_class(self):
+        return StepClass.TRAINING