diff --git a/backend/app/core/task_registry.py b/backend/app/core/task_registry.py index 27a182b..86176bb 100644 --- a/backend/app/core/task_registry.py +++ b/backend/app/core/task_registry.py @@ -22,6 +22,7 @@ def __init__(self) -> None: self._tasks: Dict[str, Dict[str, Any]] = {} self._cancellations: Dict[str, asyncio.Event] = {} self._update_locks: Dict[str, asyncio.Lock] = {} + self._install_locks: Dict[str, asyncio.Lock] = {} def register_task( self, @@ -83,5 +84,14 @@ def get_update_lock(self, engine_key: str) -> asyncio.Lock: self._update_locks[engine_key] = asyncio.Lock() return self._update_locks[engine_key] + def get_install_lock(self, engine_key: str) -> asyncio.Lock: + """Retrieve or create an exclusive async lease lock for an engine installation.""" + if not hasattr(self, "_install_locks"): + self._install_locks: Dict[str, asyncio.Lock] = {} + if engine_key not in self._install_locks: + self._install_locks[engine_key] = asyncio.Lock() + return self._install_locks[engine_key] + task_registry = TaskLifecycleRegistry() + diff --git a/backend/app/runtime/installer.py b/backend/app/runtime/installer.py index fe4e5f3..7d78d89 100644 --- a/backend/app/runtime/installer.py +++ b/backend/app/runtime/installer.py @@ -37,6 +37,9 @@ class IsolatedEngineInstaller: """Manages staged installation of local engines inside hermetic environments.""" + # Class-level set of currently active installation tasks (engine_type.value) + _active_installs: set[str] = set() + def __init__(self, engine_dir: Optional[Path] = None) -> None: self.engine_dir = engine_dir or get_default_engine_dir() @@ -55,14 +58,17 @@ def read_manifest(self, engine_type: EngineType) -> EngineInstallManifest: data = json.loads(manifest_path.read_text(encoding="utf-8")) manifest = EngineInstallManifest.model_validate(data) # Check for uncompleted previous run (interruption detection) + # Only mark as INTERRUPTED if there is NO active in-memory install worker running! if manifest.phase in ( + InstallPhase.CHECKING, InstallPhase.CREATING_VENV, InstallPhase.DOWNLOADING, InstallPhase.INSTALLING_DEPS, ): - manifest.phase = InstallPhase.INTERRUPTED - manifest.error_message = "Installation was interrupted before completion." - self.write_manifest(manifest) + if engine_type.value not in self._active_installs: + manifest.phase = InstallPhase.INTERRUPTED + manifest.error_message = "Installation was interrupted before completion." + self.write_manifest(manifest) return manifest except Exception as e: logger.warning(f"Error reading manifest for {engine_type}: {e}") @@ -91,6 +97,34 @@ def _get_pip_bin(self, runtime_dir: Path) -> Path: return runtime_dir / "Scripts" / "pip.exe" return runtime_dir / "bin" / "pip" + def _get_expected_entrypoints(self, engine_type: EngineType, engine_target: Path) -> list[Path]: + """Return expected entrypoint script paths for the engine.""" + if engine_type == EngineType.COMFYUI: + return [engine_target / "main.py"] + elif engine_type == EngineType.WEBUI: + return [engine_target / "launch.py", engine_target / "webui.py"] + return [] + + def validate_installation(self, engine_type: EngineType, engine_target: Path, runtime_target: Path) -> None: + """ + Validate that the engine installation is complete and healthy: + 1. Virtual environment Python binary exists + 2. Expected entrypoint script exists + 3. Engine requirements.txt exists + """ + python_bin = self._get_python_bin(runtime_target) + if not python_bin.is_file(): + raise RuntimeError(f"Hermetic Python binary missing at {python_bin}") + + entrypoints = self._get_expected_entrypoints(engine_type, engine_target) + if not any(ep.is_file() for ep in entrypoints): + expected_names = " or ".join(ep.name for ep in entrypoints) or "entrypoint" + raise RuntimeError(f"Engine entrypoint ({expected_names}) missing in {engine_target}") + + req_file = engine_target / "requirements.txt" + if not req_file.is_file(): + raise RuntimeError(f"Engine requirements.txt missing in {engine_target}") + async def create_isolated_venv( self, runtime_dir: Path, on_log: Optional[Callable[[str], None]] = None ) -> bool: @@ -120,68 +154,94 @@ async def install_engine( ) -> EngineInstallManifest: """ Run complete isolated installation: - 1. Check environment & disk space - 2. Create isolated sandboxed venv - 3. Clone / stage repository code + 1. Acquire per-engine install lock + 2. Check environment & create isolated sandboxed venv + 3. Clone / stage repository code into staging or target 4. Install isolated dependencies via sandboxed pip + 5. Validate entrypoint and health before marking completed """ - manifest = self.read_manifest(engine_type) - now_str = datetime.now(timezone.utc).isoformat() - manifest.created_at = now_str - manifest.phase = InstallPhase.CHECKING - self.write_manifest(manifest) - if on_progress: - on_progress(manifest) - - engine_target = Path(manifest.engine_dir) - runtime_target = Path(manifest.runtime_dir) - - try: - # Phase 1: Virtualenv Creation - manifest.phase = InstallPhase.CREATING_VENV - manifest.last_log_line = "Initializing sandboxed virtual environment..." - self.write_manifest(manifest) - if on_progress: - on_progress(manifest) - - venv_ok = await self.create_isolated_venv(runtime_target) - if not venv_ok: - raise RuntimeError(f"Failed to create virtual environment at {runtime_target}") - - python_bin = self._get_python_bin(runtime_target) - manifest.python_bin = str(python_bin) - - # Phase 2: Stage / Download code - manifest.phase = InstallPhase.DOWNLOADING - manifest.last_log_line = f"Staging {engine_type.value} repository..." - self.write_manifest(manifest) - if on_progress: - on_progress(manifest) - - repo_url = COMFYUI_GIT_REPO if engine_type == EngineType.COMFYUI else WEBUI_GIT_REPO - if not engine_target.exists(): - clone_cmd = ["git", "clone", "--depth", "1", repo_url, str(engine_target)] - proc = await asyncio.create_subprocess_exec( - *clone_cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await proc.communicate() - if proc.returncode != 0: - err_msg = stderr.decode(errors="replace").strip() - raise RuntimeError(f"Git clone failed: {err_msg}") + engine_key = engine_type.value + install_lock = task_registry.get_install_lock(engine_key) - # Phase 3: Install isolated dependencies - manifest.phase = InstallPhase.INSTALLING_DEPS - manifest.last_log_line = "Installing engine requirements into isolated environment..." + if install_lock.locked(): + manifest = self.read_manifest(engine_type) + manifest.phase = InstallPhase.FAILED + manifest.error_message = f"An installation lease is already held for {engine_type.value}." self.write_manifest(manifest) - if on_progress: - on_progress(manifest) + return manifest - req_file = engine_target / "requirements.txt" - pip_bin = self._get_pip_bin(runtime_target) + async with install_lock: + self._active_installs.add(engine_key) + try: + manifest = self.read_manifest(engine_type) + now_str = datetime.now(timezone.utc).isoformat() + manifest.created_at = now_str + manifest.phase = InstallPhase.CHECKING + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + engine_target = Path(manifest.engine_dir) + runtime_target = Path(manifest.runtime_dir) + + # Phase 1: Virtualenv Creation + manifest.phase = InstallPhase.CREATING_VENV + manifest.last_log_line = "Initializing sandboxed virtual environment..." + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + venv_ok = await self.create_isolated_venv(runtime_target) + if not venv_ok: + raise RuntimeError(f"Failed to create virtual environment at {runtime_target}") + + python_bin = self._get_python_bin(runtime_target) + manifest.python_bin = str(python_bin) + + # Phase 2: Stage / Download code + manifest.phase = InstallPhase.DOWNLOADING + manifest.last_log_line = f"Staging {engine_type.value} repository..." + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + repo_url = COMFYUI_GIT_REPO if engine_type == EngineType.COMFYUI else WEBUI_GIT_REPO + entrypoints = self._get_expected_entrypoints(engine_type, engine_target) + + # If destination directory exists but has no valid entrypoint or git repo, clean it up before cloning + if engine_target.exists(): + has_entrypoint = any(ep.is_file() for ep in entrypoints) + has_git = (engine_target / ".git").is_dir() + if not has_entrypoint or not has_git: + shutil.rmtree(engine_target, ignore_errors=True) + + if not engine_target.exists(): + clone_cmd = ["git", "clone", "--depth", "1", repo_url, str(engine_target)] + proc = await asyncio.create_subprocess_exec( + *clone_cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate() + if proc.returncode != 0: + err_msg = stderr.decode(errors="replace").strip() + raise RuntimeError(f"Git clone failed: {err_msg}") + + # Phase 3: Install isolated dependencies + manifest.phase = InstallPhase.INSTALLING_DEPS + manifest.last_log_line = "Installing engine requirements into isolated environment..." + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + req_file = engine_target / "requirements.txt" + pip_bin = self._get_pip_bin(runtime_target) + + if not req_file.is_file(): + raise RuntimeError(f"requirements.txt missing in {engine_target}") + if not pip_bin.is_file(): + raise RuntimeError(f"Hermetic pip binary missing in {runtime_target}") - if req_file.is_file() and pip_bin.is_file(): # Strict: execute ONLY pip inside runtime_target! install_cmd = [str(pip_bin), "install", "--no-warn-script-location", "-r", str(req_file)] proc = await asyncio.create_subprocess_exec( @@ -195,22 +255,27 @@ async def install_engine( err_msg = stderr.decode(errors="replace").strip() raise RuntimeError(f"Pip install failed: {err_msg[:300]}") - # Completed! - manifest.phase = InstallPhase.COMPLETED - manifest.completed_at = datetime.now(timezone.utc).isoformat() - manifest.last_log_line = f"{engine_type.value} installed successfully." - manifest.error_message = None - self.write_manifest(manifest) - if on_progress: - on_progress(manifest) + # Phase 4: Validation + self.validate_installation(engine_type, engine_target, runtime_target) - except Exception as e: - manifest.phase = InstallPhase.FAILED - manifest.error_message = str(e) - manifest.last_log_line = f"Installation error: {e}" - self.write_manifest(manifest) - if on_progress: - on_progress(manifest) + # Completed! + manifest.phase = InstallPhase.COMPLETED + manifest.completed_at = datetime.now(timezone.utc).isoformat() + manifest.last_log_line = f"{engine_type.value} installed successfully." + manifest.error_message = None + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + except Exception as e: + manifest.phase = InstallPhase.FAILED + manifest.error_message = str(e) + manifest.last_log_line = f"Installation error: {e}" + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + finally: + self._active_installs.discard(engine_key) return manifest diff --git a/backend/tests/test_m2_engines.py b/backend/tests/test_m2_engines.py index d27c5b6..4693745 100644 --- a/backend/tests/test_m2_engines.py +++ b/backend/tests/test_m2_engines.py @@ -20,6 +20,7 @@ from fastapi.testclient import TestClient from app.main import app +from app.core.task_registry import task_registry from app.runtime.engine_manager import EngineManager from app.runtime.hardware import check_hardware_readiness, detect_gpus, get_storage_readiness from app.runtime.installer import IsolatedEngineInstaller @@ -90,7 +91,7 @@ def test_installer_manifest_and_interruption(): assert manifest.phase == InstallPhase.IDLE assert manifest.engine_type == EngineType.COMFYUI - # 2. Simulate an interrupted install state + # 2. Simulate an interrupted install state (no active in-memory task) manifest.phase = InstallPhase.DOWNLOADING installer_inst.write_manifest(manifest) @@ -99,6 +100,59 @@ def test_installer_manifest_and_interruption(): assert reloaded.phase == InstallPhase.INTERRUPTED assert "interrupted" in (reloaded.error_message or "").lower() + # 4. Live polling while an install is actively running must NOT mark as interrupted + installer_inst._active_installs.add(EngineType.COMFYUI.value) + try: + manifest.phase = InstallPhase.DOWNLOADING + manifest.error_message = None + installer_inst.write_manifest(manifest) + live_manifest = installer_inst.read_manifest(EngineType.COMFYUI) + assert live_manifest.phase == InstallPhase.DOWNLOADING + assert live_manifest.error_message is None + finally: + installer_inst._active_installs.discard(EngineType.COMFYUI.value) + + +@pytest.mark.asyncio +async def test_installer_empty_dir_missing_entrypoint_rejection(tmp_path): + """Verify that an empty engine directory or missing entrypoint is rejected and does not report COMPLETED.""" + installer_inst = IsolatedEngineInstaller(engine_dir=tmp_path) + engine_target = tmp_path / "comfyui" + engine_target.mkdir(parents=True, exist_ok=True) + runtime_target = tmp_path / "runtime" + + # Mock venv creation as succeeding and creating python binary + python_bin = installer_inst._get_python_bin(runtime_target) + python_bin.parent.mkdir(parents=True, exist_ok=True) + python_bin.write_text("# mock python", encoding="utf-8") + pip_bin = installer_inst._get_pip_bin(runtime_target) + pip_bin.write_text("# mock pip", encoding="utf-8") + + with patch.object(installer_inst, "create_isolated_venv", return_value=True): + # Empty directory with no main.py and no .git should trigger cleanup & clone attempt, + # but mock git clone to do nothing, simulating missing entrypoint failure + with patch("asyncio.create_subprocess_exec") as mock_exec: + proc_mock = AsyncMock() + proc_mock.returncode = 0 + proc_mock.communicate.return_value = (b"", b"") + mock_exec.return_value = proc_mock + + manifest = await installer_inst.install_engine(EngineType.COMFYUI) + assert manifest.phase == InstallPhase.FAILED + assert "requirements.txt missing" in (manifest.error_message or "") or "entrypoint" in (manifest.error_message or "") + + +@pytest.mark.asyncio +async def test_installer_concurrent_lease_lock_protection(tmp_path): + """Verify that concurrent install calls for the same engine are rejected by lease lock.""" + installer_inst = IsolatedEngineInstaller(engine_dir=tmp_path) + lock = task_registry.get_install_lock(EngineType.COMFYUI.value) + + async with lock: + manifest = await installer_inst.install_engine(EngineType.COMFYUI) + assert manifest.phase == InstallPhase.FAILED + assert "lease is already held" in (manifest.error_message or "") + @pytest.mark.asyncio async def test_external_engine_connection_zero_ownership():