wehub-resource-sync
27 lines
863 B
Python
27 lines
863 B
Python
from mlflow.tracking import MlflowClient
|
|
|
|
|
|
def yield_artifacts(run_id, path=None):
|
|
"""Yield all artifacts in the specified run"""
|
|
client = MlflowClient()
|
|
for item in client.list_artifacts(run_id, path):
|
|
if item.is_dir:
|
|
yield from yield_artifacts(run_id, item.path)
|
|
else:
|
|
yield item.path
|
|
|
|
|
|
def fetch_logged_data(run_id):
|
|
"""Fetch params, metrics, tags, and artifacts in the specified run"""
|
|
client = MlflowClient()
|
|
data = client.get_run(run_id).data
|
|
# Exclude system tags: https://www.mlflow.org/docs/latest/tracking.html#system-tags
|
|
tags = {k: v for k, v in data.tags.items() if not k.startswith("mlflow.")}
|
|
artifacts = list(yield_artifacts(run_id))
|
|
return {
|
|
"params": data.params,
|
|
"metrics": data.metrics,
|
|
"tags": tags,
|
|
"artifacts": artifacts,
|
|
}
|