star

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

@@ -0,0 +1,390 @@
+import abc
+import json
+import logging
+import os
+import time
+import traceback
+from enum import Enum
+from typing import Any, Optional
+
+import yaml
+
+from mlflow.recipes.cards import CARD_HTML_NAME, CARD_PICKLE_NAME, BaseCard, FailureCard
+from mlflow.recipes.utils import get_recipe_name
+from mlflow.recipes.utils.step import display_html
+from mlflow.tracking import MlflowClient
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+
+_logger = logging.getLogger(__name__)
+
+
+class StepStatus(Enum):
+    """
+    Represents the execution status of a step.
+    """
+
+    # Indicates that no execution status information is available for the step,
+    # which may occur if the step has never been run or its outputs have been cleared
+    UNKNOWN = "UNKNOWN"
+    # Indicates that the step is currently running
+    RUNNING = "RUNNING"
+    # Indicates that the step completed successfully
+    SUCCEEDED = "SUCCEEDED"
+    # Indicates that the step completed with one or more failures
+    FAILED = "FAILED"
+
+
+class StepClass(Enum):
+    """
+    Represents the class of a step.
+    """
+
+    # Indicates that the step class is unknown.
+    UNKNOWN = "UNKNOWN"
+    # Indicates that the step runs at training time.
+    TRAINING = "TRAINING"
+    # Indicates that the step runs at inference time.
+    PREDICTION = "PREDICTION"
+
+
+class StepExecutionState:
+    """
+    Represents execution state for a step, including the current status and
+    the time of the last status update.
+    """
+
+    _KEY_STATUS = "recipe_step_execution_status"
+    _KEY_LAST_UPDATED_TIMESTAMP = "recipe_step_execution_last_updated_timestamp"
+    _KEY_STACK_TRACE = "recipe_step_stack_trace"
+
+    def __init__(self, status: StepStatus, last_updated_timestamp: int, stack_trace: str):
+        """
+        Args:
+            status: The execution status of the step.
+            last_updated_timestamp: The timestamp of the last execution status update, measured
+                in seconds since the UNIX epoch.
+            stack_trace: The stack trace of the last execution. None if the step execution
+                succeeds.
+        """
+        self.status = status
+        self.last_updated_timestamp = last_updated_timestamp
+        self.stack_trace = stack_trace
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Creates a dictionary representation of the step execution state.
+        """
+        return {
+            StepExecutionState._KEY_STATUS: self.status.value,
+            StepExecutionState._KEY_LAST_UPDATED_TIMESTAMP: self.last_updated_timestamp,
+            StepExecutionState._KEY_STACK_TRACE: self.stack_trace,
+        }
+
+    @classmethod
+    def from_dict(cls, state_dict) -> "StepExecutionState":
+        """
+        Creates a ``StepExecutionState`` instance from the specified execution state dictionary.
+        """
+        return cls(
+            status=StepStatus[state_dict[StepExecutionState._KEY_STATUS]],
+            last_updated_timestamp=state_dict[StepExecutionState._KEY_LAST_UPDATED_TIMESTAMP],
+            stack_trace=state_dict[StepExecutionState._KEY_STACK_TRACE],
+        )
+
+
+class BaseStep(metaclass=abc.ABCMeta):
+    """
+    Base class representing a step in an MLflow Recipe
+    """
+
+    _EXECUTION_STATE_FILE_NAME = "execution_state.json"
+
+    def __init__(self, step_config: dict[str, Any], recipe_root: str):
+        """
+        Args:
+            step_config: Dictionary of the config needed to run/implement the step.
+            recipe_root: String file path to the directory where step are defined.
+        """
+        self.step_config = step_config
+        self.recipe_root = recipe_root
+        self.recipe_name = get_recipe_name(recipe_root_path=recipe_root)
+        self.task = self.step_config.get("recipe", "regression/v1").rsplit("/", 1)[0]
+        self.step_card = None
+
+    def __str__(self):
+        return f"Step:{self.name}"
+
+    def run(self, output_directory: str):
+        """
+        Executes the step by running common setup operations and invoking
+        step-specific code (as defined in ``_run()``).
+
+        Args:
+            output_directory: String file path to the directory where step
+                outputs should be stored.
+        """
+        _logger.info(f"Running step {self.name}...")
+        start_timestamp = time.time()
+        self._initialize_databricks_spark_connection_and_hooks_if_applicable()
+        try:
+            self._update_status(status=StepStatus.RUNNING, output_directory=output_directory)
+            self._validate_and_apply_step_config()
+            self.step_card = self._run(output_directory=output_directory)
+            self._update_status(status=StepStatus.SUCCEEDED, output_directory=output_directory)
+        except Exception:
+            stack_trace = traceback.format_exc()
+            self._update_status(
+                status=StepStatus.FAILED, output_directory=output_directory, stack_trace=stack_trace
+            )
+            self.step_card = FailureCard(
+                recipe_name=self.recipe_name,
+                step_name=self.name,
+                failure_traceback=stack_trace,
+                output_directory=output_directory,
+            )
+            raise
+        finally:
+            self._serialize_card(start_timestamp, output_directory)
+
+    def inspect(self, output_directory: str):
+        """
+        Inspect the step output state by running the generic inspect information here and
+        running the step specific inspection code in the step's _inspect() method.
+
+        Args:
+            output_directory: String file path where to the directory where step
+                outputs are located.
+        """
+        card_path = os.path.join(output_directory, CARD_PICKLE_NAME)
+        if not os.path.exists(card_path):
+            _logger.info(
+                "Unable to locate runtime info for step '%s'. Re-run the step before inspect.",
+                self.name,
+            )
+            return None
+
+        card = BaseCard.load(card_path)
+        card_html_path = os.path.join(output_directory, CARD_HTML_NAME)
+        display_html(html_data=card.to_html(), html_file_path=card_html_path)
+
+    @abc.abstractmethod
+    def _run(self, output_directory: str) -> BaseCard:
+        """
+        This function is responsible for executing the step, writing outputs
+        to the specified directory, and returning results to the user. It
+        is invoked by the internal step runner.
+
+        Args:
+            output_directory: String file path to the directory where step outputs
+                should be stored.
+
+        Returns:
+            A BaseCard containing step execution information.
+        """
+
+    @abc.abstractmethod
+    def _validate_and_apply_step_config(self) -> None:
+        """
+        This function is responsible for validating and loading the step config for
+        a particular step. It is invoked by the internal step runner.
+        """
+
+    @classmethod
+    @abc.abstractmethod
+    def from_recipe_config(cls, recipe_config: dict[str, Any], recipe_root: str) -> "BaseStep":
+        """
+        Constructs a step class instance by creating a step config using the recipe
+        config.
+        Subclasses must implement this method to produce the config required to correctly
+        run the corresponding step.
+
+        Args:
+            recipe_config: Dictionary representation of the full recipe config.
+            recipe_root: String file path to the recipe root directory.
+
+        Returns:
+            class instance of the step.
+        """
+
+    @classmethod
+    def from_step_config_path(cls, step_config_path: str, recipe_root: str) -> "BaseStep":
+        """
+        Constructs a step class instance using the config specified in the
+        configuration file.
+
+        Args:
+            step_config_path: String path to the step-specific configuration
+                on the local filesystem.
+            recipe_root: String path to the recipe root directory on
+                the local filesystem.
+
+        Returns:
+            class instance of the step.
+        """
+        with open(step_config_path) as f:
+            step_config = yaml.safe_load(f)
+        return cls(step_config, recipe_root)
+
+    @property
+    @abc.abstractmethod
+    def name(self) -> str:
+        """
+        Returns back the name of the step for the current class instance. This is used
+        downstream by the execution engine to create step-specific directory structures.
+        """
+
+    @property
+    def environment(self) -> dict[str, str]:
+        """
+        Returns environment variables associated with step that should be set when the
+        step is executed.
+        """
+        return {}
+
+    def get_artifacts(self) -> list[Any]:
+        """
+        Returns the named artifacts produced by the step for the current class instance.
+        """
+        return {}
+
+    @abc.abstractmethod
+    def step_class(self) -> StepClass:
+        """
+        Returns the step class.
+        """
+
+    def get_execution_state(self, output_directory: str) -> StepExecutionState:
+        """
+        Returns the execution state of the step, which provides information about its
+        status (succeeded, failed, unknown), last update time, and, if applicable, encountered
+        stacktraces.
+
+        Args:
+            output_directory: String file path to the directory where step
+                outputs are stored.
+
+        Returns:
+            A ``StepExecutionState`` instance containing the step execution state.
+
+        """
+        execution_state_file_path = os.path.join(
+            output_directory, BaseStep._EXECUTION_STATE_FILE_NAME
+        )
+        if os.path.exists(execution_state_file_path):
+            with open(execution_state_file_path) as f:
+                return StepExecutionState.from_dict(json.load(f))
+        else:
+            return StepExecutionState(StepStatus.UNKNOWN, 0, None)
+
+    def _serialize_card(self, start_timestamp: float, output_directory: str) -> None:
+        if self.step_card is None:
+            return
+        execution_duration = time.time() - start_timestamp
+        tab = self.step_card.get_tab("Run Summary")
+        if tab is not None:
+            tab.add_markdown("EXE_DURATION", f"**Run duration (s)**: {execution_duration:.3g}")
+            tab.add_markdown(
+                "LAST_UPDATE_TIME",
+                f"**Last updated:** {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}",
+            )
+        self.step_card.save(path=output_directory)
+        self.step_card.save_as_html(path=output_directory)
+
+    def _update_status(
+        self, status: StepStatus, output_directory: str, stack_trace: Optional[str] = None
+    ) -> None:
+        execution_state = StepExecutionState(
+            status=status, last_updated_timestamp=time.time(), stack_trace=stack_trace
+        )
+        with open(os.path.join(output_directory, BaseStep._EXECUTION_STATE_FILE_NAME), "w") as f:
+            json.dump(execution_state.to_dict(), f)
+
+    def _initialize_databricks_spark_connection_and_hooks_if_applicable(self) -> None:
+        """
+        Initializes a connection to the Databricks Spark Gateway and sets up associated hooks
+        (e.g. MLflow Run creation notification hooks) if MLflow Recipes is running in the
+        Databricks Runtime.
+        """
+        if is_in_databricks_runtime():
+            try:
+                from dbruntime.spark_connection import (
+                    initialize_spark_connection,
+                    is_pinn_mode_enabled,
+                )
+                from IPython.utils.io import capture_output
+
+                with capture_output():
+                    spark_handles, entry_point = initialize_spark_connection(is_pinn_mode_enabled())
+            except Exception as e:
+                _logger.warning(
+                    "Encountered unexpected failure while initializing Spark connection. Spark"
+                    " operations may not succeed. Exception: %s",
+                    e,
+                )
+            else:
+                try:
+                    from dbruntime.MlflowCreateRunHook import get_mlflow_create_run_hook
+
+                    # `get_mlflow_create_run_hook` sets up a patch to trigger a Databricks command
+                    # notification every time an MLflow Run is created. This notification is
+                    # visible to users in notebook environments
+                    get_mlflow_create_run_hook(spark_handles["sc"], entry_point)
+                except Exception as e:
+                    _logger.warning(
+                        "Encountered unexpected failure while setting up Databricks MLflow Run"
+                        " creation hooks. Exception: %s",
+                        e,
+                    )
+
+    def _log_step_card(self, run_id: str, step_name: str) -> None:
+        """
+        Logs a step card as an artifact (destination: <step_name>/card.html) in a specified run.
+        If the step card does not exist, logging is skipped.
+
+        Args:
+            run_id: Run ID to which the step card is logged.
+            step_name: Step name.
+        """
+        from mlflow.recipes.utils.execution import get_step_output_path
+
+        local_card_path = get_step_output_path(
+            recipe_root_path=self.recipe_root,
+            step_name=step_name,
+            relative_path=CARD_HTML_NAME,
+        )
+        if os.path.exists(local_card_path):
+            MlflowClient().log_artifact(run_id, local_card_path, artifact_path=step_name)
+        else:
+            _logger.warning(
+                "Failed to log step card for step %s. Run ID: %s. Card local path: %s",
+                step_name,
+                run_id,
+                local_card_path,
+            )
+
+    @staticmethod
+    def _generate_worst_examples_dataframe(
+        dataframe,
+        predictions,
+        error,
+        target_col,
+        worst_k=10,
+    ):
+        """
+        Generate dataframe containing worst k examples with largest prediction error.
+        Dataframe contains columns of all features, prediction, error, and target_col column.
+        The prediction error is defined as absolute error between target value and
+        prediction value.
+        """
+        import numpy as np
+
+        predictions = np.array(predictions)
+        abs_error = np.absolute(error)
+        worst_k_indexes = np.argsort(abs_error)[::-1][:worst_k]
+        result_df = dataframe.iloc[worst_k_indexes].assign(
+            prediction=predictions[worst_k_indexes],
+            absolute_error=abs_error[worst_k_indexes],
+        )
+        front_columns = ["absolute_error", "prediction", target_col]
+        reordered_columns = front_columns + result_df.columns.drop(front_columns).tolist()
+        return result_df[reordered_columns].reset_index(drop=True)