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: /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)