graphdoc.data.mlflow_data.mlflow_data_helper module

class graphdoc.data.mlflow_data.mlflow_data_helper.MlflowDataHelper(mlflow_tracking_uri: str | Path, mlflow_tracking_username: str | None = None, mlflow_tracking_password: str | None = None)

Bases: object

__init__(mlflow_tracking_uri: str | Path, mlflow_tracking_username: str | None = None, mlflow_tracking_password: str | None = None)

A helper class for loading and saving models and metadata from mlflow.

Parameters:

• mlflow_tracking_uri (Union[str, Path]) – The uri of the mlflow tracking server.

• mlflow_tracking_username (Optional[str]) – The username for the mlflow tracking server.

• mlflow_tracking_password (Optional[str]) – The password for the mlflow tracking server.

update_auth_env_vars(mlflow_tracking_username: str, mlflow_tracking_password: str)

Update the authentication environment variables.

Parameters:

• mlflow_tracking_username (str) – The username for the mlflow tracking server.

• mlflow_tracking_password (str) – The password for the mlflow tracking server.

set_auth_env_vars()

Set the authentication environment variables.

latest_model_version(model_name: str)

Load the latest version of a model from mlflow.

Parameters:

model_name (str) – The name of the model to load.

Returns:

The loaded model.

model_by_name_and_version(model_name: str, model_version: str)

Load a model from mlflow by name and version.

Parameters:

• model_name (str) – The name of the model to load.

• model_version (str) – The version of the model to load.

Returns:

The loaded model.

model_by_uri(model_uri: str)

Load a model from mlflow by uri.

Parameters:

model_uri (str) – The uri of the model to load.

Returns:

The loaded model.

model_by_args(load_model_args: Dict[str, str])

Given a dictionary of arguments, load a model from mlflow. Ordering is model_by_uri, model_by_name_and_version, latest_model_version.

Parameters:

load_model_args (Dict[str, str]) – A dictionary of arguments.

Returns:

The loaded model.

save_model(model: Signature, model_signature: ModelSignature, model_name: str)

Save a model to mlflow.

Parameters:

• model (dspy.Signature) – The model to save.

• model_signature (ModelSignature) – The signature of the model.

• model_name (str) – The name of the model to save.

run_parameters(run_id: str) → dict[str, Any]

Load the parameters of a run from mlflow.

Parameters:

run_id (str) – The id of the run to load the parameters from.

Returns:

The parameters of the run.