Spaces:
Sleeping
Sleeping
File size: 3,718 Bytes
a36238b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | from collections.abc import Callable
from pathlib import Path
from typing import Any
from gradio_client import handle_file
from trackio.sqlite_storage import SQLiteStorage
from trackio.typehints import ARTIFACT_BLOB_UPLOAD_KIND
def classify_pending_uploads(buffered: dict) -> dict:
"""Partition `buffered` (`{"uploads": [...], "ids": [...]}` from
`SQLiteStorage.get_pending_uploads`) by kind.
"""
media: list[tuple[dict, int]] = []
artifact_blobs: list[tuple[dict, int]] = []
missing: dict = {"paths": [], "ids": []}
for upload, upload_id in zip(buffered["uploads"], buffered["ids"]):
fp = upload["file_path"]
if not Path(fp).exists():
missing["paths"].append(fp)
missing["ids"].append(upload_id)
elif upload.get("kind") == ARTIFACT_BLOB_UPLOAD_KIND:
artifact_blobs.append((upload, upload_id))
else:
media.append((upload, upload_id))
return {"media": media, "artifact_blobs": artifact_blobs, "missing": missing}
def _media_upload_entry(upload: dict) -> dict:
return {
"project": upload["project"],
"run": upload["run"],
"run_id": upload.get("run_id"),
"step": upload["step"],
"relative_path": upload["relative_path"],
"uploaded_file": handle_file(upload["file_path"]),
}
def _artifact_blob_upload_entry(upload: dict) -> dict:
return {
"project": upload["project"],
"digest": upload["digest"],
"uploaded_file": handle_file(upload["file_path"]),
}
def group_pending_uploads(buffered: dict) -> dict:
"""Shape classified rows for the gradio `predict` endpoints."""
classified = classify_pending_uploads(buffered)
media: dict = {"entries": [], "ids": []}
for upload, upload_id in classified["media"]:
media["entries"].append(_media_upload_entry(upload))
media["ids"].append(upload_id)
artifact_blobs: dict[str, dict] = {}
for upload, upload_id in classified["artifact_blobs"]:
group = artifact_blobs.setdefault(upload["project"], {"entries": [], "ids": []})
group["entries"].append(_artifact_blob_upload_entry(upload))
group["ids"].append(upload_id)
return {
"media": media,
"artifact_blobs": artifact_blobs,
"missing": classified["missing"],
}
def replay_pending_uploads(
buffered: dict,
project: str,
*,
predict: Callable[..., Any],
hf_token: str | None,
warn_missing: Callable[[int, str], None],
verbose: bool = False,
) -> None:
"""Route grouped `pending_uploads` rows to their endpoints, clearing each
group's rows as soon as it is sent.
"""
grouped = group_pending_uploads(buffered)
missing = grouped["missing"]
if missing["ids"]:
warn_missing(len(missing["ids"]), missing["paths"][0])
SQLiteStorage.clear_pending_uploads(project, missing["ids"])
media = grouped["media"]
if media["entries"]:
if verbose:
print(f" Syncing {len(media['entries'])} media files...")
predict(
api_name="/bulk_upload_media",
uploads=media["entries"],
hf_token=hf_token,
)
SQLiteStorage.clear_pending_uploads(project, media["ids"])
for proj, group in grouped["artifact_blobs"].items():
if verbose:
print(
f" Syncing {len(group['entries'])} artifact blobs for project '{proj}'..."
)
predict(
api_name="/bulk_upload_artifact_blob",
project=proj,
uploads=group["entries"],
hf_token=hf_token,
)
SQLiteStorage.clear_pending_uploads(project, group["ids"])
|