83 lines
2.7 KiB
Python
Executable file
83 lines
2.7 KiB
Python
Executable file
"""Verify that every YOLO smoke task was persisted completely in MLflow."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
|
|
import mlflow
|
|
from mlflow.entities import Run
|
|
from mlflow.tracking import MlflowClient
|
|
|
|
|
|
TASKS = ("detect", "segment", "classify", "pose", "obb")
|
|
REQUIRED_ARTIFACTS = {"weights/best.pt", "weights/last.pt", "results.csv"}
|
|
|
|
|
|
def artifact_paths(client: MlflowClient, run_id: str, path: str = "") -> set[str]:
|
|
result: set[str] = set()
|
|
for artifact in client.list_artifacts(run_id, path):
|
|
if artifact.is_dir:
|
|
result.update(artifact_paths(client, run_id, artifact.path))
|
|
else:
|
|
result.add(artifact.path)
|
|
return result
|
|
|
|
|
|
def latest_task_run(runs: list[Run], task: str) -> Run:
|
|
expected_name = f"{task}-smoke"
|
|
for run in runs:
|
|
if run.data.tags.get("mlflow.runName") == expected_name:
|
|
return run
|
|
raise AssertionError(f"MLflow run not found: {expected_name}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--tracking-uri",
|
|
default="sqlite:///runs/yolo26_mlflow_smoke/mlflow.db",
|
|
)
|
|
parser.add_argument("--experiment", default="yolo26-mlflow-smoke")
|
|
args = parser.parse_args()
|
|
|
|
mlflow.set_tracking_uri(args.tracking_uri)
|
|
client = MlflowClient()
|
|
experiment = client.get_experiment_by_name(args.experiment)
|
|
if experiment is None:
|
|
raise AssertionError(f"MLflow experiment not found: {args.experiment}")
|
|
|
|
runs = client.search_runs(
|
|
[experiment.experiment_id],
|
|
order_by=["start_time DESC"],
|
|
)
|
|
summary: dict[str, object] = {
|
|
"tracking_uri": args.tracking_uri,
|
|
"experiment_id": experiment.experiment_id,
|
|
"artifact_location": experiment.artifact_location,
|
|
"tasks": {},
|
|
}
|
|
task_summary: dict[str, object] = summary["tasks"] # type: ignore[assignment]
|
|
|
|
for task in TASKS:
|
|
run = latest_task_run(runs, task)
|
|
artifacts = artifact_paths(client, run.info.run_id)
|
|
missing = REQUIRED_ARTIFACTS - artifacts
|
|
assert run.info.status == "FINISHED", (task, run.info.status)
|
|
assert run.data.params, f"No parameters logged for {task}"
|
|
assert run.data.metrics, f"No metrics logged for {task}"
|
|
assert not missing, f"Missing artifacts for {task}: {sorted(missing)}"
|
|
task_summary[task] = {
|
|
"run_id": run.info.run_id,
|
|
"status": run.info.status,
|
|
"parameters": len(run.data.params),
|
|
"metrics": len(run.data.metrics),
|
|
"artifact_uri": run.info.artifact_uri,
|
|
"required_artifacts": sorted(REQUIRED_ARTIFACTS),
|
|
}
|
|
|
|
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|