"""Tests for the sync custom nodes boundary.""" from __future__ import annotations from sync_engine_test_support import * # noqa: F401,F403 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._sync_backends, "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._sync_backends, "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._sync_backends, "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._sync_backends, "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._custom_nodes._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._custom_nodes, "_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_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._sync_backends, "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._custom_nodes, "_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._sync_backends, "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_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._sync_backends, "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._custom_nodes, "_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._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" 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._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.""" 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._custom_nodes, "_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._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" 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._custom_nodes._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._custom_nodes, "_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_custom_node_archive_uses_archive_byte_digest_for_object_identity( settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """R2 identity should describe ZIP bytes rather than the source-tree digest.""" custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) source_path = package_dir / "__init__.py" source_path.write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") engine = sync_engine_module.ModalAssetSyncEngine( volume=object(), settings=_r2_sync_settings(settings_module, tmp_path), ) source_digest = engine._hash_file_group(custom_nodes_dir, [source_path]) archive_spec = sync_engine_module._CustomNodesArchiveSpec( entry_name="example", display_name="example", source_description=str(package_dir), files=(source_path,), sha256=source_digest, ) captured_identity: list[tuple[Path, str, str]] = [] def sync_file(**kwargs: Any) -> Any: """Capture the exact local payload and identities selected for sync.""" captured_identity.append( (kwargs["local_path"], kwargs["sha256"], kwargs["sync_key"]) ) return sync_engine_module._ContentAddressedSyncResult( remote_path=kwargs["remote_path"], uploaded=True, ) monkeypatch.setattr(engine, "_sync_content_addressed_file", sync_file) result = engine._sync_custom_nodes_archive_spec(custom_nodes_dir, archive_spec) archive_path, payload_digest, sync_key = captured_identity[0] assert payload_digest == hashlib.sha256(archive_path.read_bytes()).hexdigest() assert payload_digest != source_digest assert sync_key.endswith(f"custom_nodes_entry:example:{source_digest}") assert result.sha256 == payload_digest def test_indexed_custom_node_manifest_backfills_archives_and_manifest_to_r2( settings_module: Any, sync_engine_module: Any, r2_cache_module: Any, tmp_path: Path, ) -> None: """Enabling R2 should revisit every object behind an indexed manifest.""" 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 = replace( _r2_sync_settings(settings_module, tmp_path), sync_custom_nodes=True, custom_nodes_dir=custom_nodes_dir, ) volume = _R2BackfillVolume(r2_cache_module) initial_engine = sync_engine_module.ModalAssetSyncEngine( volume=volume, settings=settings, ) initial_bundle = initial_engine.sync_custom_nodes_directory() assert initial_bundle is not None assert initial_bundle.uploaded is True backfill_engine = sync_engine_module.ModalAssetSyncEngine( volume=volume, settings=settings, r2_cache=_SynchronousBackfillCache(r2_cache_module), ) backfilled_bundle = backfill_engine.sync_custom_nodes_directory() backfill_engine.wait_for_r2_writebacks() assert backfilled_bundle is not None assert backfilled_bundle.uploaded is False assert len(volume.r2_uploads) == 2 uploaded_paths = [remote_path for _, remote_path in volume.r2_uploads] assert any(remote_path.endswith(".zip") for remote_path in uploaded_paths) assert any(remote_path.endswith(".json") for remote_path in uploaded_paths)