wehub-resource-sync
54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
import os
|
|
import tempfile
|
|
from pprint import pprint
|
|
|
|
import mlflow
|
|
from mlflow.artifacts import download_artifacts
|
|
from mlflow.tracking import MlflowClient
|
|
|
|
|
|
def save_text(path, text):
|
|
with open(path, "w") as f:
|
|
f.write(text)
|
|
|
|
|
|
def log_artifacts():
|
|
# Upload artifacts
|
|
with mlflow.start_run() as run, tempfile.TemporaryDirectory() as tmp_dir:
|
|
tmp_path_a = os.path.join(tmp_dir, "a.txt")
|
|
save_text(tmp_path_a, "0")
|
|
tmp_sub_dir = os.path.join(tmp_dir, "dir")
|
|
os.makedirs(tmp_sub_dir)
|
|
tmp_path_b = os.path.join(tmp_sub_dir, "b.txt")
|
|
save_text(tmp_path_b, "1")
|
|
mlflow.log_artifact(tmp_path_a)
|
|
mlflow.log_artifacts(tmp_sub_dir, artifact_path="dir")
|
|
return run.info.run_id
|
|
|
|
|
|
def main():
|
|
assert "MLFLOW_TRACKING_URI" in os.environ
|
|
|
|
# Log artifacts
|
|
run_id1 = log_artifacts()
|
|
# Download artifacts
|
|
client = MlflowClient()
|
|
print("Downloading artifacts")
|
|
pprint(os.listdir(download_artifacts(run_id=run_id1, artifact_path="")))
|
|
pprint(os.listdir(download_artifacts(run_id=run_id1, artifact_path="dir")))
|
|
|
|
# List artifacts
|
|
print("Listing artifacts")
|
|
pprint(client.list_artifacts(run_id1))
|
|
pprint(client.list_artifacts(run_id1, "dir"))
|
|
|
|
# Log artifacts again
|
|
run_id2 = log_artifacts()
|
|
# Delete the run to test `mlflow gc` command
|
|
client.delete_run(run_id2)
|
|
print(f"Deleted run: {run_id2}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|