diff --git a/src/gpu_rent/capture.py b/src/gpu_rent/capture.py index 12eb52e..2ab0341 100644 --- a/src/gpu_rent/capture.py +++ b/src/gpu_rent/capture.py @@ -408,15 +408,13 @@ def merge_extensions_yaml( def merge_wanted_models_from_vm(cfg: Config, host: str, log: Log) -> int: """Pull VM Assistent wanted queue into local models.yaml. Returns new url count.""" from gpu_rent.manifests import MODEL_DIRS - from gpu_rent.provision import DATA - remote = f"{DATA}/.gpu-rent-wanted-models.yaml" + data_root = "/mnt/swarm_data" + remote = f"{data_root}/.gpu-rent-wanted-models.yaml" raw = run_ssh( cfg, host, - f"test -f {remote!r} && cat {remote!r} || true".replace("'", '"') - if False - else f'test -f {remote} && cat {remote} || true', + f"test -f {remote} && cat {remote} || true", check=False, timeout=20, ).strip() @@ -611,9 +609,6 @@ def capture_models( # Also merge Assistent wanted-queue (models not yet on disk / seed list). try: - from gpu_rent.provision import merge_wanted_models_from_vm - - # merge writes directly; still collect for report when dry_run=False if not dry_run: n = merge_wanted_models_from_vm(cfg, host, log) if n: diff --git a/src/gpu_rent/provision.py b/src/gpu_rent/provision.py index c496c2e..77a9680 100644 --- a/src/gpu_rent/provision.py +++ b/src/gpu_rent/provision.py @@ -750,59 +750,9 @@ def seed_assistent_personas(cfg: Config, host: str, log: Log) -> None: def merge_wanted_models_from_vm(cfg: Config, host: str, log: Log) -> int: """Pull VM wanted queue into local models.yaml. Returns count of new urls merged.""" - from gpu_rent.capture import ModelCaptureItem, merge_models_yaml + from gpu_rent.capture import merge_wanted_models_from_vm as _merge - remote = f"{DATA}/.gpu-rent-wanted-models.yaml" - raw = run_ssh( - cfg, - host, - f"test -f {shlex.quote(remote)} && cat {shlex.quote(remote)} || true", - check=False, - timeout=20, - ).strip() - if not raw: - return 0 - try: - import yaml - - data = yaml.safe_load(raw) or {} - except Exception as exc: - log(f"wanted-models: bad yaml — {exc}") - return 0 - if not isinstance(data, dict): - return 0 - items: list[ModelCaptureItem] = [] - for kind, rows in data.items(): - if kind not in MODEL_DIRS or not isinstance(rows, list): - continue - for row in rows: - if not isinstance(row, dict): - continue - url = str(row.get("url") or "").strip() - if not url: - continue - vid = row.get("version_id") - try: - vid_i = int(vid) if vid not in (None, "", 0, "0") else extract_version_id(url) - except (TypeError, ValueError): - vid_i = extract_version_id(url) - items.append( - ModelCaptureItem( - kind=kind, - url=url, - title=str(row.get("title") or Path(url).name), - version_id=vid_i, - source="wanted", - ) - ) - if not items: - return 0 - added, skipped = merge_models_yaml(Path(cfg.models_manifest), items, dry_run=False) - if added: - log(f"wanted→models.yaml: +{len(added)} (skip {len(skipped)})") - elif skipped: - log(f"wanted→models.yaml: уже есть ({len(skipped)})") - return len(added) + return _merge(cfg, host, log) def _place_wanted_cards(cfg: Config, host: str, jobs: list[dict], log: Log) -> None: