wehub-resource-sync
227 lines
7.5 KiB
Python
227 lines
7.5 KiB
Python
import pandas as pd
|
|
|
|
from mlflow.evaluation.evaluation import EvaluationEntity as EvaluationEntity
|
|
from mlflow.evaluation.utils import (
|
|
_get_assessments_dataframe_schema,
|
|
_get_evaluations_dataframe_schema,
|
|
_get_metrics_dataframe_schema,
|
|
_get_tags_dataframe_schema,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, RESOURCE_DOES_NOT_EXIST
|
|
from mlflow.tracking.client import MlflowClient
|
|
|
|
|
|
def get_evaluation(*, run_id: str, evaluation_id: str) -> EvaluationEntity:
|
|
"""
|
|
Retrieves an Evaluation object from an MLflow Run.
|
|
|
|
Args:
|
|
run_id (str): ID of the MLflow Run containing the evaluation.
|
|
evaluation_id (str): The ID of the evaluation.
|
|
|
|
Returns:
|
|
Evaluation: The Evaluation object.
|
|
"""
|
|
client = MlflowClient()
|
|
if not _contains_evaluation_artifacts(client=client, run_id=run_id):
|
|
raise MlflowException(
|
|
"The specified run does not contain any evaluations. "
|
|
"Please log evaluations to the run before retrieving them.",
|
|
error_code=RESOURCE_DOES_NOT_EXIST,
|
|
)
|
|
|
|
evaluations_file = client.download_artifacts(run_id=run_id, path="_evaluations.json")
|
|
evaluations_df = _read_evaluations_dataframe(evaluations_file)
|
|
|
|
assessments_file = client.download_artifacts(run_id=run_id, path="_assessments.json")
|
|
assessments_df = _read_assessments_dataframe(assessments_file)
|
|
|
|
metrics_file = client.download_artifacts(run_id=run_id, path="_metrics.json")
|
|
metrics_df = _read_metrics_dataframe(metrics_file)
|
|
|
|
tags_file = client.download_artifacts(run_id=run_id, path="_tags.json")
|
|
tags_df = _read_tags_dataframe(tags_file)
|
|
|
|
return _get_evaluation_from_dataframes(
|
|
run_id=run_id,
|
|
evaluation_id=evaluation_id,
|
|
evaluations_df=evaluations_df,
|
|
metrics_df=metrics_df,
|
|
assessments_df=assessments_df,
|
|
tags_df=tags_df,
|
|
)
|
|
|
|
|
|
def _contains_evaluation_artifacts(*, client: MlflowClient, run_id: str) -> bool:
|
|
return {"_evaluations.json", "_metrics.json", "_assessments.json", "_tags.json"}.issubset({
|
|
file.path for file in client.list_artifacts(run_id)
|
|
})
|
|
|
|
|
|
def _read_evaluations_dataframe(path: str) -> pd.DataFrame:
|
|
"""
|
|
Reads an evaluations DataFrame from a file.
|
|
|
|
Args:
|
|
path (str): Path to the file.
|
|
|
|
Returns:
|
|
pd.DataFrame: The evaluations DataFrame.
|
|
"""
|
|
schema = _get_evaluations_dataframe_schema()
|
|
return pd.read_json(path, orient="split", dtype=schema, convert_dates=False).replace(
|
|
pd.NA, None
|
|
)
|
|
|
|
|
|
def _read_assessments_dataframe(path: str) -> pd.DataFrame:
|
|
"""
|
|
Reads an assessments DataFrame from a file.
|
|
|
|
Args:
|
|
path (str): Path to the file.
|
|
|
|
Returns:
|
|
pd.DataFrame: The assessments DataFrame.
|
|
"""
|
|
schema = _get_assessments_dataframe_schema()
|
|
return pd.read_json(path, orient="split", dtype=schema, convert_dates=False).replace(
|
|
pd.NA, None
|
|
)
|
|
|
|
|
|
def _read_metrics_dataframe(path: str) -> pd.DataFrame:
|
|
"""
|
|
Reads a metrics DataFrame from a file.
|
|
|
|
Args:
|
|
path (str): Path to the file.
|
|
|
|
Returns:
|
|
pd.DataFrame: The metrics DataFrame.
|
|
"""
|
|
schema = _get_metrics_dataframe_schema()
|
|
return pd.read_json(path, orient="split", dtype=schema, convert_dates=False).replace(
|
|
pd.NA, None
|
|
)
|
|
|
|
|
|
def _read_tags_dataframe(path: str) -> pd.DataFrame:
|
|
"""
|
|
Reads a tags DataFrame from a file.
|
|
|
|
Args:
|
|
path (str): Path to the file.
|
|
|
|
Returns:
|
|
pd.DataFrame: The tags DataFrame.
|
|
"""
|
|
schema = _get_tags_dataframe_schema()
|
|
return pd.read_json(path, orient="split", dtype=schema, convert_dates=False).replace(
|
|
pd.NA, None
|
|
)
|
|
|
|
|
|
def _get_evaluation_from_dataframes(
|
|
*,
|
|
run_id: str,
|
|
evaluation_id: str,
|
|
evaluations_df: pd.DataFrame,
|
|
metrics_df: pd.DataFrame,
|
|
assessments_df: pd.DataFrame,
|
|
tags_df: pd.DataFrame,
|
|
) -> EvaluationEntity:
|
|
"""
|
|
Parses an Evaluation object with the specified evaluation ID from the specified DataFrames.
|
|
"""
|
|
evaluation_row = evaluations_df[evaluations_df["evaluation_id"] == evaluation_id]
|
|
if evaluation_row.empty:
|
|
raise MlflowException(
|
|
f"The specified evaluation ID '{evaluation_id}' does not exist in the run '{run_id}'.",
|
|
error_code=RESOURCE_DOES_NOT_EXIST,
|
|
)
|
|
|
|
evaluations: list[EvaluationEntity] = _dataframes_to_evaluations(
|
|
evaluations_df=evaluation_row,
|
|
metrics_df=metrics_df,
|
|
assessments_df=assessments_df,
|
|
tags_df=tags_df,
|
|
)
|
|
if len(evaluations) != 1:
|
|
raise MlflowException(
|
|
f"Expected to find a single evaluation with ID '{evaluation_id}', but found "
|
|
f"{len(evaluations)} evaluations.",
|
|
error_code=INTERNAL_ERROR,
|
|
)
|
|
|
|
return evaluations[0]
|
|
|
|
|
|
def _dataframes_to_evaluations(
|
|
evaluations_df: pd.DataFrame,
|
|
metrics_df: pd.DataFrame,
|
|
assessments_df: pd.DataFrame,
|
|
tags_df: pd.DataFrame,
|
|
) -> list[EvaluationEntity]:
|
|
"""
|
|
Converts four separate DataFrames (main evaluation data, metrics, assessments, and tags) back
|
|
into a list of Evaluation entities.
|
|
|
|
Args:
|
|
evaluations_df (pd.DataFrame): DataFrame with the main evaluation data
|
|
(excluding assessment and metrics).
|
|
metrics_df (pd.DataFrame): DataFrame with metrics.
|
|
assessments_df (pd.DataFrame): DataFrame with assessments.
|
|
tags_df (pd.DataFrame): DataFrame with tags.
|
|
|
|
Returns:
|
|
List[EvaluationEntity]: A list of Evaluation entities created from the DataFrames.
|
|
"""
|
|
# Group metrics and assessment by evaluation_id
|
|
metrics_by_eval = _group_dataframe_by_evaluation_id(metrics_df)
|
|
assessments_by_eval = _group_dataframe_by_evaluation_id(assessments_df)
|
|
tags_by_eval = _group_dataframe_by_evaluation_id(tags_df)
|
|
|
|
# Convert main DataFrame to list of dictionaries and create Evaluation objects
|
|
evaluations = []
|
|
for eval_dict in evaluations_df.to_dict(orient="records"):
|
|
evaluation_id = eval_dict["evaluation_id"]
|
|
eval_dict["metrics"] = [
|
|
{
|
|
"key": metric["key"],
|
|
"value": metric["value"],
|
|
"timestamp": metric["timestamp"],
|
|
# Evaluation metrics don't have steps, but we're reusing the MLflow Metric
|
|
# class to represent Evaluation metrics as entities in Python for now. Accordingly,
|
|
# we set the step to 0 in order to parse the evaluation metric as an MLflow Metric
|
|
# Python entity
|
|
"step": 0,
|
|
# Also discard the evaluation_id field from the evaluation metric, since this
|
|
# field is not part of the MLflow Metric Python entity
|
|
}
|
|
for metric in metrics_by_eval.get(evaluation_id, [])
|
|
]
|
|
eval_dict["assessments"] = assessments_by_eval.get(evaluation_id, [])
|
|
eval_dict["tags"] = tags_by_eval.get(evaluation_id, [])
|
|
evaluations.append(EvaluationEntity.from_dictionary(eval_dict))
|
|
|
|
return evaluations
|
|
|
|
|
|
def _group_dataframe_by_evaluation_id(df: pd.DataFrame):
|
|
"""
|
|
Groups evaluation dataframe rows by 'evaluation_id'.
|
|
|
|
Args:
|
|
df (pd.DataFrame): DataFrame to group.
|
|
|
|
Returns:
|
|
Dict[str, List]: A dictionary with 'evaluation_id' as keys and lists of entity
|
|
dictionaries as values.
|
|
"""
|
|
grouped = df.groupby("evaluation_id", group_keys=False).apply(
|
|
lambda x: x.to_dict(orient="records")
|
|
)
|
|
return grouped.to_dict()
|