"""Tests for asset syncing and custom_nodes archiving.""" from __future__ import annotations import json from pathlib import Path import threading import time import types from typing import Any import zipfile def test_sync_file_deduplicates_by_hash( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Repeated file syncs should reuse the same remote path and sync-index record.""" monkeypatch.setattr(sync_engine_module, "modal", None) asset_path = tmp_path / "model.safetensors" asset_path.write_bytes(b"model-bytes") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=None, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) first = engine.sync_file(asset_path) second = engine.sync_file(asset_path) assert first.remote_path == second.remote_path assert first.sha256 == second.sha256 assert first.uploaded is True assert second.uploaded is False assert (settings.local_storage_root / first.remote_path.lstrip("/")).exists() sync_index_path = settings.local_storage_root / "metadata" / "sync_index.json" assert sync_index_path.exists() assert first.remote_path in sync_index_path.read_text(encoding="utf-8") def test_request_asset_cache_syncs_repeated_prompt_asset_once( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Repeated references across remote nodes should share one request sync decision.""" monkeypatch.setattr(sync_engine_module, "modal", None) asset_path = tmp_path / "model.safetensors" asset_path.write_bytes(b"model-bytes") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=None, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) sync_calls: list[Path] = [] original_sync_file = engine.sync_file def recording_sync_file(local_path: Path, *args: Any, **kwargs: Any) -> Any: """Record actual request sync decisions before delegating.""" sync_calls.append(local_path.resolve()) return original_sync_file(local_path, *args, **kwargs) monkeypatch.setattr(engine, "sync_file", recording_sync_file) request_cache = engine.create_request_asset_cache( ({"model": str(asset_path)}, {"second_model": str(asset_path)}) ) first_inputs, first_assets = engine.sync_prompt_inputs( {"model": str(asset_path)}, request_cache=request_cache, ) second_inputs, second_assets = engine.sync_prompt_inputs( {"second_model": str(asset_path)}, request_cache=request_cache, ) assert sync_calls == [asset_path.resolve()] assert first_inputs["model"] == second_inputs["second_model"] assert first_assets == second_assets assert request_cache.synced_assets() == (first_assets[0],) def test_sync_file_emits_upload_status( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Uploading a new asset should emit a status message naming the uploaded file.""" monkeypatch.setattr(sync_engine_module, "modal", None) asset_path = tmp_path / "model.safetensors" asset_path.write_bytes(b"model-bytes") observed_statuses: list[tuple[str, int | None, int | None]] = [] settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=None, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) engine.sync_file( asset_path, status_callback=lambda message, current, total: observed_statuses.append( (message, current, total) ), item_index=1, total_items=2, ) assert observed_statuses == [("Uploading asset 1/2 to Modal: model.safetensors", 1, 2)] def test_sync_custom_nodes_directory_creates_archive( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """The sync engine should archive and mirror a custom_nodes directory.""" monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) bundle = engine.sync_custom_nodes_directory() assert bundle is not None assert bundle.remote_path.startswith("/custom_nodes/") assert bundle.uploaded is True assert (settings.local_storage_root / bundle.remote_path.lstrip("/")).exists() def test_sync_custom_nodes_directory_only_checks_once_per_engine_lifetime( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, 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) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) first_bundle = engine.sync_custom_nodes_directory() assert first_bundle is not None assert first_bundle.uploaded is True (package_dir / "new_file.py").write_text("print('new node')\n", encoding="utf-8") def fail_hash_directory(path: Path) -> str: """Fail the test when a second custom_nodes scan happens.""" raise AssertionError(f"Did not expect custom_nodes to be rehashed: {path}") monkeypatch.setattr(engine, "_hash_directory", fail_hash_directory) second_bundle = engine.sync_custom_nodes_directory() assert second_bundle == sync_engine_module.SyncedAsset( local_path=first_bundle.local_path, remote_path=first_bundle.remote_path, sha256=first_bundle.sha256, uploaded=False, ) def test_sync_custom_nodes_directory_emits_packaging_and_upload_status( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Building a new custom_nodes archive should report packaging and upload stages.""" monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") observed_statuses: list[tuple[str, int | None, int | None]] = [] settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) engine.sync_custom_nodes_directory( status_callback=lambda message, current, total: observed_statuses.append( (message, current, total) ) ) assert observed_statuses == [ ("Packaging custom-node code for Modal", None, None), ("Uploading custom-node code and assets to Modal", None, None), ] def test_sync_custom_nodes_directory_concurrent_calls_share_one_initial_sync( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Concurrent first-run calls should serialize behind one custom_nodes sync decision.""" monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) original_sync = engine._sync_custom_nodes_directory_uncached sync_call_count = 0 sync_call_count_lock = threading.Lock() release_first_sync = threading.Event() def wrapped_sync(*, status_callback: Any = None) -> Any: """Count uncached sync executions and block the first one briefly.""" nonlocal sync_call_count with sync_call_count_lock: sync_call_count += 1 release_first_sync.wait(timeout=5.0) return original_sync(status_callback=status_callback) monkeypatch.setattr(engine, "_sync_custom_nodes_directory_uncached", wrapped_sync) results: list[Any] = [None, None] def run_sync(index: int) -> None: """Run one custom_nodes sync call and store the result.""" results[index] = engine.sync_custom_nodes_directory() first_thread = threading.Thread(target=run_sync, args=(0,)) second_thread = threading.Thread(target=run_sync, args=(1,)) first_thread.start() time.sleep(0.1) second_thread.start() release_first_sync.set() first_thread.join(timeout=5.0) second_thread.join(timeout=5.0) assert not first_thread.is_alive() assert not second_thread.is_alive() assert sync_call_count == 1 assert results[0] is not None assert results[1] is not None assert results[0].remote_path == results[1].remote_path assert results[0].sha256 == results[1].sha256 assert results[0].uploaded is True assert results[1].uploaded is False def test_hash_directory_ignores_virtualenv_and_bytecode_artifacts( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Ignored directories and file suffixes should not affect custom_nodes hashing.""" monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) baseline_hash = engine._hash_directory(custom_nodes_dir) ignored_venv_dir = custom_nodes_dir / ".venv" / "lib" ignored_venv_dir.mkdir(parents=True) (ignored_venv_dir / "ignored.py").write_text("print('ignored')\n", encoding="utf-8") ignored_cache_dir = package_dir / "__pycache__" ignored_cache_dir.mkdir() (ignored_cache_dir / "example.cpython-312.pyc").write_bytes(b"bytecode") (package_dir / "ignored.pyc").write_bytes(b"bytecode") (package_dir / "ignored.log").write_text("log noise\n", encoding="utf-8") ignored_hash = engine._hash_directory(custom_nodes_dir) assert ignored_hash == baseline_hash def test_sync_file_reuses_cached_hash_for_unchanged_file( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, 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) asset_path = tmp_path / "model.safetensors" asset_path.write_bytes(b"model-bytes") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=None, ) first_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) first_sha = first_engine._hash_file(asset_path) second_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) def fail_open(*args: Any, **kwargs: Any) -> Any: raise AssertionError("Expected cached hash lookup to avoid reopening the unchanged file.") monkeypatch.setattr(Path, "open", fail_open) second_sha = second_engine._hash_file(asset_path) assert second_sha == first_sha def test_sync_custom_nodes_directory_reuses_cached_archive( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """An unchanged custom_nodes tree should reuse its digest-keyed local archive.""" monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) first_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) first_bundle = first_engine.sync_custom_nodes_directory() assert first_bundle is not None class NeverExistsVolume: """Volume double that forces archive reuse without any remote metadata probes.""" def exists(self, remote_path: str) -> bool: return False def put_file(self, local_path: Path, remote_path: str) -> None: return None def put_bytes(self, payload: bytes, remote_path: str) -> None: return None second_engine = sync_engine_module.ModalAssetSyncEngine( volume=NeverExistsVolume(), settings=settings, ) 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) second_bundle = second_engine.sync_custom_nodes_directory() assert second_bundle is not None assert second_bundle.sha256 == first_bundle.sha256 assert second_bundle.uploaded is False entry_hash = second_engine._custom_nodes_archive_specs(custom_nodes_dir)[0].sha256 assert second_engine._cached_custom_nodes_archive_path("example", entry_hash).exists() def test_sync_custom_nodes_separates_code_assets_and_nested_virtualenv( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Code ZIPs should exclude mounted model assets and nested virtual environments.""" monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") model_path = package_dir / "checkpoints" / "model.pth" model_path.parent.mkdir() model_path.write_bytes(b"package-model-weights") venv_path = package_dir / ".venv" / "lib" / "site-packages" / "ignored.py" venv_path.parent.mkdir(parents=True) venv_path.write_text("raise RuntimeError('must not ship')\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) bundle = engine.sync_custom_nodes_directory() assert bundle is not None manifest_path = settings.local_storage_root / bundle.remote_path.lstrip("/") manifest = json.loads(manifest_path.read_text(encoding="utf-8")) assert manifest["version"] == 2 [entry] = manifest["entries"] [asset] = entry["assets"] assert asset["relative_path"] == "example/checkpoints/model.pth" assert asset["size_bytes"] == len(b"package-model-weights") asset_path = settings.local_storage_root / asset["remote_path"].lstrip("/") assert asset_path.read_bytes() == b"package-model-weights" archive_path = settings.local_storage_root / entry["remote_path"].lstrip("/") with zipfile.ZipFile(archive_path, "r") as archive: archived_names = archive.namelist() assert "example/__init__.py" in archived_names assert "example/checkpoints/model.pth" not in archived_names assert not any(".venv" in name for name in archived_names) def test_sync_file_reuses_sync_index_record_for_existing_remote_payload( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """An indexed deterministic asset should be reused without a second upload.""" monkeypatch.setattr(sync_engine_module, "modal", None) asset_path = tmp_path / "model.safetensors" asset_path.write_bytes(b"model-bytes") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=None, ) class RecordingVolume: """Volume double that records any attempted uploads.""" def __init__(self) -> None: """Initialize captured writes.""" self.put_file_calls: list[tuple[Path, str]] = [] def put_file(self, local_path: Path, remote_path: str) -> None: self.put_file_calls.append((local_path, remote_path)) def put_bytes(self, payload: bytes, remote_path: str) -> None: raise AssertionError("Sync records should not be mirrored as marker files.") def exists(self, remote_path: str) -> bool: raise AssertionError("Sync should not probe volume metadata for indexed payloads.") volume = RecordingVolume() engine = sync_engine_module.ModalAssetSyncEngine(volume=volume, settings=settings) engine.sync_index.put( engine._asset_sync_index_key(engine._hash_file(asset_path)), {"remote_path": "/assets/existing_model.safetensors", "source": "existing"}, ) synced_asset = engine.sync_file(asset_path) assert synced_asset.uploaded is False assert synced_asset.remote_path == "/assets/existing_model.safetensors" assert volume.put_file_calls == [] def test_sync_custom_nodes_directory_reuses_indexed_remote_bundle( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, 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) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) directory_hash = engine._hash_directory(custom_nodes_dir) remote_path = engine._custom_nodes_manifest_remote_path(directory_hash) class RecordingVolume: """Volume double that records any attempted uploads.""" def __init__(self) -> None: """Initialize captured writes.""" self.put_file_calls: list[tuple[Path, str]] = [] def put_file(self, local_path: Path, candidate_path: str) -> None: self.put_file_calls.append((local_path, candidate_path)) def put_bytes(self, payload: bytes, candidate_path: str) -> None: raise AssertionError("Sync records should not be mirrored as marker files.") def exists(self, candidate_path: str) -> bool: raise AssertionError("Sync should not probe volume metadata for indexed bundles.") volume = RecordingVolume() engine = sync_engine_module.ModalAssetSyncEngine(volume=volume, settings=settings) engine.sync_index.put( engine._custom_nodes_manifest_sync_index_key(directory_hash), {"remote_path": remote_path, "source": str(custom_nodes_dir)}, ) 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) bundle = engine.sync_custom_nodes_directory() assert bundle is not None assert bundle.sha256 == directory_hash assert bundle.remote_path == remote_path assert bundle.uploaded is False assert volume.put_file_calls == [] def test_sync_custom_nodes_directory_only_rebuilds_changed_top_level_archive( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, 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) custom_nodes_dir = tmp_path / "custom_nodes" package_a = custom_nodes_dir / "example_a" package_b = custom_nodes_dir / "example_b" package_a.mkdir(parents=True) package_b.mkdir(parents=True) (package_a / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") (package_b / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) first_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) first_bundle = first_engine.sync_custom_nodes_directory() assert first_bundle is not None (package_b / "node.py").write_text("VALUE = 2\n", encoding="utf-8") class NeverExistsVolume: """Volume double that forces all deterministic uploads through without indexed cache hits.""" def exists(self, remote_path: str) -> bool: return False def put_file(self, local_path: Path, remote_path: str) -> None: return None def put_bytes(self, payload: bytes, remote_path: str) -> None: return None second_engine = sync_engine_module.ModalAssetSyncEngine( volume=NeverExistsVolume(), settings=settings, ) rebuilt_entries: list[str] = [] original_create_archive = second_engine._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.""" del root_path 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) second_bundle = second_engine.sync_custom_nodes_directory() assert second_bundle is not None assert second_bundle.sha256 != first_bundle.sha256 assert rebuilt_entries == ["example_b"] def test_sync_custom_nodes_directory_builds_multiple_archives_in_parallel( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Fresh per-package custom_nodes archives should build in parallel.""" monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_a = custom_nodes_dir / "example_a" package_b = custom_nodes_dir / "example_b" package_a.mkdir(parents=True) package_b.mkdir(parents=True) (package_a / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") (package_b / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) original_create_archive = engine._create_archive_from_files thread_ids: set[int] = set() started_count = 0 started_lock = threading.Lock() overlap_event = threading.Event() def record_create_archive(root_path: Path, files: list[Path], archive_path: Path) -> Path: """Block briefly until two archive builds overlap so the test can observe parallel execution.""" nonlocal started_count del root_path with started_lock: started_count += 1 thread_ids.add(threading.get_ident()) if started_count >= 2: overlap_event.set() assert overlap_event.wait(0.2), "Expected per-package archive builds to overlap." time.sleep(0.02) return original_create_archive(custom_nodes_dir, files, archive_path) monkeypatch.setattr(engine, "_create_archive_from_files", record_create_archive) bundle = engine.sync_custom_nodes_directory() assert bundle is not None assert len(thread_ids) >= 2 def test_remote_mode_uses_modal_volume_backend_when_sdk_is_available( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Remote execution mode should upload into a real Modal volume when the SDK is available.""" class FakeBatch: """Capture files uploaded through Modal batch_upload.""" def __init__(self, uploads: list[tuple[Any, str]]) -> None: """Store uploaded file references.""" self.uploads = uploads def __enter__(self) -> "FakeBatch": """Return the active batch context.""" return self def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: """Close the fake batch context.""" return None def put_file(self, source: Any, remote_path: str) -> None: """Record the uploaded file or bytes payload.""" self.uploads.append((source, remote_path)) class FakeVolume: """Minimal Modal volume double.""" def __init__(self) -> None: """Initialize the fake upload state.""" self.paths: set[str] = set() self.uploads: list[tuple[Any, str]] = [] def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: """Return a non-empty listing when the path has already been uploaded.""" return [remote_path] if remote_path in self.paths else [] def batch_upload(self) -> FakeBatch: """Return a fake batch uploader.""" return FakeBatch(self.uploads) fake_volume = FakeVolume() class FakeModal: """Minimal modal SDK double that returns a stable volume handle.""" class Dict: """Namespace for sync-index lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> dict[str, Any]: """Return a plain dict-backed sync index.""" del name, create_if_missing return {} class Volume: """Namespace for volume lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: """Return the fake volume for any lookup.""" return fake_volume monkeypatch.setattr(sync_engine_module, "modal", FakeModal) settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=None, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) assert type(engine.volume).__name__ == "ModalVolumeBackend" assert type(engine.sync_index).__name__ == "ModalDictSyncIndex" asset_path = tmp_path / "encoder.safetensors" asset_path.write_bytes(b"weights") synced = engine.sync_file(asset_path) uploaded_remote_paths = [remote_path for _, remote_path in fake_volume.uploads] assert synced.remote_path in uploaded_remote_paths assert all(not remote_path.startswith("/hashes/") for remote_path in uploaded_remote_paths) def test_remote_sync_index_discards_stale_volume_epoch_and_reuploads_missing_payloads( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """A fresh volume should ignore stale sync-index records from an older volume epoch.""" class FakeBatch: """Capture files uploaded through Modal batch_upload.""" def __init__(self, volume: "FakeVolume") -> None: """Store the backing volume for uploaded-path tracking.""" self.volume = volume def __enter__(self) -> "FakeBatch": """Return the active batch context.""" return self def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: """Close the fake batch context.""" return None def put_file(self, source: Any, remote_path: str) -> None: """Record the uploaded remote path as present in the fake volume.""" del source self.volume.paths.add(remote_path) self.volume.uploads.append(remote_path) class FakeVolume: """Minimal Modal volume double with listdir accounting.""" def __init__(self) -> None: """Initialize fake storage state.""" self.paths: set[str] = set() self.uploads: list[str] = [] def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: """Return a listing for known paths.""" del recursive return [remote_path] if remote_path in self.paths else [] def batch_upload(self) -> FakeBatch: """Return a fake uploader.""" return FakeBatch(self) fake_volume = FakeVolume() fake_dict: dict[str, Any] = {} class FakeModal: """Minimal Modal SDK double exposing both Volume and Dict backends.""" exception = type( "FakeExceptionNamespace", (), { "NotFoundError": FileNotFoundError, "ResourceExhaustedError": RuntimeError, }, ) class Dict: """Namespace for sync-index lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> dict[str, Any]: """Return the shared fake dict store for any lookup.""" del name, create_if_missing return fake_dict class Volume: """Namespace for volume lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: """Return the fake volume for any lookup.""" del name, create_if_missing return fake_volume monkeypatch.setattr(sync_engine_module, "modal", FakeModal) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=custom_nodes_dir, ) engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) entry_hash = engine._hash_directory(package_dir) fake_dict[f"{settings.volume_name}:current_volume_epoch"] = { "epoch": "stale-epoch", "remote_path": "/sync_index_epochs/stale-epoch.json", "sentinel_path": "/sync_index_epochs/stale-epoch.json", } fake_dict[ f"{settings.volume_name}:epoch:stale-epoch:custom_nodes_entry:example:{entry_hash}" ] = { "remote_path": engine._custom_nodes_archive_remote_path("example", entry_hash), "source": str(package_dir), } bundle = engine.sync_custom_nodes_directory() assert bundle is not None assert bundle.uploaded is True assert any(path.startswith("/sync_index_epochs/") for path in fake_volume.uploads) assert any(path.startswith("/custom_nodes/entries/example/") for path in fake_volume.uploads) assert fake_dict[f"{settings.volume_name}:current_volume_epoch"]["epoch"] != "stale-epoch" def test_modal_volume_backend_treats_missing_path_as_cache_miss( sync_engine_module: Any, monkeypatch: Any, ) -> None: """A missing Modal volume path should behave like a normal absent marker file.""" class FakeNotFoundError(Exception): """Stand-in for modal.exception.NotFoundError.""" class FakeVolume: """Minimal Modal volume double that raises on missing listdir.""" def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: """Simulate Modal's missing-path behavior.""" raise FakeNotFoundError(remote_path) class FakeModal: """Minimal modal SDK double exposing the exception namespace.""" exception = type("FakeExceptionNamespace", (), {"NotFoundError": FakeNotFoundError}) class Volume: """Namespace for volume lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: """Return a fake volume that always reports missing paths.""" return FakeVolume() monkeypatch.setattr(sync_engine_module, "modal", FakeModal) backend = sync_engine_module.ModalVolumeBackend("volume") assert backend.exists("/hashes/missing.done") is False def test_modal_volume_backend_caches_exists_results_and_uploaded_paths( sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Repeated existence checks for the same path should avoid repeated Modal metadata calls.""" class FakeBatch: """Capture one uploaded path.""" def __init__(self, volume: "FakeVolume") -> None: """Store the backing volume.""" self.volume = volume def __enter__(self) -> "FakeBatch": """Return the active batch context.""" return self def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: """Close the fake batch context.""" return None def put_file(self, source: Any, remote_path: str) -> None: """Record the uploaded path as present.""" self.volume.paths.add(remote_path) class FakeVolume: """Minimal Modal volume double with listdir accounting.""" def __init__(self) -> None: """Initialize fake storage state.""" self.paths: set[str] = set() self.listdir_calls = 0 def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: """Return a listing for known paths while counting calls.""" self.listdir_calls += 1 return [remote_path] if remote_path in self.paths else [] def batch_upload(self) -> FakeBatch: """Return a fake uploader.""" return FakeBatch(self) fake_volume = FakeVolume() class FakeModal: """Minimal modal SDK double exposing a fake volume.""" exception = type( "FakeExceptionNamespace", (), { "NotFoundError": FileNotFoundError, }, ) class Volume: """Namespace for volume lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: """Return the fake volume for any lookup.""" return fake_volume monkeypatch.setattr(sync_engine_module, "modal", FakeModal) backend = sync_engine_module.ModalVolumeBackend("volume") assert backend.exists("/hashes/present.done") is False assert backend.exists("/hashes/present.done") is False assert fake_volume.listdir_calls == 1 local_path = tmp_path / "marker.txt" local_path.write_text("marker", encoding="utf-8") backend.put_file(local_path, "/hashes/present.done") assert backend.exists("/hashes/present.done") is True assert fake_volume.listdir_calls == 1 def test_modal_volume_backend_ignores_file_exists_upload_race( sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Content-addressed Modal uploads should be idempotent when the path already exists remotely.""" class FakeBatch: """Raise the same collision Modal raises when closing a duplicate upload batch.""" def __enter__(self) -> "FakeBatch": """Return the active batch context.""" return self def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: """Surface the duplicate remote file at batch close.""" raise FileExistsError("/custom_nodes/entries/package/hash_custom_nodes_bundle.zip") def put_file(self, source: Any, remote_path: str) -> None: """Accept the upload request before the batch close detects the duplicate.""" del source, remote_path class FakeVolume: """Minimal Modal volume double with duplicate-upload behavior.""" def __init__(self) -> None: """Initialize fake storage state.""" self.listdir_calls = 0 def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: """Return a listing for known paths while counting calls.""" del recursive self.listdir_calls += 1 return [remote_path] def batch_upload(self) -> FakeBatch: """Return a fake uploader.""" return FakeBatch() fake_volume = FakeVolume() class FakeModal: """Minimal modal SDK double exposing a duplicate-prone fake volume.""" exception = type( "FakeExceptionNamespace", (), { "NotFoundError": FileNotFoundError, }, ) class Volume: """Namespace for volume lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: """Return the fake volume for any lookup.""" del name, create_if_missing return fake_volume monkeypatch.setattr(sync_engine_module, "modal", FakeModal) backend = sync_engine_module.ModalVolumeBackend("volume") local_path = tmp_path / "bundle.zip" local_path.write_bytes(b"bundle") backend.put_file(local_path, "/custom_nodes/entries/package/hash_custom_nodes_bundle.zip") backend.put_bytes(b"{}", "/custom_nodes/manifests/hash_manifest.json") assert backend.exists("/custom_nodes/entries/package/hash_custom_nodes_bundle.zip") is True assert backend.exists("/custom_nodes/manifests/hash_manifest.json") is True assert fake_volume.listdir_calls == 0 def test_modal_volume_backend_retries_rate_limited_calls( sync_engine_module: Any, monkeypatch: Any, ) -> None: """Transient Modal rate limiting should back off and retry instead of failing immediately.""" class FakeResourceExhaustedError(Exception): """Stand-in for modal.exception.ResourceExhaustedError.""" class FakeVolume: """Minimal Modal volume double that rate limits once.""" def __init__(self) -> None: """Initialize the listdir attempt counter.""" self.listdir_calls = 0 def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: """Raise once, then succeed.""" self.listdir_calls += 1 if self.listdir_calls == 1: raise FakeResourceExhaustedError("rate limited") return [remote_path] fake_volume = FakeVolume() class FakeModal: """Minimal modal SDK double exposing retryable error types.""" exception = type( "FakeExceptionNamespace", (), { "NotFoundError": FileNotFoundError, "ResourceExhaustedError": FakeResourceExhaustedError, }, ) class Volume: """Namespace for volume lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: """Return the fake volume for any lookup.""" return fake_volume monotonic_time = 100.0 sleep_calls: list[float] = [] def fake_monotonic() -> float: """Return the controllable monotonic clock.""" return monotonic_time def fake_sleep(seconds: float) -> None: """Advance the controllable monotonic clock instead of actually sleeping.""" nonlocal monotonic_time sleep_calls.append(seconds) monotonic_time += seconds monkeypatch.setattr(sync_engine_module, "modal", FakeModal) monkeypatch.setattr( sync_engine_module, "time", types.SimpleNamespace(sleep=fake_sleep, monotonic=fake_monotonic), ) backend = sync_engine_module.ModalVolumeBackend("volume") assert backend.exists("/hashes/present.done") is True assert fake_volume.listdir_calls == 2 assert 0.25 in sleep_calls def test_modal_volume_backend_applies_shared_rate_limit_backoff_across_calls( sync_engine_module: Any, monkeypatch: Any, ) -> None: """One rate-limited Modal volume call should publish a shared backoff window for later calls.""" class FakeVolume: """Minimal Modal volume double that always succeeds.""" def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: """Return one successful listing.""" return [remote_path] fake_volume = FakeVolume() class FakeModal: """Minimal modal SDK double exposing retryable error types.""" exception = type( "FakeExceptionNamespace", (), { "NotFoundError": FileNotFoundError, "ResourceExhaustedError": RuntimeError, }, ) class Volume: """Namespace for volume lookups.""" @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: """Return the fake volume for any lookup.""" return fake_volume monotonic_time = 100.0 sleep_calls: list[float] = [] def fake_monotonic() -> float: """Return the controllable monotonic clock.""" return monotonic_time def fake_sleep(seconds: float) -> None: """Advance the controllable monotonic clock instead of actually sleeping.""" nonlocal monotonic_time sleep_calls.append(seconds) monotonic_time += seconds monkeypatch.setattr(sync_engine_module, "modal", FakeModal) monkeypatch.setattr( sync_engine_module, "time", types.SimpleNamespace(sleep=fake_sleep, monotonic=fake_monotonic), ) backend = sync_engine_module.ModalVolumeBackend("volume") assert backend._record_shared_rate_limit_backoff() == 0.25 assert backend.exists("/hashes/second.done") is True assert 0.25 in sleep_calls