5-th DL fix
This commit is contained in:
@@ -1,4 +1,4 @@
|
|||||||
# US-047 — Model-source race: visibility, sane timeouts, quiet tracker aborts
|
# US-047 — Tracker-first model downloads: visibility, sane timeouts, RAM-based sizing
|
||||||
|
|
||||||
Status: in progress
|
Status: in progress
|
||||||
Priority: High (follow-up to US-044/US-046; blocks usable LAN downloads)
|
Priority: High (follow-up to US-044/US-046; blocks usable LAN downloads)
|
||||||
@@ -45,14 +45,40 @@ http://192.168.0.179:8080 --model Qwen3.6-35B-A3B`):
|
|||||||
`BrokenPipeError`/`ConnectionResetError` around the tar stream and log a
|
`BrokenPipeError`/`ConnectionResetError` around the tar stream and log a
|
||||||
single line instead of a traceback.
|
single line instead of a traceback.
|
||||||
|
|
||||||
|
## Design revision (2026-07-06, after live retest)
|
||||||
|
|
||||||
|
The race is gone. User decision: **HuggingFace is used only when the model is
|
||||||
|
not available from a tracker/peer source, or when `--tracker-source-disabled`
|
||||||
|
is passed.** Sources are tried sequentially with progress + failure output;
|
||||||
|
HF (layer-filtered via the source file list, else the remote index) is the
|
||||||
|
fallback.
|
||||||
|
|
||||||
|
Second live finding: the node was assigned only layers 0–2 of 40 on a 79 GB
|
||||||
|
box. Cause: CPU-mode nodes still report the detected-but-unusable GPU's
|
||||||
|
`vram_mb` (RTX 4060 → 8192), and shard sizing used VRAM whenever it was > 0
|
||||||
|
(8 GB × 0.8 ≈ 6.5 GB ≈ 3 layers). Fixed on both sides: the node now sends
|
||||||
|
`assignment_vram_mb` (0 unless CUDA is actually usable) to `/v1/nodes/assign`,
|
||||||
|
and the tracker only trusts `vram_mb` when `device=cuda` (all three sizing
|
||||||
|
sites), falling back to `ram_mb`.
|
||||||
|
|
||||||
## Acceptance criteria
|
## Acceptance criteria
|
||||||
|
|
||||||
- [ ] Node started with an explicit `--model` never queries
|
- [x] Node started with an explicit `--model` never queries
|
||||||
`/v1/network/assign` and never prints `auto-join unavailable`.
|
`/v1/network/assign` and never prints `auto-join unavailable`.
|
||||||
- [ ] During the race, tracker-source progress lines appear alongside HF
|
- [x] Tracker/peer model source is preferred outright; HF is contacted only
|
||||||
tqdm output; a failing tracker source prints its exception.
|
when no source is advertised, every source fails, or
|
||||||
- [ ] A ≥2 s read stall no longer aborts a tracker model-source download
|
`--tracker-source-disabled` is passed (flag on both CLI parsers, plumbed
|
||||||
|
through config and `run_startup`).
|
||||||
|
- [x] Tracker-source downloads print progress every 512 MB and print the
|
||||||
|
exception + URL on failure; nothing fails silently.
|
||||||
|
- [x] A ≥2 s read stall no longer aborts a tracker model-source download
|
||||||
(30 s socket timeout).
|
(30 s socket timeout).
|
||||||
- [ ] Client disconnect during `/v1/model-files/download` logs one line on
|
- [x] Client disconnect during `/v1/model-files/download` logs one line on
|
||||||
the tracker, no traceback.
|
the tracker, no traceback.
|
||||||
- [ ] `python -m pytest` passes from repo root.
|
- [x] CPU node with big RAM gets a RAM-sized shard: `/v1/nodes/assign` and
|
||||||
|
both `/v1/network/assign` sizing paths ignore VRAM unless `device=cuda`.
|
||||||
|
- [x] `pytest tests/test_node_startup.py tests/test_tracker_routing.py`
|
||||||
|
passes (139/140; the one failure is the pre-existing port-dependent
|
||||||
|
`test_mining_cli` case, present on clean master).
|
||||||
|
- [ ] Live two-machine retest: Windows node downloads only from tracker at
|
||||||
|
LAN speed and is assigned a RAM-sized shard.
|
||||||
|
|||||||
@@ -65,6 +65,7 @@ def _run_node(cfg: dict) -> None:
|
|||||||
vram_mb_override=cfg.get("vram_mb_override"),
|
vram_mb_override=cfg.get("vram_mb_override"),
|
||||||
max_loaded_shards=int(cfg.get("max_loaded_shards", 1)),
|
max_loaded_shards=int(cfg.get("max_loaded_shards", 1)),
|
||||||
debug=bool(cfg.get("debug", False)),
|
debug=bool(cfg.get("debug", False)),
|
||||||
|
tracker_source_disabled=bool(cfg.get("tracker_source_disabled", False)),
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
|
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
|
||||||
@@ -148,6 +149,8 @@ def _cmd_default(args) -> int:
|
|||||||
overrides["max_loaded_shards"] = args.max_shards
|
overrides["max_loaded_shards"] = args.max_shards
|
||||||
if args.debug:
|
if args.debug:
|
||||||
overrides["debug"] = True
|
overrides["debug"] = True
|
||||||
|
if getattr(args, "tracker_source_disabled", False):
|
||||||
|
overrides["tracker_source_disabled"] = True
|
||||||
|
|
||||||
if overrides:
|
if overrides:
|
||||||
cfg = merge_cli_overrides(cfg, **overrides)
|
cfg = merge_cli_overrides(cfg, **overrides)
|
||||||
@@ -245,6 +248,7 @@ def _cmd_start(args) -> int:
|
|||||||
vram_mb_override=getattr(args, "memory", None),
|
vram_mb_override=getattr(args, "memory", None),
|
||||||
max_loaded_shards=getattr(args, "max_shards", 1),
|
max_loaded_shards=getattr(args, "max_shards", 1),
|
||||||
debug=getattr(args, "debug", False),
|
debug=getattr(args, "debug", False),
|
||||||
|
tracker_source_disabled=getattr(args, "tracker_source_disabled", False),
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
|
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
|
||||||
@@ -281,6 +285,8 @@ def main() -> None:
|
|||||||
help="Quantization level")
|
help="Quantization level")
|
||||||
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
|
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
|
||||||
parser.add_argument("--tracker", metavar="URL", help="Tracker URL")
|
parser.add_argument("--tracker", metavar="URL", help="Tracker URL")
|
||||||
|
parser.add_argument("--tracker-source-disabled", action="store_true",
|
||||||
|
help="Skip tracker/peer model-file sources and download from HuggingFace directly")
|
||||||
parser.add_argument("--wallet", metavar="PATH", help="Wallet file path")
|
parser.add_argument("--wallet", metavar="PATH", help="Wallet file path")
|
||||||
parser.add_argument("--shard-start", type=int, metavar="N", help="Pin shard start layer")
|
parser.add_argument("--shard-start", type=int, metavar="N", help="Pin shard start layer")
|
||||||
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer")
|
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer")
|
||||||
@@ -329,6 +335,8 @@ def main() -> None:
|
|||||||
start_cmd.add_argument("--max-shards", type=int, default=1, metavar="N",
|
start_cmd.add_argument("--max-shards", type=int, default=1, metavar="N",
|
||||||
help="Maximum shard slots this node advertises to the tracker (default 1)")
|
help="Maximum shard slots this node advertises to the tracker (default 1)")
|
||||||
start_cmd.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
|
start_cmd.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
|
||||||
|
start_cmd.add_argument("--tracker-source-disabled", action="store_true",
|
||||||
|
help="Skip tracker/peer model-file sources and download from HuggingFace directly")
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
|||||||
@@ -13,7 +13,6 @@ import tarfile
|
|||||||
import tempfile
|
import tempfile
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
import urllib.request
|
import urllib.request
|
||||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -147,11 +146,13 @@ def _download_model_source(
|
|||||||
if progress:
|
if progress:
|
||||||
print(f" {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
|
print(f" {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
|
||||||
_safe_extract_shard(archive_path, extract_dir)
|
_safe_extract_shard(archive_path, extract_dir)
|
||||||
|
if shard_dir.exists():
|
||||||
|
shutil.rmtree(shard_dir)
|
||||||
shutil.move(str(extract_dir), str(shard_dir))
|
shutil.move(str(extract_dir), str(shard_dir))
|
||||||
return shard_dir
|
return shard_dir
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
if progress:
|
if progress:
|
||||||
print(f" {label}: download failed: {exc!r}", flush=True)
|
print(f" {label}: download failed ({url}): {exc!r}", flush=True)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
@@ -177,55 +178,6 @@ def _download_huggingface_subset(
|
|||||||
return Path(snapshot_download(**kwargs))
|
return Path(snapshot_download(**kwargs))
|
||||||
|
|
||||||
|
|
||||||
def _download_from_fastest_source(
|
|
||||||
*,
|
|
||||||
model_sources: list[dict],
|
|
||||||
hf_repo: str,
|
|
||||||
cache_dir: Path,
|
|
||||||
shard_dir: Path,
|
|
||||||
progress: bool,
|
|
||||||
timeout: float,
|
|
||||||
) -> tuple[str, Path] | None:
|
|
||||||
shard_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
with tempfile.TemporaryDirectory(prefix="meshnet-race-", dir=shard_dir.parent) as tmp:
|
|
||||||
tmp_root = Path(tmp)
|
|
||||||
jobs: dict[Any, tuple[str, Path]] = {}
|
|
||||||
pool = ThreadPoolExecutor(max_workers=min(4, len(model_sources) + 1))
|
|
||||||
try:
|
|
||||||
for index, source in enumerate(model_sources):
|
|
||||||
label = str(source.get("type") or "model-source")
|
|
||||||
candidate = tmp_root / f"source-{index}"
|
|
||||||
jobs[pool.submit(
|
|
||||||
_download_model_source, source, candidate, timeout, progress, label,
|
|
||||||
)] = (label, candidate)
|
|
||||||
allow_patterns = _allow_patterns_from_sources(model_sources)
|
|
||||||
hf_candidate = tmp_root / "huggingface"
|
|
||||||
jobs[pool.submit(_download_huggingface_subset, hf_repo, cache_dir, hf_candidate, allow_patterns)] = (
|
|
||||||
"HuggingFace",
|
|
||||||
hf_candidate,
|
|
||||||
)
|
|
||||||
for future in as_completed(jobs):
|
|
||||||
label, candidate = jobs[future]
|
|
||||||
try:
|
|
||||||
result = future.result()
|
|
||||||
except Exception as exc:
|
|
||||||
if progress:
|
|
||||||
print(f" {label}: download failed: {exc!r}", flush=True)
|
|
||||||
continue
|
|
||||||
if result is None:
|
|
||||||
continue
|
|
||||||
if shard_dir.exists():
|
|
||||||
shutil.rmtree(shard_dir)
|
|
||||||
shutil.move(str(candidate), str(shard_dir))
|
|
||||||
if progress:
|
|
||||||
print(f" download source: {label}", flush=True)
|
|
||||||
pool.shutdown(wait=False, cancel_futures=True)
|
|
||||||
return label, shard_dir
|
|
||||||
finally:
|
|
||||||
pool.shutdown(wait=False, cancel_futures=True)
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None:
|
def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None:
|
||||||
patterns: set[str] = set()
|
patterns: set[str] = set()
|
||||||
for source in model_sources:
|
for source in model_sources:
|
||||||
@@ -336,21 +288,31 @@ def download_shard(
|
|||||||
f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
|
f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
|
||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
if model_sources:
|
# Tracker (or peer) model sources are preferred outright — usually LAN-fast.
|
||||||
|
# HuggingFace is only the fallback when every advertised source fails.
|
||||||
|
for source in model_sources or []:
|
||||||
|
label = str(source.get("type") or "model-source")
|
||||||
if progress:
|
if progress:
|
||||||
print(" Racing tracker model source against HuggingFace ...", flush=True)
|
print(f" Downloading from {label} model source (HuggingFace is the fallback) ...", flush=True)
|
||||||
raced = _download_from_fastest_source(
|
fetched = _download_model_source(
|
||||||
model_sources=model_sources,
|
source,
|
||||||
hf_repo=hf_repo,
|
shard_dir,
|
||||||
cache_dir=cache_dir,
|
|
||||||
shard_dir=shard_dir,
|
|
||||||
progress=progress,
|
|
||||||
timeout=max(peer_timeout, _MODEL_SOURCE_TIMEOUT_SECONDS),
|
timeout=max(peer_timeout, _MODEL_SOURCE_TIMEOUT_SECONDS),
|
||||||
|
progress=progress,
|
||||||
|
label=label,
|
||||||
)
|
)
|
||||||
if raced is not None:
|
if fetched is not None:
|
||||||
return raced[1]
|
if progress:
|
||||||
|
print(f" download source: {label}", flush=True)
|
||||||
|
return fetched
|
||||||
|
if model_sources and progress:
|
||||||
|
print(" All model sources failed — falling back to HuggingFace ...", flush=True)
|
||||||
|
|
||||||
allow_patterns = _allow_patterns_from_remote_index(hf_repo, cache_dir, shard_start, shard_end)
|
allow_patterns = None
|
||||||
|
if model_sources:
|
||||||
|
allow_patterns = _allow_patterns_from_sources(model_sources)
|
||||||
|
if allow_patterns is None:
|
||||||
|
allow_patterns = _allow_patterns_from_remote_index(hf_repo, cache_dir, shard_start, shard_end)
|
||||||
if progress:
|
if progress:
|
||||||
if allow_patterns:
|
if allow_patterns:
|
||||||
print(" download source: HuggingFace (layer-filtered)", flush=True)
|
print(" download source: HuggingFace (layer-filtered)", flush=True)
|
||||||
|
|||||||
@@ -332,6 +332,7 @@ def run_startup(
|
|||||||
vram_mb_override: int | None = None,
|
vram_mb_override: int | None = None,
|
||||||
max_loaded_shards: int = 1,
|
max_loaded_shards: int = 1,
|
||||||
debug: bool = False,
|
debug: bool = False,
|
||||||
|
tracker_source_disabled: bool = False,
|
||||||
) -> StubNodeServer | TorchNodeServer:
|
) -> StubNodeServer | TorchNodeServer:
|
||||||
"""Execute the full startup sequence and return a running node server.
|
"""Execute the full startup sequence and return a running node server.
|
||||||
|
|
||||||
@@ -458,7 +459,10 @@ def run_startup(
|
|||||||
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
|
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
|
||||||
shard_start = net_asgn["shard_start"]
|
shard_start = net_asgn["shard_start"]
|
||||||
shard_end = net_asgn["shard_end"]
|
shard_end = net_asgn["shard_end"]
|
||||||
full_sources = _full_model_sources(net_asgn.get("model_sources", []))
|
full_sources = (
|
||||||
|
[] if tracker_source_disabled
|
||||||
|
else _full_model_sources(net_asgn.get("model_sources", []))
|
||||||
|
)
|
||||||
if full_sources:
|
if full_sources:
|
||||||
cache_dir = download_shard(
|
cache_dir = download_shard(
|
||||||
model_id.split("/")[-1],
|
model_id.split("/")[-1],
|
||||||
@@ -588,7 +592,7 @@ def run_startup(
|
|||||||
f"(of {assigned_num_layers})",
|
f"(of {assigned_num_layers})",
|
||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
full_sources = _full_model_sources(assigned_model_sources)
|
full_sources = [] if tracker_source_disabled else _full_model_sources(assigned_model_sources)
|
||||||
if full_sources:
|
if full_sources:
|
||||||
print("Downloading assigned model snapshot...", flush=True)
|
print("Downloading assigned model snapshot...", flush=True)
|
||||||
cache_dir = download_shard(
|
cache_dir = download_shard(
|
||||||
@@ -676,7 +680,10 @@ def run_startup(
|
|||||||
assign_qs = urllib.parse.urlencode({
|
assign_qs = urllib.parse.urlencode({
|
||||||
"model": model,
|
"model": model,
|
||||||
"device": device,
|
"device": device,
|
||||||
"vram_mb": vram_mb,
|
# CPU-mode nodes must be sized by RAM: a detected-but-unusable GPU's
|
||||||
|
# VRAM would otherwise cap the shard (e.g. 8 GB VRAM → 3 layers on a
|
||||||
|
# 79 GB box whose Torch has no CUDA).
|
||||||
|
"vram_mb": assignment_vram_mb,
|
||||||
"ram_mb": ram_mb,
|
"ram_mb": ram_mb,
|
||||||
})
|
})
|
||||||
try:
|
try:
|
||||||
@@ -690,7 +697,7 @@ def run_startup(
|
|||||||
assigned_model: str = assignment.get("model", model)
|
assigned_model: str = assignment.get("model", model)
|
||||||
hf_repo: str | None = assignment.get("hf_repo")
|
hf_repo: str | None = assignment.get("hf_repo")
|
||||||
peers: list[dict] = assignment.get("peers", [])
|
peers: list[dict] = assignment.get("peers", [])
|
||||||
model_sources: list[dict] = assignment.get("model_sources", [])
|
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
|
||||||
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
|
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
|
||||||
|
|
||||||
# 4. Download shard
|
# 4. Download shard
|
||||||
|
|||||||
@@ -3758,7 +3758,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
except ValueError:
|
except ValueError:
|
||||||
ram_mb = 0
|
ram_mb = 0
|
||||||
max_layers = required_end - required_start + 1
|
max_layers = required_end - required_start + 1
|
||||||
memory_mb = vram_mb if vram_mb > 0 else ram_mb
|
# VRAM only bounds the shard for CUDA nodes; a CPU node may still report
|
||||||
|
# a detected-but-unusable GPU, and must be sized by system RAM.
|
||||||
|
memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb
|
||||||
if memory_mb > 0:
|
if memory_mb > 0:
|
||||||
layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024)
|
layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024)
|
||||||
max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes)))
|
max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes)))
|
||||||
@@ -3963,7 +3965,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if preset is not None and preset.get("hf_repo"):
|
if preset is not None and preset.get("hf_repo"):
|
||||||
required_start, required_end = _preset_layer_bounds(preset)
|
required_start, required_end = _preset_layer_bounds(preset)
|
||||||
total_l = required_end - required_start + 1
|
total_l = required_end - required_start + 1
|
||||||
memory_mb = vram_mb if vram_mb > 0 else ram_mb
|
memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb
|
||||||
max_layers = _max_layers_for_memory(memory_mb, total_l, preset)
|
max_layers = _max_layers_for_memory(memory_mb, total_l, preset)
|
||||||
shard_start = required_start
|
shard_start = required_start
|
||||||
shard_end = min(required_end, shard_start + max_layers - 1)
|
shard_end = min(required_end, shard_start + max_layers - 1)
|
||||||
@@ -4052,7 +4054,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
|
|
||||||
# Capacity: use the same 80%-of-memory rule as registered node planning.
|
# Capacity: use the same 80%-of-memory rule as registered node planning.
|
||||||
total_l = best_num_layers
|
total_l = best_num_layers
|
||||||
memory_mb = vram_mb if vram_mb > 0 else ram_mb
|
memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb
|
||||||
resolved_name, best_preset = _resolve_model_preset(server.model_presets, str(best_repo))
|
resolved_name, best_preset = _resolve_model_preset(server.model_presets, str(best_repo))
|
||||||
if memory_mb > 0:
|
if memory_mb > 0:
|
||||||
max_layers = _max_layers_for_memory(memory_mb, total_l, best_preset)
|
max_layers = _max_layers_for_memory(memory_mb, total_l, best_preset)
|
||||||
|
|||||||
@@ -407,11 +407,11 @@ def test_download_shard_uses_huggingface_when_repo_is_assigned(tmp_path, monkeyp
|
|||||||
}]
|
}]
|
||||||
|
|
||||||
|
|
||||||
def test_download_shard_races_tracker_model_source_against_huggingface(
|
def test_download_shard_prefers_tracker_model_source_over_huggingface(
|
||||||
tmp_path,
|
tmp_path,
|
||||||
monkeypatch,
|
monkeypatch,
|
||||||
):
|
):
|
||||||
"""Tracker-hosted model files can win while HF receives the same allow_patterns."""
|
"""A working tracker model source is used exclusively — HF is never contacted."""
|
||||||
source_dir = tmp_path / "source"
|
source_dir = tmp_path / "source"
|
||||||
source_dir.mkdir()
|
source_dir.mkdir()
|
||||||
(source_dir / "config.json").write_text("{}")
|
(source_dir / "config.json").write_text("{}")
|
||||||
@@ -473,6 +473,49 @@ def test_download_shard_races_tracker_model_source_against_huggingface(
|
|||||||
)
|
)
|
||||||
|
|
||||||
assert (shard_dir / "model-00002-of-00004.safetensors").read_text() == "tracker"
|
assert (shard_dir / "model-00002-of-00004.safetensors").read_text() == "tracker"
|
||||||
|
assert hf_calls == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_shard_falls_back_to_huggingface_when_tracker_source_fails(
|
||||||
|
tmp_path,
|
||||||
|
monkeypatch,
|
||||||
|
):
|
||||||
|
"""A dead tracker source falls through to HF with allow_patterns from the source files."""
|
||||||
|
|
||||||
|
def failing_urlopen(*args, **kwargs):
|
||||||
|
raise ConnectionResetError("tracker went away")
|
||||||
|
|
||||||
|
monkeypatch.setattr(urllib.request, "urlopen", failing_urlopen)
|
||||||
|
hf_calls = []
|
||||||
|
|
||||||
|
def fake_snapshot_download(**kwargs):
|
||||||
|
hf_calls.append(kwargs)
|
||||||
|
local_dir = Path(kwargs["local_dir"])
|
||||||
|
local_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
(local_dir / "model-00002-of-00004.safetensors").write_text("hf")
|
||||||
|
return str(local_dir)
|
||||||
|
|
||||||
|
monkeypatch.setitem(
|
||||||
|
sys.modules,
|
||||||
|
"huggingface_hub",
|
||||||
|
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
|
||||||
|
)
|
||||||
|
|
||||||
|
shard_dir = download_shard(
|
||||||
|
"tiny-llama",
|
||||||
|
2,
|
||||||
|
3,
|
||||||
|
cache_dir=tmp_path / "cache",
|
||||||
|
hf_repo="org/tiny-llama-shards",
|
||||||
|
model_sources=[{
|
||||||
|
"type": "tracker",
|
||||||
|
"url": "http://tracker/v1/model-files/download?model=tiny-llama",
|
||||||
|
"files": ["config.json", "model-00002-of-00004.safetensors"],
|
||||||
|
}],
|
||||||
|
progress=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert (shard_dir / "model-00002-of-00004.safetensors").read_text() == "hf"
|
||||||
assert hf_calls[0]["allow_patterns"] == ["config.json", "model-00002-of-00004.safetensors"]
|
assert hf_calls[0]["allow_patterns"] == ["config.json", "model-00002-of-00004.safetensors"]
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user