graphdoc.data.mlflow_data package

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

Bases: object

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

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)[source]

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()[source]

Set the authentication environment variables.

latest_model_version(model_name: str)[source]

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)[source]

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)[source]

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])[source]

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)[source]

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][source]

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.

Submodules