N-th fix od model DL

This commit is contained in:
Dobromir Popov
2026-07-06 23:41:06 +03:00
parent 4bfdc814e2
commit d83224a62f
5 changed files with 224 additions and 18 deletions

View File

@@ -3836,12 +3836,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"shard_end": shard_end,
})
full_query = urllib.parse.urlencode({"model": model, "full": 1})
full_files = _snapshot_regular_files(snapshot_dir)
file_sizes: dict[str, int] = {}
for rel in set(files) | set(full_files):
try:
file_sizes[rel] = (snapshot_dir / rel).resolve().stat().st_size
except OSError:
continue
return [{
"type": "tracker",
"endpoint": base_url,
"url": f"{base_url}/v1/model-files/download?{query}",
"full_url": f"{base_url}/v1/model-files/download?{full_query}",
"files": files,
"full_files": full_files,
"file_sizes": file_sizes,
}]
def _handle_model_files_download(self, parsed: urllib.parse.ParseResult) -> None:
@@ -3877,6 +3886,42 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if not rel_files:
self._send_json(404, {"error": "no local files matched the assigned shard"})
return
# Single-file mode: ?file=<rel> streams one file with Content-Length so
# clients can verify completeness and retry per file instead of
# restarting one giant tar stream after a network hiccup.
single = params.get("file", [""])[0]
if single:
if single not in set(rel_files):
self._send_json(404, {"error": f"file not in requested shard: {single!r}"})
return
file_path = snapshot_dir / single
try:
# resolve() dereferences HF snapshot symlinks into blobs/.
size = file_path.resolve().stat().st_size
except OSError:
self._send_json(404, {"error": f"file unreadable: {single!r}"})
return
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(size))
self.send_header("X-Meshnet-Model-Source", "tracker")
self.end_headers()
try:
with file_path.open("rb") as f:
while True:
chunk = f.read(1024 * 1024)
if not chunk:
break
self.wfile.write(chunk)
except (BrokenPipeError, ConnectionResetError):
print(
f"model-file download aborted by client "
f"({self.client_address[0]}, model={resolved_name}, file={single})",
flush=True,
)
return
self.send_response(200)
self.send_header("Content-Type", "application/x-tar")
self.send_header("X-Meshnet-Model-Source", "tracker")