Split asset sync engine into collaborators
This commit is contained in:
@@ -79,6 +79,11 @@
|
||||
!/ssh_executor.py
|
||||
!/ssh_runtime.py
|
||||
!/sync_engine.py
|
||||
!/sync_backends.py
|
||||
!/sync_custom_nodes.py
|
||||
!/sync_hashing.py
|
||||
!/sync_protocols.py
|
||||
!/sync_r2_transfer.py
|
||||
!/vast_api.py
|
||||
!/vast_config_node.py
|
||||
!/vast_executor.py
|
||||
|
||||
@@ -520,6 +520,8 @@ The prompt interception layer is organized as a compatibility aggregator over fo
|
||||
|
||||
Resident LLM inference uses the same ownership boundaries: `llm_types.py` defines backend-neutral values, `llm_inputs.py` normalizes multimodal and file inputs, `vllm_instrumentation.py` owns process-wide vLLM execution policy and Triton telemetry, and `llm_backend_transformers.py`, `llm_backend_llamacpp.py`, and `llm_backend_vllm.py` isolate backend implementations. `modal_llm_runtime.py` remains the model-residency coordinator, while `local_llm_runtime.py` consumes shared types and input preparation without importing private helpers from the Modal backend.
|
||||
|
||||
Asset synchronization is composed from focused collaborators: `sync_protocols.py` holds provider-neutral contracts and values, `sync_backends.py` owns local and Modal persistence adapters, `sync_hashing.py` owns the persistent content-hash cache, `sync_custom_nodes.py` builds and transfers custom-node archives and manifests, and `sync_r2_transfer.py` owns worker preflight plus idle-gated R2 read-through and write-back. `sync_engine.py` coordinates those collaborators and retains compatibility entrypoints.
|
||||
|
||||
### Vast.ai API Simulator
|
||||
|
||||
Vast development and CI need no live credentials or billable instances. Start the stateful local simulator:
|
||||
|
||||
@@ -676,3 +676,6 @@
|
||||
- Backend-neutral LLM request values and input preparation should not live in a provider-named runtime. Moving them to `llm_types.py` and `llm_inputs.py` lets Apple-local inference depend only on the shared contract while the Modal runtime remains a residency coordinator.
|
||||
- vLLM's execution-mode controller and Triton listener installation are process-wide mutable instrumentation, not model-manager state. They must move together—including listener paths, patched-entrypoint references, locks, and singleton controller—and tests must patch that owner rather than a compatibility import in `modal_llm_runtime.py`.
|
||||
- Resident inference backends have sharply different lifecycle dependencies: Transformers owns synchronous tensor preparation, llama.cpp owns its private subprocess and HTTP server, and vLLM owns an asynchronous event loop plus EngineCore policy. Separate modules keep those imports and cleanup paths isolated while a backend factory preserves the existing coordinator contract.
|
||||
- Sync protocols and result values form a lower dependency layer than storage implementations. Keeping them in `sync_protocols.py` lets local, Modal, SSH, Vast, and R2-capable backends share runtime-checkable capabilities without importing the engine coordinator.
|
||||
- Persistent hash metadata, custom-node scan state, and the R2 idle queue are three different mutable lifecycles. Compose dedicated owners for each and keep only typed delegation methods on `ModalAssetSyncEngine`; otherwise a nominal file split leaves locks and caches owned by the compatibility module.
|
||||
- Tests that replace Modal SDK guards, clocks, archive builders, or uncached custom-node passes must patch the composed owner that resolves those names. Patching a class or helper re-export in `sync_engine.py` does not alter `sync_backends.py` or a live `CustomNodesSynchronizer` instance.
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
"""Local and Modal storage/index backends for asset synchronization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
import modal # type: ignore
|
||||
except ModuleNotFoundError: # pragma: no cover - exercised without Modal SDK.
|
||||
modal = None
|
||||
|
||||
|
||||
def _modal_volume_worker_count() -> int:
|
||||
"""Return the worker count used for local Modal volume SDK calls."""
|
||||
return 4
|
||||
|
||||
|
||||
_MODAL_VOLUME_EXECUTOR = ThreadPoolExecutor(max_workers=_modal_volume_worker_count())
|
||||
|
||||
class _ModalSdkCaller:
|
||||
"""Shared retry and backoff helper for Modal SDK calls."""
|
||||
|
||||
def __init__(self, *, target_kind: str) -> None:
|
||||
"""Initialize shared retry bookkeeping for one Modal SDK target."""
|
||||
self._target_kind = target_kind
|
||||
self._rate_limit_lock = threading.Lock()
|
||||
self._rate_limit_until_monotonic = 0.0
|
||||
self._rate_limit_backoff_seconds = 0.0
|
||||
|
||||
def _resource_exhausted_error_types(self) -> tuple[type[BaseException], ...]:
|
||||
"""Return the Modal SDK exception types that indicate transient rate limiting."""
|
||||
if modal is None:
|
||||
return ()
|
||||
exception_namespace = getattr(modal, "exception", None)
|
||||
if exception_namespace is None:
|
||||
return ()
|
||||
error_type = getattr(exception_namespace, "ResourceExhaustedError", None)
|
||||
if isinstance(error_type, type) and issubclass(error_type, BaseException):
|
||||
return (error_type,)
|
||||
return ()
|
||||
|
||||
def _wait_for_shared_rate_limit_backoff(self) -> None:
|
||||
"""Pause until the shared Modal backoff window expires."""
|
||||
while True:
|
||||
with self._rate_limit_lock:
|
||||
remaining_seconds = self._rate_limit_until_monotonic - time.monotonic()
|
||||
if remaining_seconds <= 0.0:
|
||||
return
|
||||
time.sleep(remaining_seconds)
|
||||
|
||||
def _record_shared_rate_limit_backoff(self) -> float:
|
||||
"""Increase and publish the shared Modal backoff window."""
|
||||
with self._rate_limit_lock:
|
||||
next_backoff_seconds = (
|
||||
0.25
|
||||
if self._rate_limit_backoff_seconds <= 0.0
|
||||
else min(self._rate_limit_backoff_seconds * 2.0, 8.0)
|
||||
)
|
||||
self._rate_limit_backoff_seconds = next_backoff_seconds
|
||||
self._rate_limit_until_monotonic = max(
|
||||
self._rate_limit_until_monotonic,
|
||||
time.monotonic() + next_backoff_seconds,
|
||||
)
|
||||
return next_backoff_seconds
|
||||
|
||||
def _clear_shared_rate_limit_backoff_if_expired(self) -> None:
|
||||
"""Reset the shared backoff after the cooldown window has fully elapsed."""
|
||||
with self._rate_limit_lock:
|
||||
if time.monotonic() >= self._rate_limit_until_monotonic:
|
||||
self._rate_limit_backoff_seconds = 0.0
|
||||
self._rate_limit_until_monotonic = 0.0
|
||||
|
||||
def _run_sdk_call(self, callback: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
"""Run one Modal SDK call with shared retry and backoff semantics."""
|
||||
retryable_errors = self._resource_exhausted_error_types()
|
||||
max_attempts = 5
|
||||
for attempt_index in range(max_attempts):
|
||||
self._wait_for_shared_rate_limit_backoff()
|
||||
future = _MODAL_VOLUME_EXECUTOR.submit(callback, *args, **kwargs)
|
||||
try:
|
||||
result = future.result()
|
||||
self._clear_shared_rate_limit_backoff_if_expired()
|
||||
return result
|
||||
except retryable_errors:
|
||||
if attempt_index >= max_attempts - 1:
|
||||
raise
|
||||
backoff_seconds = self._record_shared_rate_limit_backoff()
|
||||
logger.warning(
|
||||
"Modal %s call %s hit rate limiting on attempt %d/%d; applying shared retry backoff of %.2fs.",
|
||||
self._target_kind,
|
||||
getattr(callback, "__name__", repr(callback)),
|
||||
attempt_index + 1,
|
||||
max_attempts,
|
||||
backoff_seconds,
|
||||
)
|
||||
raise RuntimeError("Modal SDK call retry loop exited unexpectedly.")
|
||||
|
||||
|
||||
class LocalMirrorVolume:
|
||||
"""Simple filesystem-backed volume used for tests and dry runs."""
|
||||
|
||||
def __init__(self, root: Path) -> None:
|
||||
"""Initialize the local mirror volume root."""
|
||||
self.root = root
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def exists(self, remote_path: str) -> bool:
|
||||
"""Return whether a file already exists in the local mirror."""
|
||||
return self._resolve(remote_path).exists()
|
||||
|
||||
def put_file(self, local_path: Path, remote_path: str) -> None:
|
||||
"""Copy a local file into the mirror volume."""
|
||||
target = self._resolve(remote_path)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_bytes(local_path.read_bytes())
|
||||
|
||||
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
||||
"""Write bytes into the mirror volume."""
|
||||
target = self._resolve(remote_path)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_bytes(payload)
|
||||
|
||||
def _resolve(self, remote_path: str) -> Path:
|
||||
"""Resolve a remote storage path relative to the mirror root."""
|
||||
return self.root / remote_path.lstrip("/")
|
||||
|
||||
|
||||
class LocalFileSyncIndex:
|
||||
"""JSON-backed sync index used for local mirrors and tests."""
|
||||
|
||||
def __init__(self, root: Path) -> None:
|
||||
"""Initialize the on-disk metadata store."""
|
||||
self._index_path = root / "metadata" / "sync_index.json"
|
||||
self._lock = threading.Lock()
|
||||
self._records = self._load_records()
|
||||
|
||||
def get(self, key: str) -> dict[str, Any] | None:
|
||||
"""Return one stored sync record when it exists."""
|
||||
with self._lock:
|
||||
payload = self._records.get(key)
|
||||
return dict(payload) if isinstance(payload, dict) else None
|
||||
|
||||
def put(self, key: str, value: dict[str, Any]) -> None:
|
||||
"""Persist one sync record to the local metadata file."""
|
||||
with self._lock:
|
||||
self._records[key] = dict(value)
|
||||
self._save_records()
|
||||
|
||||
def _load_records(self) -> dict[str, dict[str, Any]]:
|
||||
"""Load the persisted sync index when available."""
|
||||
if not self._index_path.exists():
|
||||
return {}
|
||||
try:
|
||||
payload = json.loads(self._index_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
logger.warning("Sync index at %s is unreadable; rebuilding it from scratch.", self._index_path)
|
||||
return {}
|
||||
if not isinstance(payload, dict):
|
||||
return {}
|
||||
return {
|
||||
str(key): value
|
||||
for key, value in payload.items()
|
||||
if isinstance(value, dict)
|
||||
}
|
||||
|
||||
def _save_records(self) -> None:
|
||||
"""Write the current sync index to disk."""
|
||||
self._index_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._index_path.write_text(json.dumps(self._records, sort_keys=True), encoding="utf-8")
|
||||
|
||||
|
||||
class ModalDictSyncIndex(_ModalSdkCaller):
|
||||
"""Modal Dict-backed sync index for remote content-addressed uploads."""
|
||||
|
||||
def __init__(self, dict_name: str) -> None:
|
||||
"""Resolve a named Modal Dict lazily from the local SDK client."""
|
||||
if modal is None:
|
||||
raise RuntimeError("Modal SDK is required for ModalDictSyncIndex.")
|
||||
super().__init__(target_kind="dict")
|
||||
self._dict = modal.Dict.from_name(dict_name, create_if_missing=True)
|
||||
self._missing = object()
|
||||
self._cache: dict[str, object] = {}
|
||||
self._cache_lock = threading.Lock()
|
||||
|
||||
def get(self, key: str) -> dict[str, Any] | None:
|
||||
"""Return one stored sync record when it exists."""
|
||||
with self._cache_lock:
|
||||
cached_value = self._cache.get(key, self._missing)
|
||||
if cached_value is not self._missing:
|
||||
return dict(cached_value) if isinstance(cached_value, dict) else None
|
||||
payload = self._run_sdk_call(self._dict.get, key)
|
||||
normalized_payload = dict(payload) if isinstance(payload, dict) else None
|
||||
with self._cache_lock:
|
||||
self._cache[key] = dict(normalized_payload) if normalized_payload is not None else None
|
||||
return dict(normalized_payload) if normalized_payload is not None else None
|
||||
|
||||
def put(self, key: str, value: dict[str, Any]) -> None:
|
||||
"""Persist one sync record to the shared Modal Dict."""
|
||||
normalized_value = dict(value)
|
||||
|
||||
def write_record() -> None:
|
||||
self._dict[key] = normalized_value
|
||||
|
||||
self._run_sdk_call(write_record)
|
||||
with self._cache_lock:
|
||||
self._cache[key] = dict(normalized_value)
|
||||
|
||||
|
||||
class ModalVolumeBackend(_ModalSdkCaller):
|
||||
"""Modal Volume-backed storage for real remote execution."""
|
||||
|
||||
def __init__(self, volume_name: str) -> None:
|
||||
"""Resolve a named Modal volume lazily from the local SDK client."""
|
||||
if modal is None:
|
||||
raise RuntimeError("Modal SDK is required for ModalVolumeBackend.")
|
||||
super().__init__(target_kind="volume")
|
||||
self._volume = modal.Volume.from_name(volume_name, create_if_missing=True)
|
||||
self._exists_cache: dict[str, bool] = {}
|
||||
self._exists_cache_lock = threading.Lock()
|
||||
|
||||
def exists(self, remote_path: str) -> bool:
|
||||
"""Return whether a file already exists in the Modal volume."""
|
||||
with self._exists_cache_lock:
|
||||
cached_result = self._exists_cache.get(remote_path)
|
||||
if cached_result is not None:
|
||||
return cached_result
|
||||
try:
|
||||
exists = len(self._run_sdk_call(self._volume.listdir, remote_path, recursive=False)) > 0
|
||||
except modal.exception.NotFoundError:
|
||||
exists = False
|
||||
with self._exists_cache_lock:
|
||||
self._exists_cache[remote_path] = exists
|
||||
return exists
|
||||
|
||||
def put_file(self, local_path: Path, remote_path: str) -> None:
|
||||
"""Upload a local file into the Modal volume."""
|
||||
def upload() -> None:
|
||||
with self._volume.batch_upload() as batch:
|
||||
batch.put_file(local_path, remote_path)
|
||||
|
||||
try:
|
||||
self._run_sdk_call(upload)
|
||||
except FileExistsError:
|
||||
logger.info(
|
||||
"Treating Modal volume upload for %s as successful because the content-addressed path already exists.",
|
||||
remote_path,
|
||||
)
|
||||
with self._exists_cache_lock:
|
||||
self._exists_cache[remote_path] = True
|
||||
|
||||
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
||||
"""Upload bytes into the Modal volume."""
|
||||
def upload() -> None:
|
||||
with self._volume.batch_upload() as batch:
|
||||
batch.put_file(io.BytesIO(payload), remote_path)
|
||||
|
||||
try:
|
||||
self._run_sdk_call(upload)
|
||||
except FileExistsError:
|
||||
logger.info(
|
||||
"Treating Modal volume byte upload for %s as successful because the content-addressed path already exists.",
|
||||
remote_path,
|
||||
)
|
||||
with self._exists_cache_lock:
|
||||
self._exists_cache[remote_path] = True
|
||||
|
||||
@@ -0,0 +1,638 @@
|
||||
"""Custom-node archive, asset, manifest, and backfill synchronization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Protocol
|
||||
import zipfile
|
||||
|
||||
if __package__:
|
||||
from .settings import ModalSyncSettings
|
||||
from .sync_backends import _modal_volume_worker_count
|
||||
from .sync_hashing import _SKIP_DIRS, _SKIP_FILE_SUFFIXES
|
||||
from .sync_protocols import (
|
||||
CancellationCheck,
|
||||
SyncedAsset,
|
||||
SyncStatusCallback,
|
||||
_ContentAddressedSyncResult,
|
||||
_CustomNodeAssetSpec,
|
||||
_CustomNodeAssetSyncResult,
|
||||
_CustomNodesArchiveSpec,
|
||||
_CustomNodesArchiveSyncResult,
|
||||
)
|
||||
from .sync_r2_transfer import R2WriteBackCancelled, _R2_WRITE_BACK_COORDINATOR
|
||||
else: # pragma: no cover - flat import inside the Modal container.
|
||||
from settings import ModalSyncSettings
|
||||
from sync_backends import _modal_volume_worker_count
|
||||
from sync_hashing import _SKIP_DIRS, _SKIP_FILE_SUFFIXES
|
||||
from sync_protocols import (
|
||||
CancellationCheck,
|
||||
SyncedAsset,
|
||||
SyncStatusCallback,
|
||||
_ContentAddressedSyncResult,
|
||||
_CustomNodeAssetSpec,
|
||||
_CustomNodeAssetSyncResult,
|
||||
_CustomNodesArchiveSpec,
|
||||
_CustomNodesArchiveSyncResult,
|
||||
)
|
||||
from sync_r2_transfer import R2WriteBackCancelled, _R2_WRITE_BACK_COORDINATOR
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_CUSTOM_NODES_MANIFEST_VERSION = 2
|
||||
MODEL_FILE_EXTENSIONS = frozenset({
|
||||
".bin", ".ckpt", ".engine", ".gguf", ".onnx", ".pt", ".pth",
|
||||
".safetensors", ".vae",
|
||||
})
|
||||
_CUSTOM_NODE_ASSET_SUFFIXES = MODEL_FILE_EXTENSIONS
|
||||
|
||||
|
||||
def _emit_sync_status(
|
||||
status_callback: SyncStatusCallback | None,
|
||||
message: str,
|
||||
current: int | None = None,
|
||||
total: int | None = None,
|
||||
) -> None:
|
||||
"""Emit one human-readable sync status update when available."""
|
||||
if status_callback is not None:
|
||||
status_callback(message, current, total)
|
||||
|
||||
|
||||
def _custom_nodes_sync_worker_count() -> int:
|
||||
"""Return the worker count used for parallel custom-node synchronization."""
|
||||
return max(4, min(16, os.cpu_count() or 1))
|
||||
|
||||
|
||||
class CustomNodesSyncHost(Protocol):
|
||||
"""Define the narrow engine services consumed by custom-node synchronization."""
|
||||
|
||||
settings: ModalSyncSettings
|
||||
|
||||
def _destination_label(self) -> str:
|
||||
"""Return a user-facing synchronization destination."""
|
||||
|
||||
def _hash_directory(self, path: Path) -> str:
|
||||
"""Return the stable digest for one directory."""
|
||||
|
||||
def _hash_file(self, path: Path) -> str:
|
||||
"""Return the content digest for one file."""
|
||||
|
||||
def _hash_file_group(self, root: Path, files: list[Path]) -> str:
|
||||
"""Return the stable digest for selected files."""
|
||||
|
||||
def _iter_files(self, path: Path) -> list[Path]:
|
||||
"""Return syncable files below one directory."""
|
||||
|
||||
def _lookup_sync_record(self, sync_key: str) -> dict[str, Any] | None:
|
||||
"""Return an indexed synchronization record."""
|
||||
|
||||
def _r2_writeback_enabled(self) -> bool:
|
||||
"""Return whether R2 write-back is configured."""
|
||||
|
||||
def _r2_writeback_job_prefix(self) -> tuple[str, str, str]:
|
||||
"""Return the stable write-back namespace."""
|
||||
|
||||
def _sync_content_addressed_file(
|
||||
self,
|
||||
**kwargs: Any,
|
||||
) -> _ContentAddressedSyncResult:
|
||||
"""Synchronize one content-addressed file."""
|
||||
|
||||
def _sync_index_scope_prefix(self) -> str:
|
||||
"""Return the active persistent-index scope."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class CustomNodesSynchronizer:
|
||||
"""Own custom-node synchronization state and archive orchestration."""
|
||||
|
||||
_host: CustomNodesSyncHost
|
||||
_custom_nodes_sync_lock: threading.Lock = field(init=False, default_factory=threading.Lock)
|
||||
_custom_nodes_sync_checked: bool = field(init=False, default=False)
|
||||
_custom_nodes_bundle_cache: SyncedAsset | None = field(init=False, default=None)
|
||||
|
||||
def sync_custom_nodes_directory(
|
||||
self,
|
||||
*,
|
||||
status_callback: SyncStatusCallback | None = None,
|
||||
) -> SyncedAsset | None:
|
||||
"""Mirror custom_nodes once per ComfyUI process and reuse that result afterward."""
|
||||
with self._custom_nodes_sync_lock:
|
||||
if self._custom_nodes_sync_checked:
|
||||
logger.info(
|
||||
"Skipping custom_nodes rescan because this ComfyUI process already resolved it once."
|
||||
)
|
||||
return self._clone_cached_custom_nodes_bundle()
|
||||
|
||||
bundle = self._sync_custom_nodes_directory_uncached(status_callback=status_callback)
|
||||
self._custom_nodes_bundle_cache = bundle
|
||||
self._custom_nodes_sync_checked = True
|
||||
return bundle
|
||||
|
||||
def _clone_cached_custom_nodes_bundle(self) -> SyncedAsset | None:
|
||||
"""Return the cached custom_nodes sync result as a no-upload reuse decision."""
|
||||
if self._custom_nodes_bundle_cache is None:
|
||||
return None
|
||||
return SyncedAsset(
|
||||
local_path=self._custom_nodes_bundle_cache.local_path,
|
||||
remote_path=self._custom_nodes_bundle_cache.remote_path,
|
||||
sha256=self._custom_nodes_bundle_cache.sha256,
|
||||
uploaded=False,
|
||||
)
|
||||
|
||||
def _sync_custom_nodes_directory_uncached(
|
||||
self,
|
||||
*,
|
||||
status_callback: SyncStatusCallback | None = None,
|
||||
revisit_indexed_for_r2: bool = False,
|
||||
writeback_cancellation_check: CancellationCheck | None = None,
|
||||
) -> SyncedAsset | None:
|
||||
"""Mirror custom_nodes as a manifest plus per-package archives when available."""
|
||||
if writeback_cancellation_check is not None and writeback_cancellation_check():
|
||||
raise R2WriteBackCancelled("Custom-node R2 backfill yielded before scanning.")
|
||||
custom_nodes_dir = self._host.settings.custom_nodes_dir
|
||||
if custom_nodes_dir is None or not custom_nodes_dir.exists():
|
||||
logger.info("No custom_nodes directory detected for mirroring.")
|
||||
return None
|
||||
|
||||
sync_started_at = time.perf_counter()
|
||||
logger.info("Hashing custom_nodes directory at %s", custom_nodes_dir)
|
||||
directory_hash = self._host._hash_directory(custom_nodes_dir)
|
||||
if writeback_cancellation_check is not None and writeback_cancellation_check():
|
||||
raise R2WriteBackCancelled("Custom-node R2 backfill yielded after hashing.")
|
||||
logger.info(
|
||||
"Finished hashing custom_nodes directory in %.3fs with digest %s.",
|
||||
time.perf_counter() - sync_started_at,
|
||||
directory_hash,
|
||||
)
|
||||
manifest_sync_key = self._custom_nodes_manifest_sync_index_key(directory_hash)
|
||||
manifest_record = self._host._lookup_sync_record(manifest_sync_key)
|
||||
if manifest_record is not None and (
|
||||
not self._host._r2_writeback_enabled() or not revisit_indexed_for_r2
|
||||
):
|
||||
remote_path = str(manifest_record["remote_path"])
|
||||
if self._host._r2_writeback_enabled():
|
||||
self._schedule_custom_nodes_r2_backfill(directory_hash)
|
||||
logger.info(
|
||||
"Custom_nodes manifest already mirrored at %s after %.3fs total sync time.",
|
||||
remote_path,
|
||||
time.perf_counter() - sync_started_at,
|
||||
)
|
||||
return SyncedAsset(
|
||||
local_path=custom_nodes_dir,
|
||||
remote_path=remote_path,
|
||||
sha256=directory_hash,
|
||||
uploaded=False,
|
||||
)
|
||||
if manifest_record is not None:
|
||||
logger.info(
|
||||
"Revisiting indexed custom_nodes entries during idle time so R2 "
|
||||
"write-back can backfill missing objects."
|
||||
)
|
||||
remote_path = self._custom_nodes_manifest_remote_path(directory_hash)
|
||||
|
||||
archive_specs, asset_specs = self._custom_nodes_bundle_specs(custom_nodes_dir)
|
||||
if writeback_cancellation_check is not None and writeback_cancellation_check():
|
||||
raise R2WriteBackCancelled("Custom-node R2 backfill yielded after planning.")
|
||||
if not archive_specs:
|
||||
logger.info("Custom_nodes directory %s contained no syncable files.", custom_nodes_dir)
|
||||
return None
|
||||
|
||||
if any(
|
||||
not self._cached_custom_nodes_archive_path(
|
||||
archive_spec.entry_name,
|
||||
archive_spec.sha256,
|
||||
).exists()
|
||||
for archive_spec in archive_specs
|
||||
):
|
||||
_emit_sync_status(
|
||||
status_callback,
|
||||
f"Packaging custom-node code for {self._host._destination_label()}",
|
||||
)
|
||||
_emit_sync_status(
|
||||
status_callback,
|
||||
f"Uploading custom-node code and assets to {self._host._destination_label()}",
|
||||
)
|
||||
|
||||
archive_results = self._sync_custom_nodes_archives_parallel(
|
||||
custom_nodes_dir=custom_nodes_dir,
|
||||
archive_specs=archive_specs,
|
||||
)
|
||||
asset_results = self._sync_custom_node_assets_parallel(asset_specs)
|
||||
if writeback_cancellation_check is not None and writeback_cancellation_check():
|
||||
raise R2WriteBackCancelled("Custom-node R2 backfill yielded after packaging.")
|
||||
uploaded = any(
|
||||
result.uploaded for result in [*archive_results, *asset_results]
|
||||
)
|
||||
assets_by_entry: dict[str, list[dict[str, Any]]] = {}
|
||||
for asset_result in asset_results:
|
||||
assets_by_entry.setdefault(asset_result.entry_name, []).append(
|
||||
{
|
||||
"relative_path": asset_result.relative_path,
|
||||
"sha256": asset_result.sha256,
|
||||
"size_bytes": asset_result.size_bytes,
|
||||
"remote_path": asset_result.remote_path,
|
||||
}
|
||||
)
|
||||
manifest_entries = [
|
||||
{
|
||||
"entry_name": archive_result.entry_name,
|
||||
"display_name": archive_result.display_name,
|
||||
"sha256": archive_result.sha256,
|
||||
"remote_path": archive_result.remote_path,
|
||||
"assets": assets_by_entry.get(archive_result.entry_name, []),
|
||||
}
|
||||
for archive_result in archive_results
|
||||
]
|
||||
|
||||
manifest_path = self._cached_custom_nodes_manifest_path(directory_hash)
|
||||
if manifest_path.exists():
|
||||
logger.info(
|
||||
"Reusing cached custom_nodes manifest %s for digest %s.",
|
||||
manifest_path,
|
||||
directory_hash,
|
||||
)
|
||||
else:
|
||||
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
manifest_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"version": _CUSTOM_NODES_MANIFEST_VERSION,
|
||||
"bundle_sha256": directory_hash,
|
||||
"entries": manifest_entries,
|
||||
},
|
||||
sort_keys=True,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
manifest_sync_result = self._host._sync_content_addressed_file(
|
||||
local_path=manifest_path,
|
||||
remote_path=remote_path,
|
||||
sha256=self._host._hash_file(manifest_path),
|
||||
sync_key=manifest_sync_key,
|
||||
source_description=str(custom_nodes_dir),
|
||||
upload_status_message=(
|
||||
f"Uploading custom-node manifest to {self._host._destination_label()}"
|
||||
),
|
||||
)
|
||||
uploaded = uploaded or manifest_sync_result.uploaded
|
||||
|
||||
logger.info(
|
||||
"Finished custom_nodes sync to %s in %.3fs total.",
|
||||
manifest_sync_result.remote_path,
|
||||
time.perf_counter() - sync_started_at,
|
||||
)
|
||||
return SyncedAsset(
|
||||
local_path=custom_nodes_dir,
|
||||
remote_path=manifest_sync_result.remote_path,
|
||||
sha256=directory_hash,
|
||||
uploaded=uploaded,
|
||||
)
|
||||
|
||||
def _schedule_custom_nodes_r2_backfill(self, directory_hash: str) -> None:
|
||||
"""Defer indexed custom-node traversal until foreground workflows are idle."""
|
||||
if not self._host._r2_writeback_enabled():
|
||||
return
|
||||
job_key = (
|
||||
*self._host._r2_writeback_job_prefix(),
|
||||
"custom-nodes-backfill",
|
||||
directory_hash,
|
||||
)
|
||||
_R2_WRITE_BACK_COORDINATOR.submit(
|
||||
job_key,
|
||||
lambda cancellation_check: self._sync_custom_nodes_directory_uncached(
|
||||
revisit_indexed_for_r2=True,
|
||||
writeback_cancellation_check=cancellation_check,
|
||||
),
|
||||
)
|
||||
|
||||
def _sync_custom_nodes_archives_parallel(
|
||||
self,
|
||||
*,
|
||||
custom_nodes_dir: Path,
|
||||
archive_specs: list[_CustomNodesArchiveSpec],
|
||||
) -> list[_CustomNodesArchiveSyncResult]:
|
||||
"""Build and upload per-package custom_nodes archives in parallel."""
|
||||
max_workers = min(len(archive_specs), _custom_nodes_sync_worker_count())
|
||||
if max_workers <= 1:
|
||||
return [
|
||||
self._sync_custom_nodes_archive_spec(custom_nodes_dir, archive_spec)
|
||||
for archive_spec in archive_specs
|
||||
]
|
||||
|
||||
results_by_entry_name: dict[str, _CustomNodesArchiveSyncResult] = {}
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures_by_entry_name: dict[str, Future[_CustomNodesArchiveSyncResult]] = {
|
||||
archive_spec.entry_name: executor.submit(
|
||||
self._sync_custom_nodes_archive_spec,
|
||||
custom_nodes_dir,
|
||||
archive_spec,
|
||||
)
|
||||
for archive_spec in archive_specs
|
||||
}
|
||||
for entry_name, future in futures_by_entry_name.items():
|
||||
results_by_entry_name[entry_name] = future.result()
|
||||
return [
|
||||
results_by_entry_name[archive_spec.entry_name]
|
||||
for archive_spec in archive_specs
|
||||
]
|
||||
|
||||
def _sync_custom_nodes_archive_spec(
|
||||
self,
|
||||
custom_nodes_dir: Path,
|
||||
archive_spec: _CustomNodesArchiveSpec,
|
||||
) -> _CustomNodesArchiveSyncResult:
|
||||
"""Build and upload one per-package custom_nodes archive."""
|
||||
archive_path = self._cached_custom_nodes_archive_path(
|
||||
archive_spec.entry_name,
|
||||
archive_spec.sha256,
|
||||
)
|
||||
archive_remote_path = self._custom_nodes_archive_remote_path(
|
||||
archive_spec.entry_name,
|
||||
archive_spec.sha256,
|
||||
)
|
||||
if archive_path.exists():
|
||||
logger.info(
|
||||
"Reusing cached custom_nodes archive %s for entry=%s digest=%s.",
|
||||
archive_path,
|
||||
archive_spec.display_name,
|
||||
archive_spec.sha256,
|
||||
)
|
||||
else:
|
||||
archive_started_at = time.perf_counter()
|
||||
logger.info(
|
||||
"Creating custom_nodes archive for entry=%s from %d files.",
|
||||
archive_spec.display_name,
|
||||
len(archive_spec.files),
|
||||
)
|
||||
self._create_archive_from_files(
|
||||
custom_nodes_dir,
|
||||
list(archive_spec.files),
|
||||
archive_path,
|
||||
)
|
||||
logger.info(
|
||||
"Created custom_nodes archive %s for entry=%s in %.3fs.",
|
||||
archive_path,
|
||||
archive_spec.display_name,
|
||||
time.perf_counter() - archive_started_at,
|
||||
)
|
||||
|
||||
archive_sha256 = self._host._hash_file(archive_path)
|
||||
entry_uploaded = self._host._sync_content_addressed_file(
|
||||
local_path=archive_path,
|
||||
remote_path=archive_remote_path,
|
||||
sha256=archive_sha256,
|
||||
sync_key=self._custom_nodes_entry_sync_index_key(
|
||||
archive_spec.entry_name,
|
||||
archive_spec.sha256,
|
||||
),
|
||||
source_description=archive_spec.source_description,
|
||||
)
|
||||
return _CustomNodesArchiveSyncResult(
|
||||
entry_name=archive_spec.entry_name,
|
||||
display_name=archive_spec.display_name,
|
||||
sha256=archive_sha256,
|
||||
remote_path=entry_uploaded.remote_path,
|
||||
uploaded=entry_uploaded.uploaded,
|
||||
)
|
||||
|
||||
def _sync_custom_node_assets_parallel(
|
||||
self,
|
||||
asset_specs: list[_CustomNodeAssetSpec],
|
||||
) -> list[_CustomNodeAssetSyncResult]:
|
||||
"""Upload package-owned model assets without embedding them in code ZIPs."""
|
||||
if not asset_specs:
|
||||
return []
|
||||
max_workers = min(len(asset_specs), _modal_volume_worker_count())
|
||||
if max_workers <= 1:
|
||||
return [self._sync_custom_node_asset_spec(spec) for spec in asset_specs]
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(self._sync_custom_node_asset_spec, spec)
|
||||
for spec in asset_specs
|
||||
]
|
||||
return [future.result() for future in futures]
|
||||
|
||||
def _sync_custom_node_asset_spec(
|
||||
self,
|
||||
asset_spec: _CustomNodeAssetSpec,
|
||||
) -> _CustomNodeAssetSyncResult:
|
||||
"""Upload one package-owned model asset to its content-addressed path."""
|
||||
remote_path = self._custom_nodes_asset_remote_path(asset_spec)
|
||||
sync_result = self._host._sync_content_addressed_file(
|
||||
local_path=asset_spec.local_path,
|
||||
remote_path=remote_path,
|
||||
sha256=asset_spec.sha256,
|
||||
sync_key=self._custom_nodes_asset_sync_index_key(asset_spec.sha256),
|
||||
source_description=str(asset_spec.local_path),
|
||||
)
|
||||
return _CustomNodeAssetSyncResult(
|
||||
entry_name=asset_spec.entry_name,
|
||||
relative_path=asset_spec.relative_path,
|
||||
sha256=asset_spec.sha256,
|
||||
size_bytes=asset_spec.size_bytes,
|
||||
remote_path=sync_result.remote_path,
|
||||
uploaded=sync_result.uploaded,
|
||||
)
|
||||
|
||||
def _custom_nodes_manifest_sync_index_key(self, directory_hash: str) -> str:
|
||||
"""Return the sync-index key for one whole-tree custom_nodes manifest digest."""
|
||||
return (
|
||||
f"{self._host._sync_index_scope_prefix()}:custom_nodes_manifest:"
|
||||
f"v{_CUSTOM_NODES_MANIFEST_VERSION}:{directory_hash}"
|
||||
)
|
||||
|
||||
def _custom_nodes_entry_sync_index_key(self, entry_name: str, entry_hash: str) -> str:
|
||||
"""Return the sync-index key for one top-level custom_nodes entry archive digest."""
|
||||
return (
|
||||
f"{self._host._sync_index_scope_prefix()}:custom_nodes_entry:"
|
||||
f"{self._custom_nodes_entry_slug(entry_name)}:{entry_hash}"
|
||||
)
|
||||
|
||||
def _custom_nodes_asset_sync_index_key(self, asset_hash: str) -> str:
|
||||
"""Return the sync-index key for one package-owned model asset digest."""
|
||||
return f"{self._host._sync_index_scope_prefix()}:custom_nodes_asset:{asset_hash}"
|
||||
|
||||
def _create_archive(self, path: Path, archive_path: Path) -> Path:
|
||||
"""Create a zip archive for the given directory tree."""
|
||||
files = sorted(self._host._iter_files(path), key=lambda item: item.relative_to(path).as_posix())
|
||||
return self._create_archive_from_files(path, files, archive_path)
|
||||
|
||||
def _create_archive_from_files(
|
||||
self,
|
||||
root_path: Path,
|
||||
files: list[Path],
|
||||
archive_path: Path,
|
||||
) -> Path:
|
||||
"""Create a zip archive from a selected file list rooted under one directory."""
|
||||
archive_started_at = time.perf_counter()
|
||||
archive_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
logger.info("Archiving %d files from %s into %s", len(files), root_path, archive_path)
|
||||
with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
|
||||
for child in sorted(files, key=lambda item: item.relative_to(root_path).as_posix()):
|
||||
archive.write(child, arcname=child.relative_to(root_path))
|
||||
|
||||
logger.info(
|
||||
"Finished archive build for %s in %.3fs.",
|
||||
root_path,
|
||||
time.perf_counter() - archive_started_at,
|
||||
)
|
||||
return archive_path
|
||||
|
||||
def _cached_custom_nodes_archive_path(self, entry_name: str, entry_hash: str) -> Path:
|
||||
"""Return the deterministic local path for one digest-keyed custom_nodes slice archive."""
|
||||
return (
|
||||
self._host.settings.local_storage_root
|
||||
/ "custom_nodes_archives"
|
||||
/ self._custom_nodes_entry_slug(entry_name)
|
||||
/ f"{entry_hash}_{self._host.settings.custom_nodes_archive_name}"
|
||||
)
|
||||
|
||||
def _cached_custom_nodes_manifest_path(self, directory_hash: str) -> Path:
|
||||
"""Return the deterministic local path for a whole-tree custom_nodes manifest."""
|
||||
return (
|
||||
self._host.settings.local_storage_root
|
||||
/ "custom_nodes_manifests"
|
||||
/ f"{directory_hash}_custom_nodes_bundle_manifest_v{_CUSTOM_NODES_MANIFEST_VERSION}.json"
|
||||
)
|
||||
|
||||
def _custom_nodes_manifest_remote_path(self, directory_hash: str) -> str:
|
||||
"""Return the remote storage path for a whole-tree custom_nodes manifest."""
|
||||
return (
|
||||
f"/custom_nodes/manifests/{directory_hash}_custom_nodes_bundle_"
|
||||
f"manifest_v{_CUSTOM_NODES_MANIFEST_VERSION}.json"
|
||||
)
|
||||
|
||||
def _custom_nodes_archive_remote_path(self, entry_name: str, entry_hash: str) -> str:
|
||||
"""Return the remote storage path for one content-addressed custom_nodes slice archive."""
|
||||
return (
|
||||
f"/custom_nodes/entries/{self._custom_nodes_entry_slug(entry_name)}/"
|
||||
f"{entry_hash}_{self._host.settings.custom_nodes_archive_name}"
|
||||
)
|
||||
|
||||
def _custom_nodes_asset_remote_path(self, asset_spec: _CustomNodeAssetSpec) -> str:
|
||||
"""Return the content-addressed remote path for one custom-node model asset."""
|
||||
return (
|
||||
f"/custom_nodes/assets/{self._custom_nodes_entry_slug(asset_spec.entry_name)}/"
|
||||
f"{asset_spec.sha256}_{asset_spec.local_path.name}"
|
||||
)
|
||||
|
||||
def _custom_nodes_entry_slug(self, entry_name: str) -> str:
|
||||
"""Return a filesystem-safe slug for one top-level custom_nodes entry name."""
|
||||
normalized_name = entry_name.strip() or "root_files"
|
||||
return "".join(
|
||||
character if character.isalnum() or character in {"-", "_", "."} else "_"
|
||||
for character in normalized_name
|
||||
)
|
||||
|
||||
def _custom_nodes_bundle_specs(
|
||||
self,
|
||||
custom_nodes_dir: Path,
|
||||
) -> tuple[list[_CustomNodesArchiveSpec], list[_CustomNodeAssetSpec]]:
|
||||
"""Split custom-node source archives from package-owned model assets."""
|
||||
resolved_root = custom_nodes_dir.resolve()
|
||||
root_files: list[Path] = []
|
||||
archive_specs: list[_CustomNodesArchiveSpec] = []
|
||||
asset_specs: list[_CustomNodeAssetSpec] = []
|
||||
|
||||
for child in sorted(resolved_root.iterdir(), key=lambda item: item.name):
|
||||
if child.name in _SKIP_DIRS:
|
||||
continue
|
||||
if child.is_file():
|
||||
if child.suffix.lower() in _SKIP_FILE_SUFFIXES:
|
||||
continue
|
||||
root_files.append(child)
|
||||
continue
|
||||
if child.is_dir():
|
||||
entry_files = sorted(
|
||||
self._host._iter_files(child),
|
||||
key=lambda item: item.relative_to(resolved_root).as_posix(),
|
||||
)
|
||||
code_files, entry_assets = self._partition_custom_node_files(entry_files)
|
||||
if not code_files:
|
||||
continue
|
||||
archive_specs.append(
|
||||
_CustomNodesArchiveSpec(
|
||||
entry_name=child.name,
|
||||
display_name=child.name,
|
||||
source_description=str(child),
|
||||
files=tuple(code_files),
|
||||
sha256=self._host._hash_file_group(resolved_root, code_files),
|
||||
)
|
||||
)
|
||||
asset_specs.extend(
|
||||
self._custom_node_asset_specs(
|
||||
resolved_root,
|
||||
child.name,
|
||||
entry_assets,
|
||||
)
|
||||
)
|
||||
|
||||
if root_files:
|
||||
code_files, root_assets = self._partition_custom_node_files(root_files)
|
||||
else:
|
||||
code_files, root_assets = [], []
|
||||
if code_files:
|
||||
archive_specs.append(
|
||||
_CustomNodesArchiveSpec(
|
||||
entry_name="root_files",
|
||||
display_name="root files",
|
||||
source_description=str(resolved_root),
|
||||
files=tuple(code_files),
|
||||
sha256=self._host._hash_file_group(resolved_root, code_files),
|
||||
)
|
||||
)
|
||||
asset_specs.extend(
|
||||
self._custom_node_asset_specs(
|
||||
resolved_root,
|
||||
"root_files",
|
||||
root_assets,
|
||||
)
|
||||
)
|
||||
|
||||
return archive_specs, asset_specs
|
||||
|
||||
def _custom_nodes_archive_specs(self, custom_nodes_dir: Path) -> list[_CustomNodesArchiveSpec]:
|
||||
"""Return code-only archive specs for compatibility with existing callers."""
|
||||
archive_specs, _ = self._custom_nodes_bundle_specs(custom_nodes_dir)
|
||||
return archive_specs
|
||||
|
||||
def _partition_custom_node_files(
|
||||
self,
|
||||
files: list[Path],
|
||||
) -> tuple[list[Path], list[Path]]:
|
||||
"""Partition custom-node files into code resources and mounted model assets."""
|
||||
code_files: list[Path] = []
|
||||
asset_files: list[Path] = []
|
||||
for file_path in files:
|
||||
if file_path.suffix.lower() in _CUSTOM_NODE_ASSET_SUFFIXES:
|
||||
asset_files.append(file_path)
|
||||
else:
|
||||
code_files.append(file_path)
|
||||
return code_files, asset_files
|
||||
|
||||
def _custom_node_asset_specs(
|
||||
self,
|
||||
custom_nodes_root: Path,
|
||||
entry_name: str,
|
||||
asset_files: list[Path],
|
||||
) -> list[_CustomNodeAssetSpec]:
|
||||
"""Build deterministic metadata for package-owned model assets."""
|
||||
return [
|
||||
_CustomNodeAssetSpec(
|
||||
entry_name=entry_name,
|
||||
local_path=asset_path,
|
||||
relative_path=asset_path.relative_to(custom_nodes_root).as_posix(),
|
||||
sha256=self._host._hash_file(asset_path),
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
)
|
||||
for asset_path in asset_files
|
||||
]
|
||||
+226
-1594
File diff suppressed because it is too large
Load Diff
+192
@@ -0,0 +1,192 @@
|
||||
"""Persistent hashing and filesystem filtering for asset synchronization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
if __package__:
|
||||
from .settings import ModalSyncSettings
|
||||
from .sync_protocols import CancellationCheck, SyncCancelledError
|
||||
else: # pragma: no cover - flat import inside the Modal container.
|
||||
from settings import ModalSyncSettings
|
||||
from sync_protocols import CancellationCheck, SyncCancelledError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SKIP_DIRS = {
|
||||
".git", ".mypy_cache", ".nox", ".pytest_cache", ".ruff_cache",
|
||||
".tox", ".venv", "__pycache__", "venv",
|
||||
}
|
||||
_SKIP_FILE_SUFFIXES = {".log", ".pyc", ".pyd", ".pyo", ".so", ".swp", ".tmp"}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SyncHasher:
|
||||
"""Own persistent content hashes and cooperative cancellation checks."""
|
||||
|
||||
settings: ModalSyncSettings
|
||||
cancellation_check: CancellationCheck | None = None
|
||||
_hash_cache: dict[str, dict[str, Any]] = field(init=False, default_factory=dict)
|
||||
_hash_cache_dirty: bool = field(init=False, default=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Load the persistent hash cache once for this synchronization engine."""
|
||||
self._hash_cache = self._load_hash_cache()
|
||||
|
||||
def _hash_file(self, path: Path) -> str:
|
||||
"""Compute the SHA256 digest for a file."""
|
||||
self._raise_if_cancelled()
|
||||
resolved_path = path.resolve()
|
||||
stat_result = resolved_path.stat()
|
||||
cache_key = str(resolved_path)
|
||||
cache_entry = self._hash_cache.get(cache_key)
|
||||
if (
|
||||
cache_entry is not None
|
||||
and cache_entry.get("kind") == "file"
|
||||
and cache_entry.get("size") == stat_result.st_size
|
||||
and cache_entry.get("mtime_ns") == stat_result.st_mtime_ns
|
||||
):
|
||||
return str(cache_entry["sha256"])
|
||||
|
||||
digest = hashlib.sha256()
|
||||
with resolved_path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
self._raise_if_cancelled()
|
||||
digest.update(chunk)
|
||||
sha256 = digest.hexdigest()
|
||||
self._hash_cache[cache_key] = {
|
||||
"kind": "file",
|
||||
"size": stat_result.st_size,
|
||||
"mtime_ns": stat_result.st_mtime_ns,
|
||||
"sha256": sha256,
|
||||
}
|
||||
self._mark_hash_cache_dirty()
|
||||
return sha256
|
||||
|
||||
def _raise_if_cancelled(self) -> None:
|
||||
"""Abort queue-time synchronization when its prompt was cancelled."""
|
||||
if self.cancellation_check is not None and self.cancellation_check():
|
||||
raise SyncCancelledError("Remote workflow preparation was cancelled.")
|
||||
|
||||
def _hash_directory(self, path: Path) -> str:
|
||||
"""Compute a stable SHA256 digest for a directory tree."""
|
||||
hash_started_at = time.perf_counter()
|
||||
resolved_path = path.resolve()
|
||||
digest = hashlib.sha256()
|
||||
files = sorted(self._iter_files(resolved_path), key=lambda item: item.relative_to(resolved_path).as_posix())
|
||||
logger.info("Hashing %d files under %s", len(files), resolved_path)
|
||||
fingerprint = self._directory_fingerprint(resolved_path, files)
|
||||
cache_key = f"dir::{resolved_path}"
|
||||
cache_entry = self._hash_cache.get(cache_key)
|
||||
if (
|
||||
cache_entry is not None
|
||||
and cache_entry.get("kind") == "dir"
|
||||
and cache_entry.get("fingerprint") == fingerprint
|
||||
):
|
||||
logger.info(
|
||||
"Reused cached directory hash for %s over %d files in %.3fs.",
|
||||
resolved_path,
|
||||
len(files),
|
||||
time.perf_counter() - hash_started_at,
|
||||
)
|
||||
return str(cache_entry["sha256"])
|
||||
|
||||
for child in files:
|
||||
relative_path = child.relative_to(resolved_path).as_posix()
|
||||
digest.update(relative_path.encode("utf-8"))
|
||||
digest.update(b"\0")
|
||||
digest.update(self._hash_file(child).encode("ascii"))
|
||||
digest.update(b"\0")
|
||||
sha256 = digest.hexdigest()
|
||||
self._hash_cache[cache_key] = {
|
||||
"kind": "dir",
|
||||
"fingerprint": fingerprint,
|
||||
"sha256": sha256,
|
||||
}
|
||||
self._mark_hash_cache_dirty()
|
||||
logger.info(
|
||||
"Computed directory hash for %s over %d files in %.3fs.",
|
||||
resolved_path,
|
||||
len(files),
|
||||
time.perf_counter() - hash_started_at,
|
||||
)
|
||||
return sha256
|
||||
|
||||
def _hash_file_group(self, root: Path, files: list[Path]) -> str:
|
||||
"""Compute a stable digest for a selected file subset rooted under one directory."""
|
||||
digest = hashlib.sha256()
|
||||
for child in sorted(files, key=lambda item: item.relative_to(root).as_posix()):
|
||||
relative_path = child.relative_to(root).as_posix()
|
||||
digest.update(relative_path.encode("utf-8"))
|
||||
digest.update(b"\0")
|
||||
digest.update(self._hash_file(child).encode("ascii"))
|
||||
digest.update(b"\0")
|
||||
return digest.hexdigest()
|
||||
def _directory_fingerprint(self, root: Path, files: list[Path]) -> str:
|
||||
"""Return a metadata-only fingerprint for a directory tree."""
|
||||
digest = hashlib.sha256()
|
||||
for child in files:
|
||||
stat_result = child.stat()
|
||||
digest.update(child.relative_to(root).as_posix().encode("utf-8"))
|
||||
digest.update(b"\0")
|
||||
digest.update(str(stat_result.st_size).encode("ascii"))
|
||||
digest.update(b"\0")
|
||||
digest.update(str(stat_result.st_mtime_ns).encode("ascii"))
|
||||
digest.update(b"\0")
|
||||
return digest.hexdigest()
|
||||
|
||||
def _hash_cache_path(self) -> Path:
|
||||
"""Return the on-disk metadata cache path."""
|
||||
return self.settings.local_storage_root / "metadata" / "hash_cache.json"
|
||||
|
||||
def _load_hash_cache(self) -> dict[str, dict[str, Any]]:
|
||||
"""Load the persistent hash cache from disk when available."""
|
||||
cache_path = self._hash_cache_path()
|
||||
if not cache_path.exists():
|
||||
return {}
|
||||
try:
|
||||
payload = json.loads(cache_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
logger.warning("Hash cache at %s is unreadable; rebuilding it from scratch.", cache_path)
|
||||
return {}
|
||||
if not isinstance(payload, dict):
|
||||
return {}
|
||||
return {
|
||||
str(key): value
|
||||
for key, value in payload.items()
|
||||
if isinstance(value, dict)
|
||||
}
|
||||
|
||||
def _mark_hash_cache_dirty(self) -> None:
|
||||
"""Persist the hash cache after it changes."""
|
||||
self._hash_cache_dirty = True
|
||||
self._save_hash_cache()
|
||||
|
||||
def _save_hash_cache(self) -> None:
|
||||
"""Write the persistent hash cache to disk."""
|
||||
if not self._hash_cache_dirty:
|
||||
return
|
||||
cache_path = self._hash_cache_path()
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cache_path.write_text(json.dumps(self._hash_cache, sort_keys=True), encoding="utf-8")
|
||||
self._hash_cache_dirty = False
|
||||
|
||||
def _iter_files(self, path: Path) -> list[Path]:
|
||||
"""Yield files from a directory tree while skipping cache folders."""
|
||||
files: list[Path] = []
|
||||
for root, dirnames, filenames in os.walk(path):
|
||||
dirnames[:] = [name for name in dirnames if name not in _SKIP_DIRS]
|
||||
for filename in filenames:
|
||||
child = Path(root) / filename
|
||||
if child.suffix.lower() in _SKIP_FILE_SUFFIXES:
|
||||
continue
|
||||
files.append(child)
|
||||
return files
|
||||
|
||||
@@ -0,0 +1,269 @@
|
||||
"""Protocols and value objects shared by asset synchronization components."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Protocol, runtime_checkable
|
||||
|
||||
if __package__:
|
||||
from .huggingface_assets import HuggingFaceAssetSource
|
||||
from .r2_cache import (
|
||||
R2DownloadRequest,
|
||||
R2UploadPlan,
|
||||
R2UploadResult,
|
||||
R2WorkerPreflightRequest,
|
||||
)
|
||||
else: # pragma: no cover - flat import inside the Modal container.
|
||||
from huggingface_assets import HuggingFaceAssetSource
|
||||
from r2_cache import (
|
||||
R2DownloadRequest,
|
||||
R2UploadPlan,
|
||||
R2UploadResult,
|
||||
R2WorkerPreflightRequest,
|
||||
)
|
||||
|
||||
SyncStatusCallback = Callable[[str, int | None, int | None], None]
|
||||
CancellationCheck = Callable[[], bool]
|
||||
|
||||
class VolumeBackend(Protocol):
|
||||
"""Minimal storage interface needed by the sync engine."""
|
||||
|
||||
def exists(self, remote_path: str) -> bool:
|
||||
"""Return whether the remote path already exists."""
|
||||
|
||||
def put_file(self, local_path: Path, remote_path: str) -> None:
|
||||
"""Upload a local file into the remote storage backend."""
|
||||
|
||||
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
||||
"""Upload raw bytes into the remote storage backend."""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class HuggingFaceMaterializingBackend(Protocol):
|
||||
"""Optional backend capability for direct verified Hugging Face acquisition."""
|
||||
|
||||
def materialize_huggingface_file(
|
||||
self,
|
||||
source: HuggingFaceAssetSource,
|
||||
remote_path: str,
|
||||
*,
|
||||
token: str | None,
|
||||
) -> bool:
|
||||
"""Materialize one immutable file and report whether it succeeded."""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class CancellableVolumeBackend(Protocol):
|
||||
"""Optional backend capability for interruptible file uploads."""
|
||||
|
||||
def put_file_cancellable(
|
||||
self,
|
||||
local_path: Path,
|
||||
remote_path: str,
|
||||
*,
|
||||
cancellation_check: CancellationCheck,
|
||||
) -> None:
|
||||
"""Upload a file while cooperatively observing cancellation."""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class CancellableHuggingFaceMaterializingBackend(Protocol):
|
||||
"""Optional backend capability for interruptible remote downloads."""
|
||||
|
||||
def materialize_huggingface_file_cancellable(
|
||||
self,
|
||||
source: HuggingFaceAssetSource,
|
||||
remote_path: str,
|
||||
*,
|
||||
token: str | None,
|
||||
cancellation_check: CancellationCheck,
|
||||
) -> bool:
|
||||
"""Materialize one immutable file while observing cancellation."""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class R2MaterializingBackend(Protocol):
|
||||
"""Optional backend capability for signed R2 downloads and write-back."""
|
||||
|
||||
def materialize_r2_file(
|
||||
self,
|
||||
request: R2DownloadRequest,
|
||||
remote_path: str,
|
||||
*,
|
||||
cancellation_check: CancellationCheck | None = None,
|
||||
) -> None:
|
||||
"""Download and verify one content-addressed R2 object."""
|
||||
|
||||
def upload_r2_file(
|
||||
self,
|
||||
plan: R2UploadPlan,
|
||||
remote_path: str,
|
||||
) -> R2UploadResult:
|
||||
"""Upload one remote file through a controller-issued signed plan."""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class CancellableR2WriteBackBackend(Protocol):
|
||||
"""Optional backend capability for preemptible background R2 uploads."""
|
||||
|
||||
def upload_r2_file_cancellable(
|
||||
self,
|
||||
plan: R2UploadPlan,
|
||||
remote_path: str,
|
||||
*,
|
||||
cancellation_check: CancellationCheck,
|
||||
) -> R2UploadResult:
|
||||
"""Upload one remote file while yielding promptly to workflow activity."""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class R2WorkerPreflightBackend(Protocol):
|
||||
"""Optional backend capability for testing effective worker-side R2 access."""
|
||||
|
||||
def preflight_r2_access(
|
||||
self,
|
||||
request: R2WorkerPreflightRequest,
|
||||
*,
|
||||
cancellation_check: CancellationCheck | None = None,
|
||||
) -> None:
|
||||
"""Verify one read-only presigned request from the worker environment."""
|
||||
|
||||
|
||||
class SyncCancelledError(RuntimeError):
|
||||
"""Raised when queue-time asset preparation is cancelled."""
|
||||
|
||||
|
||||
class SyncIndexBackend(Protocol):
|
||||
"""Minimal metadata index interface used to deduplicate deterministic uploads."""
|
||||
|
||||
def get(self, key: str) -> dict[str, Any] | None:
|
||||
"""Return one stored sync record when it exists."""
|
||||
|
||||
def put(self, key: str, value: dict[str, Any]) -> None:
|
||||
"""Persist one sync record under the provided key."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SyncedAsset:
|
||||
"""Description of a local asset mirrored into the remote storage root."""
|
||||
|
||||
local_path: Path
|
||||
remote_path: str
|
||||
sha256: str
|
||||
uploaded: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class AssetSyncRequestCache:
|
||||
"""Deduplicate asset hashing, index lookup, and upload within one queued prompt."""
|
||||
|
||||
planned_paths: tuple[Path, ...]
|
||||
_assets_by_path: dict[Path, SyncedAsset] = field(default_factory=dict)
|
||||
_positions_by_path: dict[Path, int] = field(init=False, default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Index stable request progress positions once."""
|
||||
self._positions_by_path = {
|
||||
path.resolve(): index
|
||||
for index, path in enumerate(self.planned_paths, start=1)
|
||||
}
|
||||
|
||||
def get(self, local_path: Path) -> SyncedAsset | None:
|
||||
"""Return a previously synced asset for this request when available."""
|
||||
return self._assets_by_path.get(local_path.resolve())
|
||||
|
||||
def put(self, synced_asset: SyncedAsset) -> None:
|
||||
"""Remember one request-scoped sync result by its resolved local path."""
|
||||
self._assets_by_path[synced_asset.local_path.resolve()] = synced_asset
|
||||
|
||||
def progress(self, local_path: Path) -> tuple[int, int]:
|
||||
"""Return the stable one-based progress position for a planned asset."""
|
||||
resolved_path = local_path.resolve()
|
||||
return self._positions_by_path[resolved_path], len(self.planned_paths)
|
||||
|
||||
def synced_assets(self) -> tuple[SyncedAsset, ...]:
|
||||
"""Return unique sync results in request planning order."""
|
||||
return tuple(
|
||||
self._assets_by_path[path.resolve()]
|
||||
for path in self.planned_paths
|
||||
if path.resolve() in self._assets_by_path
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CustomNodesArchiveSpec:
|
||||
"""Deterministic archive spec for one top-level custom_nodes payload slice."""
|
||||
|
||||
entry_name: str
|
||||
display_name: str
|
||||
source_description: str
|
||||
files: tuple[Path, ...]
|
||||
sha256: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CustomNodesArchiveSyncResult:
|
||||
"""Describe the sync result for one top-level custom_nodes archive."""
|
||||
|
||||
entry_name: str
|
||||
display_name: str
|
||||
sha256: str
|
||||
remote_path: str
|
||||
uploaded: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CustomNodeAssetSpec:
|
||||
"""Describe one package-owned model asset excluded from a code archive."""
|
||||
|
||||
entry_name: str
|
||||
local_path: Path
|
||||
relative_path: str
|
||||
sha256: str
|
||||
size_bytes: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CustomNodeAssetSyncResult:
|
||||
"""Describe one content-addressed package asset and its upload result."""
|
||||
|
||||
entry_name: str
|
||||
relative_path: str
|
||||
sha256: str
|
||||
size_bytes: int
|
||||
remote_path: str
|
||||
uploaded: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ContentAddressedSyncResult:
|
||||
"""Describe the outcome of one content-addressed file sync decision."""
|
||||
|
||||
remote_path: str
|
||||
uploaded: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ContentAddressedSyncSpec:
|
||||
"""Collect one file's immutable identity, destination, and progress context."""
|
||||
|
||||
local_path: Path
|
||||
remote_path: str
|
||||
sha256: str
|
||||
sync_key: str
|
||||
source_description: str
|
||||
status_callback: SyncStatusCallback | None
|
||||
upload_status_message: str | None
|
||||
status_current: int | None
|
||||
status_total: int | None
|
||||
huggingface_source: HuggingFaceAssetSource | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _R2MaterializationOutcome:
|
||||
"""Return an optional hit and whether failed cached bytes need replacement."""
|
||||
|
||||
result: _ContentAddressedSyncResult | None
|
||||
refresh_required: bool = False
|
||||
|
||||
@@ -0,0 +1,488 @@
|
||||
"""Idle-gated, deduplicated Cloudflare R2 write-back coordination."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import OrderedDict
|
||||
from contextlib import AbstractContextManager, nullcontext
|
||||
from dataclasses import dataclass, field
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable, Protocol
|
||||
|
||||
if __package__:
|
||||
from .r2_cache import (
|
||||
R2CacheClient,
|
||||
R2CacheError,
|
||||
R2DownloadRequest,
|
||||
R2UploadPlan,
|
||||
)
|
||||
from .sync_protocols import (
|
||||
CancellationCheck,
|
||||
CancellableR2WriteBackBackend,
|
||||
R2MaterializingBackend,
|
||||
R2WorkerPreflightBackend,
|
||||
SyncStatusCallback,
|
||||
_ContentAddressedSyncResult,
|
||||
_ContentAddressedSyncSpec,
|
||||
_R2MaterializationOutcome,
|
||||
)
|
||||
else: # pragma: no cover - flat import inside the Modal container.
|
||||
from r2_cache import R2CacheClient, R2CacheError, R2DownloadRequest, R2UploadPlan
|
||||
from sync_protocols import (
|
||||
CancellationCheck,
|
||||
CancellableR2WriteBackBackend,
|
||||
R2MaterializingBackend,
|
||||
R2WorkerPreflightBackend,
|
||||
SyncStatusCallback,
|
||||
_ContentAddressedSyncResult,
|
||||
_ContentAddressedSyncSpec,
|
||||
_R2MaterializationOutcome,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class R2WriteBackCancelled(RuntimeError):
|
||||
"""Signal that an idle write-back yielded to foreground workflow activity."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _R2WriteBackJob:
|
||||
"""Describe one deduplicated best-effort cache population operation."""
|
||||
|
||||
key: tuple[str, ...]
|
||||
callback: Callable[[CancellationCheck], None] = field(compare=False, repr=False)
|
||||
|
||||
|
||||
class R2WriteBackCoordinator:
|
||||
"""Run cache population only while no remote workflow is active."""
|
||||
|
||||
def __init__(self, *, max_workers: int = 2, max_pending_jobs: int = 1024) -> None:
|
||||
"""Start bounded daemon workers for opportunistic write-back jobs."""
|
||||
if max_workers <= 0 or max_pending_jobs <= 0:
|
||||
raise ValueError("R2 write-back coordinator bounds must be positive.")
|
||||
self._condition = threading.Condition()
|
||||
self._pending: OrderedDict[tuple[str, ...], _R2WriteBackJob] = OrderedDict()
|
||||
self._known_keys: set[tuple[str, ...]] = set()
|
||||
self._active_prompts: set[str] = set()
|
||||
self._active_cancellations: set[threading.Event] = set()
|
||||
self._active_jobs = 0
|
||||
self._max_pending_jobs = max_pending_jobs
|
||||
self._workers = tuple(
|
||||
threading.Thread(
|
||||
target=self._worker,
|
||||
name=f"comfy-r2-writeback-{worker_index + 1}",
|
||||
daemon=True,
|
||||
)
|
||||
for worker_index in range(max_workers)
|
||||
)
|
||||
for worker in self._workers:
|
||||
worker.start()
|
||||
|
||||
def begin_prompt(self, prompt_id: str) -> None:
|
||||
"""Reserve foreground priority and preempt active cache transfers."""
|
||||
normalized_prompt_id = str(prompt_id).strip()
|
||||
if not normalized_prompt_id:
|
||||
raise ValueError("R2 write-back prompt ID must not be empty.")
|
||||
with self._condition:
|
||||
self._active_prompts.add(normalized_prompt_id)
|
||||
for cancellation in self._active_cancellations:
|
||||
cancellation.set()
|
||||
self._condition.notify_all()
|
||||
|
||||
def finish_prompt(self, prompt_id: str) -> None:
|
||||
"""Release one prompt reservation and wake idle cache workers."""
|
||||
normalized_prompt_id = str(prompt_id).strip()
|
||||
if not normalized_prompt_id:
|
||||
return
|
||||
with self._condition:
|
||||
self._active_prompts.discard(normalized_prompt_id)
|
||||
self._condition.notify_all()
|
||||
|
||||
def reset_prompt_reservations_for_tests(self) -> None:
|
||||
"""Clear synthetic prompt reservations left by isolated queue tests."""
|
||||
with self._condition:
|
||||
self._active_prompts.clear()
|
||||
for cancellation in self._active_cancellations:
|
||||
cancellation.set()
|
||||
self._condition.notify_all()
|
||||
|
||||
def submit(
|
||||
self,
|
||||
key: tuple[str, ...],
|
||||
callback: Callable[[CancellationCheck], None],
|
||||
) -> bool:
|
||||
"""Queue one deduplicated job without waiting for capacity or execution."""
|
||||
if not key or any(not str(part) for part in key):
|
||||
raise ValueError("R2 write-back job key must contain non-empty values.")
|
||||
with self._condition:
|
||||
if key in self._known_keys:
|
||||
return False
|
||||
if len(self._pending) >= self._max_pending_jobs:
|
||||
logger.warning(
|
||||
"Dropping R2 write-back because the idle queue reached %d jobs key=%s.",
|
||||
self._max_pending_jobs,
|
||||
key[-1],
|
||||
)
|
||||
return False
|
||||
job = _R2WriteBackJob(key=key, callback=callback)
|
||||
self._pending[key] = job
|
||||
self._known_keys.add(key)
|
||||
self._condition.notify()
|
||||
return True
|
||||
|
||||
def wait_for_idle(self, timeout_seconds: float | None = None) -> bool:
|
||||
"""Wait until pending and active jobs drain, primarily for tests and shutdown."""
|
||||
deadline = (
|
||||
None if timeout_seconds is None else time.monotonic() + timeout_seconds
|
||||
)
|
||||
with self._condition:
|
||||
while self._pending or self._active_jobs:
|
||||
if deadline is None:
|
||||
self._condition.wait()
|
||||
continue
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0.0:
|
||||
return False
|
||||
self._condition.wait(remaining)
|
||||
return True
|
||||
|
||||
def _worker(self) -> None:
|
||||
"""Run queued jobs only during foreground-idle windows."""
|
||||
while True:
|
||||
with self._condition:
|
||||
self._condition.wait_for(
|
||||
lambda: bool(self._pending) and not self._active_prompts
|
||||
)
|
||||
key, job = self._pending.popitem(last=False)
|
||||
cancellation = threading.Event()
|
||||
self._active_cancellations.add(cancellation)
|
||||
self._active_jobs += 1
|
||||
cancelled = False
|
||||
try:
|
||||
job.callback(cancellation.is_set)
|
||||
except R2WriteBackCancelled:
|
||||
cancelled = True
|
||||
except (
|
||||
AssertionError,
|
||||
KeyError,
|
||||
OSError,
|
||||
RuntimeError,
|
||||
TypeError,
|
||||
ValueError,
|
||||
) as exc:
|
||||
logger.warning(
|
||||
"Unhandled R2 write-back job failure key=%s: %s",
|
||||
key[-1],
|
||||
exc,
|
||||
)
|
||||
finally:
|
||||
with self._condition:
|
||||
self._active_cancellations.discard(cancellation)
|
||||
self._active_jobs -= 1
|
||||
if cancelled:
|
||||
self._pending.setdefault(key, job)
|
||||
else:
|
||||
self._known_keys.discard(key)
|
||||
self._condition.notify_all()
|
||||
|
||||
|
||||
_R2_WRITE_BACK_COORDINATOR = R2WriteBackCoordinator()
|
||||
|
||||
|
||||
def begin_r2_writeback_prompt(prompt_id: str) -> None:
|
||||
"""Pause opportunistic cache population for one foreground prompt."""
|
||||
_R2_WRITE_BACK_COORDINATOR.begin_prompt(prompt_id)
|
||||
|
||||
|
||||
def finish_r2_writeback_prompt(prompt_id: str) -> None:
|
||||
"""Resume cache population after one foreground prompt leaves the queue."""
|
||||
_R2_WRITE_BACK_COORDINATOR.finish_prompt(prompt_id)
|
||||
|
||||
|
||||
def _never_cancel() -> bool:
|
||||
"""Return false for write-backs without a cooperative cancellation source."""
|
||||
return False
|
||||
|
||||
|
||||
def _emit_sync_status(
|
||||
status_callback: SyncStatusCallback | None,
|
||||
message: str,
|
||||
current: int | None = None,
|
||||
total: int | None = None,
|
||||
) -> None:
|
||||
"""Emit one synchronization status update when available."""
|
||||
if status_callback is not None:
|
||||
status_callback(message, current, total)
|
||||
|
||||
|
||||
def _format_r2_download_status(
|
||||
asset_name: str,
|
||||
*,
|
||||
item_index: int | None,
|
||||
total_items: int | None,
|
||||
) -> str:
|
||||
"""Return one queue-time worker-side R2 download status."""
|
||||
if item_index is not None and total_items is not None and total_items > 1:
|
||||
return (
|
||||
f"Downloading asset {item_index}/{total_items} from Cloudflare R2: "
|
||||
f"{asset_name}"
|
||||
)
|
||||
return f"Downloading asset from Cloudflare R2: {asset_name}"
|
||||
|
||||
|
||||
class R2TransferHost(Protocol):
|
||||
"""Define engine services used by R2 transfer coordination."""
|
||||
|
||||
def _destination_label(self) -> str:
|
||||
"""Return a user-facing worker destination."""
|
||||
|
||||
def _raise_if_cancelled(self) -> None:
|
||||
"""Raise when foreground synchronization was cancelled."""
|
||||
|
||||
def _store_sync_record(
|
||||
self,
|
||||
*,
|
||||
sync_key: str,
|
||||
remote_path: str,
|
||||
source_description: str,
|
||||
) -> None:
|
||||
"""Persist one successful synchronization record."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class R2TransferManager:
|
||||
"""Own R2 preflight, read-through, and background write-back behavior."""
|
||||
|
||||
_host: R2TransferHost
|
||||
volume: Any
|
||||
r2_cache: R2CacheClient | None = None
|
||||
cancellation_check: CancellationCheck | None = None
|
||||
r2_writeback_activity: Callable[[], AbstractContextManager[None]] | None = None
|
||||
|
||||
def preflight_r2_access(
|
||||
self,
|
||||
*,
|
||||
status_callback: SyncStatusCallback | None = None,
|
||||
) -> None:
|
||||
"""Fail before asset transfer when a worker cannot use configured R2 URLs."""
|
||||
self._host._raise_if_cancelled()
|
||||
if self.r2_cache is None or not isinstance(
|
||||
self.volume,
|
||||
R2WorkerPreflightBackend,
|
||||
):
|
||||
return
|
||||
destination = self._host._destination_label()
|
||||
_emit_sync_status(
|
||||
status_callback,
|
||||
f"Checking Cloudflare R2 access from {destination}",
|
||||
)
|
||||
try:
|
||||
request = self.r2_cache.worker_preflight_request()
|
||||
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
raise R2CacheError(
|
||||
"Cloudflare R2 controller validation failed before the worker "
|
||||
f"preflight: {exc}"
|
||||
) from exc
|
||||
try:
|
||||
self.volume.preflight_r2_access(
|
||||
request,
|
||||
cancellation_check=self.cancellation_check,
|
||||
)
|
||||
except InterruptedError as exc:
|
||||
self._host._raise_if_cancelled()
|
||||
raise R2CacheError(
|
||||
f"Cloudflare R2 worker preflight was interrupted for {destination}."
|
||||
) from exc
|
||||
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
diagnostic = str(exc)
|
||||
hint = (
|
||||
" Remove Client IP Address Filtering from the bucket-scoped R2 "
|
||||
"API token; dynamic worker egress addresses cannot be reliably "
|
||||
"allowlisted."
|
||||
if "r2_code=AccessDenied" in diagnostic
|
||||
else ""
|
||||
)
|
||||
raise R2CacheError(
|
||||
f"Cloudflare R2 is reachable from the controller but not from "
|
||||
f"{destination}.{hint} Safe worker diagnostic: {diagnostic}"
|
||||
) from exc
|
||||
self._host._raise_if_cancelled()
|
||||
logger.info(
|
||||
"Validated Cloudflare R2 access from %s before asset transfer.",
|
||||
destination,
|
||||
)
|
||||
|
||||
def _materialize_r2_source(
|
||||
self,
|
||||
spec: _ContentAddressedSyncSpec,
|
||||
size_bytes: int,
|
||||
) -> _R2MaterializationOutcome:
|
||||
"""Try one signed worker-side R2 download and preserve upload fallback."""
|
||||
if self.r2_cache is None or not isinstance(self.volume, R2MaterializingBackend):
|
||||
return _R2MaterializationOutcome(result=None)
|
||||
request: R2DownloadRequest | None = None
|
||||
try:
|
||||
request = self.r2_cache.download_request(spec.sha256, size_bytes)
|
||||
if request is None:
|
||||
return _R2MaterializationOutcome(result=None)
|
||||
_emit_sync_status(
|
||||
spec.status_callback,
|
||||
_format_r2_download_status(
|
||||
spec.local_path.name,
|
||||
item_index=spec.status_current,
|
||||
total_items=spec.status_total,
|
||||
),
|
||||
spec.status_current,
|
||||
spec.status_total,
|
||||
)
|
||||
self.volume.materialize_r2_file(
|
||||
request,
|
||||
spec.remote_path,
|
||||
cancellation_check=self.cancellation_check,
|
||||
)
|
||||
except InterruptedError as exc:
|
||||
raise SyncCancelledError("Remote workflow preparation was cancelled.") from exc
|
||||
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
logger.warning(
|
||||
"Cloudflare R2 materialization failed for SHA-256 %s; "
|
||||
"falling back to the existing upload path: %s",
|
||||
spec.sha256,
|
||||
exc,
|
||||
)
|
||||
return _R2MaterializationOutcome(
|
||||
result=None,
|
||||
refresh_required=request is not None,
|
||||
)
|
||||
self._host._store_sync_record(
|
||||
sync_key=spec.sync_key,
|
||||
remote_path=spec.remote_path,
|
||||
source_description=f"Cloudflare R2 SHA-256 {spec.sha256}",
|
||||
)
|
||||
logger.info(
|
||||
"Materialized SHA-256 %s from Cloudflare R2 at %s.",
|
||||
spec.sha256,
|
||||
spec.remote_path,
|
||||
)
|
||||
return _R2MaterializationOutcome(
|
||||
result=_ContentAddressedSyncResult(
|
||||
remote_path=spec.remote_path,
|
||||
uploaded=True,
|
||||
)
|
||||
)
|
||||
|
||||
def _schedule_r2_writeback(
|
||||
self,
|
||||
*,
|
||||
sha256: str,
|
||||
size_bytes: int,
|
||||
remote_path: str,
|
||||
force: bool = False,
|
||||
) -> None:
|
||||
"""Queue one best-effort R2 write without delaying foreground progress."""
|
||||
if not self._r2_writeback_enabled():
|
||||
return
|
||||
assert self.r2_cache is not None
|
||||
job_key = (
|
||||
*self._r2_writeback_job_prefix(),
|
||||
sha256,
|
||||
str(size_bytes),
|
||||
remote_path,
|
||||
"force" if force else "normal",
|
||||
)
|
||||
_R2_WRITE_BACK_COORDINATOR.submit(
|
||||
job_key,
|
||||
lambda cancellation_check: self._write_back_r2_file(
|
||||
sha256,
|
||||
size_bytes,
|
||||
remote_path,
|
||||
force,
|
||||
cancellation_check=cancellation_check,
|
||||
),
|
||||
)
|
||||
|
||||
def _r2_writeback_job_prefix(self) -> tuple[str, str, str]:
|
||||
"""Return a stable cache namespace used to deduplicate background jobs."""
|
||||
if self.r2_cache is None:
|
||||
raise RuntimeError("R2 write-back job requested without an R2 cache.")
|
||||
configuration = getattr(self.r2_cache, "configuration", None)
|
||||
if configuration is not None:
|
||||
return (
|
||||
str(configuration.endpoint_url),
|
||||
str(configuration.bucket),
|
||||
str(configuration.key_prefix),
|
||||
)
|
||||
return (
|
||||
type(self.r2_cache).__module__,
|
||||
type(self.r2_cache).__qualname__,
|
||||
str(id(self.r2_cache)),
|
||||
)
|
||||
|
||||
def _r2_writeback_enabled(self) -> bool:
|
||||
"""Return whether this engine can populate its configured R2 cache."""
|
||||
return bool(
|
||||
self.r2_cache is not None
|
||||
and self.r2_cache.write_back_mode != "off"
|
||||
and isinstance(self.volume, R2MaterializingBackend)
|
||||
)
|
||||
|
||||
def _write_back_r2_file(
|
||||
self,
|
||||
sha256: str,
|
||||
size_bytes: int,
|
||||
remote_path: str,
|
||||
force: bool = False,
|
||||
*,
|
||||
cancellation_check: CancellationCheck = _never_cancel,
|
||||
) -> None:
|
||||
"""Upload one worker-resident file to R2 without exposing permanent keys."""
|
||||
if self.r2_cache is None or not isinstance(
|
||||
self.volume,
|
||||
R2MaterializingBackend,
|
||||
):
|
||||
return
|
||||
plan: R2UploadPlan | None = None
|
||||
try:
|
||||
if cancellation_check():
|
||||
raise R2WriteBackCancelled("R2 write-back yielded before activity.")
|
||||
activity = (
|
||||
self.r2_writeback_activity()
|
||||
if self.r2_writeback_activity is not None
|
||||
else nullcontext()
|
||||
)
|
||||
with activity:
|
||||
if cancellation_check():
|
||||
raise R2WriteBackCancelled("R2 write-back yielded before signing.")
|
||||
plan = self.r2_cache.prepare_upload(sha256, size_bytes, force=force)
|
||||
if plan is None:
|
||||
return
|
||||
if cancellation_check():
|
||||
raise R2WriteBackCancelled("R2 write-back yielded after signing.")
|
||||
if isinstance(self.volume, CancellableR2WriteBackBackend):
|
||||
result = self.volume.upload_r2_file_cancellable(
|
||||
plan,
|
||||
remote_path,
|
||||
cancellation_check=cancellation_check,
|
||||
)
|
||||
else:
|
||||
result = self.volume.upload_r2_file(plan, remote_path)
|
||||
if cancellation_check():
|
||||
raise R2WriteBackCancelled(
|
||||
"R2 write-back yielded before completion."
|
||||
)
|
||||
self.r2_cache.complete_upload(plan, result)
|
||||
logger.info("Wrote SHA-256 %s back to Cloudflare R2.", sha256)
|
||||
except (InterruptedError, R2WriteBackCancelled) as exc:
|
||||
if plan is not None:
|
||||
self.r2_cache.abort_upload(plan)
|
||||
raise R2WriteBackCancelled(str(exc)) from exc
|
||||
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
if plan is not None:
|
||||
self.r2_cache.abort_upload(plan)
|
||||
logger.warning("Cloudflare R2 write-back failed for SHA-256 %s: %s", sha256, exc)
|
||||
|
||||
def wait_for_r2_writebacks(self) -> None:
|
||||
"""Wait for this engine's currently scheduled R2 writes to finish."""
|
||||
_R2_WRITE_BACK_COORDINATOR.wait_for_idle()
|
||||
|
||||
+26
-2
@@ -565,6 +565,30 @@ def sync_engine_module(extension_package: object) -> object:
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.sync_engine")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sync_backends_module(extension_package: object) -> object:
|
||||
"""Return local and Modal synchronization backends."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.sync_backends")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sync_hashing_module(extension_package: object) -> object:
|
||||
"""Return the persistent synchronization hashing module."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.sync_hashing")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sync_custom_nodes_module(extension_package: object) -> object:
|
||||
"""Return the custom-node synchronization module."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.sync_custom_nodes")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def sync_r2_transfer_module(extension_package: object) -> object:
|
||||
"""Return the R2 transfer and write-back coordination module."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.sync_r2_transfer")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def modal_cloud_module() -> object:
|
||||
"""Return the stable Modal cloud entry module."""
|
||||
@@ -598,9 +622,9 @@ def reset_modal_environment(
|
||||
or request.node.get_closest_marker("live_vast") is not None
|
||||
):
|
||||
return
|
||||
sync_engine_module = sys.modules.get(f"{PACKAGE_NAME}.sync_engine")
|
||||
sync_r2_transfer_module = sys.modules.get(f"{PACKAGE_NAME}.sync_r2_transfer")
|
||||
coordinator = getattr(
|
||||
sync_engine_module,
|
||||
sync_r2_transfer_module,
|
||||
"_R2_WRITE_BACK_COORDINATOR",
|
||||
None,
|
||||
)
|
||||
|
||||
+49
-33
@@ -22,7 +22,7 @@ def test_sync_file_deduplicates_by_hash(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Repeated file syncs should reuse the same remote path and sync-index record."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
asset_path = tmp_path / "model.safetensors"
|
||||
asset_path.write_bytes(b"model-bytes")
|
||||
|
||||
@@ -65,7 +65,7 @@ def test_request_asset_cache_syncs_repeated_prompt_asset_once(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Repeated references across remote nodes should share one request sync decision."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
asset_path = tmp_path / "model.safetensors"
|
||||
asset_path.write_bytes(b"model-bytes")
|
||||
settings = settings_module.ModalSyncSettings(
|
||||
@@ -121,7 +121,7 @@ def test_sync_file_emits_upload_status(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Uploading a new asset should emit a status message naming the uploaded file."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
asset_path = tmp_path / "model.safetensors"
|
||||
asset_path.write_bytes(b"model-bytes")
|
||||
observed_statuses: list[tuple[str, int | None, int | None]] = []
|
||||
@@ -164,7 +164,7 @@ def test_sync_file_uses_self_hosted_destination_in_ssh_mode(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""SSH uploads should not describe the destination as Modal."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
asset_path = tmp_path / "model.safetensors"
|
||||
asset_path.write_bytes(b"model-bytes")
|
||||
observed_statuses: list[tuple[str, int | None, int | None]] = []
|
||||
@@ -501,7 +501,7 @@ def test_sync_custom_nodes_directory_creates_archive(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""The sync engine should archive and mirror a custom_nodes directory."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -541,7 +541,7 @@ def test_sync_custom_nodes_directory_only_checks_once_per_engine_lifetime(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""The same sync engine should not rescan custom_nodes after the first successful sync."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -595,7 +595,7 @@ def test_sync_custom_nodes_directory_emits_packaging_and_upload_status(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Building a new custom_nodes archive should report packaging and upload stages."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -640,7 +640,7 @@ def test_sync_custom_nodes_directory_concurrent_calls_share_one_initial_sync(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Concurrent first-run calls should serialize behind one custom_nodes sync decision."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -665,7 +665,7 @@ def test_sync_custom_nodes_directory_concurrent_calls_share_one_initial_sync(
|
||||
)
|
||||
|
||||
engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
|
||||
original_sync = engine._sync_custom_nodes_directory_uncached
|
||||
original_sync = engine._custom_nodes._sync_custom_nodes_directory_uncached
|
||||
sync_call_count = 0
|
||||
sync_call_count_lock = threading.Lock()
|
||||
release_first_sync = threading.Event()
|
||||
@@ -678,7 +678,11 @@ def test_sync_custom_nodes_directory_concurrent_calls_share_one_initial_sync(
|
||||
release_first_sync.wait(timeout=5.0)
|
||||
return original_sync(status_callback=status_callback)
|
||||
|
||||
monkeypatch.setattr(engine, "_sync_custom_nodes_directory_uncached", wrapped_sync)
|
||||
monkeypatch.setattr(
|
||||
engine._custom_nodes,
|
||||
"_sync_custom_nodes_directory_uncached",
|
||||
wrapped_sync,
|
||||
)
|
||||
|
||||
results: list[Any] = [None, None]
|
||||
|
||||
@@ -713,7 +717,7 @@ def test_hash_directory_ignores_virtualenv_and_bytecode_artifacts(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Ignored directories and file suffixes should not affect custom_nodes hashing."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -761,7 +765,7 @@ def test_sync_file_reuses_cached_hash_for_unchanged_file(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Unchanged files should reuse the persisted file digest instead of re-reading the payload."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
asset_path = tmp_path / "model.safetensors"
|
||||
asset_path.write_bytes(b"model-bytes")
|
||||
|
||||
@@ -804,7 +808,7 @@ def test_sync_custom_nodes_directory_reuses_cached_archive(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""An unchanged custom_nodes tree should reuse its digest-keyed local archive."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -852,7 +856,11 @@ def test_sync_custom_nodes_directory_reuses_cached_archive(
|
||||
def fail_create_archive(*args: Any, **kwargs: Any) -> Any:
|
||||
raise AssertionError("Expected cached custom_nodes archive to be reused.")
|
||||
|
||||
monkeypatch.setattr(second_engine, "_create_archive", fail_create_archive)
|
||||
monkeypatch.setattr(
|
||||
second_engine._custom_nodes,
|
||||
"_create_archive",
|
||||
fail_create_archive,
|
||||
)
|
||||
second_bundle = second_engine.sync_custom_nodes_directory()
|
||||
|
||||
assert second_bundle is not None
|
||||
@@ -869,7 +877,7 @@ def test_sync_custom_nodes_separates_code_assets_and_nested_virtualenv(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Code ZIPs should exclude mounted model assets and nested virtual environments."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -927,7 +935,7 @@ def test_sync_file_reuses_sync_index_record_for_existing_remote_payload(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""An indexed deterministic asset should be reused without a second upload."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
asset_path = tmp_path / "model.safetensors"
|
||||
asset_path.write_bytes(b"model-bytes")
|
||||
|
||||
@@ -986,7 +994,7 @@ def test_sync_custom_nodes_directory_reuses_indexed_remote_bundle(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""An indexed hash-named remote custom_nodes bundle should be reused without rebuilding or reuploading."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -1040,7 +1048,7 @@ def test_sync_custom_nodes_directory_reuses_indexed_remote_bundle(
|
||||
def fail_create_archive(*args: Any, **kwargs: Any) -> Any:
|
||||
raise AssertionError("Expected the indexed hash-named remote bundle to be reused.")
|
||||
|
||||
monkeypatch.setattr(engine, "_create_archive", fail_create_archive)
|
||||
monkeypatch.setattr(engine._custom_nodes, "_create_archive", fail_create_archive)
|
||||
bundle = engine.sync_custom_nodes_directory()
|
||||
|
||||
assert bundle is not None
|
||||
@@ -1057,7 +1065,7 @@ def test_sync_custom_nodes_directory_only_rebuilds_changed_top_level_archive(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Changing one custom_nodes package should only rebuild that package archive plus the manifest."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_a = custom_nodes_dir / "example_a"
|
||||
package_b = custom_nodes_dir / "example_b"
|
||||
@@ -1107,7 +1115,7 @@ def test_sync_custom_nodes_directory_only_rebuilds_changed_top_level_archive(
|
||||
settings=settings,
|
||||
)
|
||||
rebuilt_entries: list[str] = []
|
||||
original_create_archive = second_engine._create_archive_from_files
|
||||
original_create_archive = second_engine._custom_nodes._create_archive_from_files
|
||||
|
||||
def record_create_archive(root_path: Path, files: list[Path], archive_path: Path) -> Path:
|
||||
"""Record which top-level package archive had to be rebuilt."""
|
||||
@@ -1115,7 +1123,11 @@ def test_sync_custom_nodes_directory_only_rebuilds_changed_top_level_archive(
|
||||
rebuilt_entries.append(files[0].relative_to(custom_nodes_dir).parts[0])
|
||||
return original_create_archive(custom_nodes_dir, files, archive_path)
|
||||
|
||||
monkeypatch.setattr(second_engine, "_create_archive_from_files", record_create_archive)
|
||||
monkeypatch.setattr(
|
||||
second_engine._custom_nodes,
|
||||
"_create_archive_from_files",
|
||||
record_create_archive,
|
||||
)
|
||||
second_bundle = second_engine.sync_custom_nodes_directory()
|
||||
|
||||
assert second_bundle is not None
|
||||
@@ -1130,7 +1142,7 @@ def test_sync_custom_nodes_directory_builds_multiple_archives_in_parallel(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Fresh per-package custom_nodes archives should build in parallel."""
|
||||
monkeypatch.setattr(sync_engine_module, "modal", None)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", None)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_a = custom_nodes_dir / "example_a"
|
||||
package_b = custom_nodes_dir / "example_b"
|
||||
@@ -1158,7 +1170,7 @@ def test_sync_custom_nodes_directory_builds_multiple_archives_in_parallel(
|
||||
)
|
||||
|
||||
engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
|
||||
original_create_archive = engine._create_archive_from_files
|
||||
original_create_archive = engine._custom_nodes._create_archive_from_files
|
||||
thread_ids: set[int] = set()
|
||||
started_count = 0
|
||||
started_lock = threading.Lock()
|
||||
@@ -1177,7 +1189,11 @@ def test_sync_custom_nodes_directory_builds_multiple_archives_in_parallel(
|
||||
time.sleep(0.02)
|
||||
return original_create_archive(custom_nodes_dir, files, archive_path)
|
||||
|
||||
monkeypatch.setattr(engine, "_create_archive_from_files", record_create_archive)
|
||||
monkeypatch.setattr(
|
||||
engine._custom_nodes,
|
||||
"_create_archive_from_files",
|
||||
record_create_archive,
|
||||
)
|
||||
bundle = engine.sync_custom_nodes_directory()
|
||||
|
||||
assert bundle is not None
|
||||
@@ -1249,7 +1265,7 @@ def test_remote_mode_uses_modal_volume_backend_when_sdk_is_available(
|
||||
"""Return the fake volume for any lookup."""
|
||||
return fake_volume
|
||||
|
||||
monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", FakeModal)
|
||||
|
||||
settings = settings_module.ModalSyncSettings(
|
||||
app_name="app",
|
||||
@@ -1361,7 +1377,7 @@ def test_remote_sync_index_discards_stale_volume_epoch_and_reuploads_missing_pay
|
||||
del name, create_if_missing
|
||||
return fake_volume
|
||||
|
||||
monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", FakeModal)
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
package_dir = custom_nodes_dir / "example"
|
||||
package_dir.mkdir(parents=True)
|
||||
@@ -1437,7 +1453,7 @@ def test_modal_volume_backend_treats_missing_path_as_cache_miss(
|
||||
"""Return a fake volume that always reports missing paths."""
|
||||
return FakeVolume()
|
||||
|
||||
monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", FakeModal)
|
||||
backend = sync_engine_module.ModalVolumeBackend("volume")
|
||||
|
||||
assert backend.exists("/hashes/missing.done") is False
|
||||
@@ -1507,7 +1523,7 @@ def test_modal_volume_backend_caches_exists_results_and_uploaded_paths(
|
||||
"""Return the fake volume for any lookup."""
|
||||
return fake_volume
|
||||
|
||||
monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", FakeModal)
|
||||
backend = sync_engine_module.ModalVolumeBackend("volume")
|
||||
|
||||
assert backend.exists("/hashes/present.done") is False
|
||||
@@ -1583,7 +1599,7 @@ def test_modal_volume_backend_ignores_file_exists_upload_race(
|
||||
del name, create_if_missing
|
||||
return fake_volume
|
||||
|
||||
monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", FakeModal)
|
||||
backend = sync_engine_module.ModalVolumeBackend("volume")
|
||||
local_path = tmp_path / "bundle.zip"
|
||||
local_path.write_bytes(b"bundle")
|
||||
@@ -1654,9 +1670,9 @@ def test_modal_volume_backend_retries_rate_limited_calls(
|
||||
sleep_calls.append(seconds)
|
||||
monotonic_time += seconds
|
||||
|
||||
monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", FakeModal)
|
||||
monkeypatch.setattr(
|
||||
sync_engine_module,
|
||||
sync_engine_module._sync_backends,
|
||||
"time",
|
||||
types.SimpleNamespace(sleep=fake_sleep, monotonic=fake_monotonic),
|
||||
)
|
||||
@@ -1715,9 +1731,9 @@ def test_modal_volume_backend_applies_shared_rate_limit_backoff_across_calls(
|
||||
sleep_calls.append(seconds)
|
||||
monotonic_time += seconds
|
||||
|
||||
monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
|
||||
monkeypatch.setattr(sync_engine_module._sync_backends, "modal", FakeModal)
|
||||
monkeypatch.setattr(
|
||||
sync_engine_module,
|
||||
sync_engine_module._sync_backends,
|
||||
"time",
|
||||
types.SimpleNamespace(sleep=fake_sleep, monotonic=fake_monotonic),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user