The MLflow client is the programmatic interface for querying experiments, filtering runs, retrieving model artifacts, registering models, and accessing model versions. It turns model selection and lifecycle management into repeatable application code rather than a UI-only task. section 4.1.3
A typical model-selection workflow uses the client to:
- Retrieve an experiment with get_experiment_by_name. section 4.1.3
- Filter and rank runs with search_runs, a metric filter, and an order_by clause. section 4.1.3
- Construct a run-based model URI and register the selected model with register_model. section 4.1.3
- Retrieve model-version information and reference datasets inside production monitoring pipelines. section 6.2.2
Register the highest-scoring eligible run
Create an MlflowClient, find the income-classifier experiment, select the run whose test ROC AUC exceeds 0.8 and ranks highest, build its runs URI, and register it as random-forest-classifier. section 4.1.3
Listing 4.7 Retrieving MLflow runs and registering the model
from mlflow import MlflowClient
mlflow_client = MlflowClient()
experiment_name = "income-classifier"
experiment = mlflow_client.get_experiment_by_name(experiment_name)
run_object = mlflow_client.search_runs(
experiment_ids=experiment.experiment_id,
filter_string="metrics.roc_auc_score_test > 0.8",
max_results=1,
order_by=["metrics.roc_auc_score_test DESC"]
)[0] #1
model_uri = f"runs:/{run_object.info.run_id}/{experiment_name}" #2
mlflow.register_model(
model_uri,
"random-forest-classifier"
) #3
Use the runnable listing to inspect the client calls, metric filter, descending ranking, model URI, and registration name. section 4.1.3
Listing 6.9 KFP component for detecting data drift
def detect_drift(
model_name: str,
model_stage: str,
mlflow_host: str,
minio_host: str,
access_key: str,
secret_key: str,
reference_dataset_name: str,
feature_dataset_path: str,
evidently_workspace_url: str,
evidently_ui_project_name: str,
):
os.environ["AWS_ACCESS_KEY_ID"] = access_key
os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
if not minio_host.startswith("http"):
os.environ["AWS_ENDPOINT_URL"] = "http://" + minio_host
else:
os.environ["AWS_ENDPOINT_URL"] = minio_host
mlflow.set_tracking_uri(mlflow_host)
mlflow_client = MlflowClient(mlflow_host) #1
evidently_workspace = \
RemoteWorkspace(evidently_workspace_url) #2
model_run_id = None
for model in mlflow_client.search_model_versions(f"name='{model_name}'"):
if model.current_stage == model_stage:
model_run_id = model.run_id
break
if not model_run_id:
raise ValueError(
f"Model in stage {model_stage} "
f"not found for {model_name}."
)
run = mlflow.get_run(model_run_id)
dataset_source = None
for dataset_input in run.inputs.dataset_inputs:
for tag in dataset_input.tags:
if tag.value == reference_dataset_name:
dataset_source = json.loads(
dataset_input.dataset.source
)["uri"]
break
if not dataset_source:
raise ValueError(
f"Reference dataset {reference_dataset_name} "
"not found."
)
bucket_name = dataset_source.split("/")[0]
object_name = "/".join(dataset_source.split("/")[1:])
file_path = object_name.split("/")[-1]
download_file_from_minio(
minio_host=minio_host,
access_key=access_key,
secret_key=secret_key,
bucket_name=bucket_name,
object_name=object_name,
file_path=file_path,
) #3
reference_df = pd.read_csv(file_path)
feature_df = pd.read_parquet(feature_dataset_path) #4
feature_df.drop(columns=["user_id"], errors="ignore", inplace=True)
report = Report(
metrics=[
DatasetDriftMetric(),
DatasetMissingValuesMetric(),
ColumnDriftMetric(column_name="Education"),
ColumnSummaryMetric(column_name="Education"),
ColumnDriftMetric(column_name="Marital-Status"),
ColumnSummaryMetric(column_name="Marital-Status"),
ColumnDriftMetric(column_name="Native_country"),
ColumnSummaryMetric(column_name="Native_country"),
ColumnDriftMetric(column_name="Occupation"),
ColumnSummaryMetric(column_name="Occupation"),
ColumnDriftMetric(column_name="Race"),
ColumnSummaryMetric(column_name="Race"),
ColumnDriftMetric(column_name="Relationship"),
ColumnSummaryMetric(column_name="Relationship"),
ColumnDriftMetric(column_name="Sex"),
ColumnSummaryMetric(column_name="Sex"),
ColumnDriftMetric(column_name="Workclass"),
ColumnSummaryMetric(column_name="Workclass"),
],
) #5
report.run(
reference_data=reference_df,
current_data=feature_df
) #6
report_file_path = "drift_report.json"
with open(report_file_path, "w") as f:
f.write(report.json())
print("report written")
project = get_evidently_project(
evidently_workspace,
evidently_ui_project_name
)
print("project retreived")
evidently_workspace.add_report(project.id, report) #7
project.save()
Use the pipeline listing to see the client configured against an MLflow host and used to locate the model run supporting drift detection. section 6.2.2
