fix(omnidreams): robust native build locking + memory-aware nvcc fanout

integrations/omnidreams/omnidreams/native/omnidreams_singleview.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import contextlib
+import errno
 import hashlib
 import importlib.util
 import json
@@ -28,6 +29,13 @@
 from types import ModuleType
 from typing import Any, Iterator
 
+try:  # POSIX only; absent on Windows.
+    import fcntl
+except ImportError:  # pragma: no cover - non-POSIX platforms
+    fcntl = None  # type: ignore[assignment]
+
+from loguru import logger
+
 from omnidreams.native.acceleration import (
     NativeAccelerationConfig,
     NativeAvailabilityCheck,
@@ -50,6 +58,17 @@
 _VAE_STREAMING_DIR = _SOURCE_DIR / "vae_streaming"
 _PYTORCH_MAX_JOBS_ENV = "MAX_JOBS"
 _DEFAULT_MAX_JOBS_CAP = 8
+# Heavy CUTLASS/SageAttention translation units peak at several GB of host RAM
+# each under -O3. On memory-constrained hosts, running one nvcc per core can
+# exhaust RAM and swap and lock up the machine, so we also cap concurrency by
+# available memory. Tunable via the env var below.
+_NATIVE_MEM_PER_JOB_ENV = "OMNIDREAMS_SINGLEVIEW_NATIVE_MEM_PER_JOB_GB"
+_DEFAULT_MEM_PER_JOB_GB = 8.0
+# PyTorch JIT (FileBaton) and Ninja lock files left behind by an interrupted
+# build (Ctrl-C, OOM, crash). When we hold the cross-process build guard these
+# are definitively stale and are reclaimed.
+_STALE_LOCK_NAMES = ("lock", ".ninja_lock")
+_BUILD_GUARD_NAME = ".build.flock"
 _NATIVE_CUDA_ARCH_LIST_ENV = "OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST"
 _PYTORCH_CUDA_ARCH_LIST_ENV = "TORCH_CUDA_ARCH_LIST"
 _DEFAULT_CUDA_ARCH_LIST = "12.0a"
@@ -225,12 +244,45 @@ def _validate_max_jobs(value: int | str) -> str:
     return str(jobs)
 
 
+def _total_ram_gb() -> float | None:
+    """Return total physical RAM in GiB, or ``None`` if it cannot be queried."""
+
+    try:
+        return (os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) / (1024**3)
+    except (AttributeError, ValueError, OSError):  # non-POSIX or unsupported key
+        return None
+
+
+def _mem_per_job_gb() -> float:
+    raw = os.environ.get(_NATIVE_MEM_PER_JOB_ENV)
+    if not raw:
+        return _DEFAULT_MEM_PER_JOB_GB
+    try:
+        value = float(raw)
+    except ValueError:
+        return _DEFAULT_MEM_PER_JOB_GB
+    return value if value > 0 else _DEFAULT_MEM_PER_JOB_GB
+
+
+def _memory_job_cap() -> int | None:
+    """Max parallel compile jobs the host RAM can absorb, or ``None`` if unknown."""
+
+    ram_gb = _total_ram_gb()
+    if ram_gb is None:
+        return None
+    return max(1, int(ram_gb // _mem_per_job_gb()))
+
+
 def _resolved_max_jobs(max_jobs: int | str | None) -> str | None:
     if max_jobs is not None:
         return _validate_max_jobs(max_jobs)
     if os.environ.get(_PYTORCH_MAX_JOBS_ENV):
         return None
-    return str(min(os.cpu_count() or 1, _DEFAULT_MAX_JOBS_CAP))
+    cap = min(os.cpu_count() or 1, _DEFAULT_MAX_JOBS_CAP)
+    memory_cap = _memory_job_cap()
+    if memory_cap is not None:
+        cap = min(cap, memory_cap)
+    return str(max(1, cap))
 
 
 def _resolved_cuda_arch_list() -> str | None:
@@ -292,6 +344,76 @@ def _scoped_cuda_arch_list() -> Iterator[None]:
             os.environ[_PYTORCH_CUDA_ARCH_LIST_ENV] = previous
 
 
+class NativeBuildBusyError(RuntimeError):
+    """Raised when another live process is already compiling the extension."""
+
+
+def _reclaim_stale_locks(build_dir: Path) -> None:
+    """Delete leftover JIT/Ninja lock files in ``build_dir``.
+
+    Only safe to call while holding the build guard: that guarantees no other
+    live process is mid-build for this directory, so any ``lock``/``.ninja_lock``
+    is an orphan from an interrupted run (Ctrl-C, OOM, crash) rather than an
+    active baton. Without this, PyTorch's ``FileBaton`` would poll the orphaned
+    lock forever and the build would stall indefinitely.
+    """
+
+    for name in _STALE_LOCK_NAMES:
+        lock_path = build_dir / name
+        try:
+            lock_path.unlink()
+        except FileNotFoundError:
+            continue
+        except OSError as exc:  # pragma: no cover - unexpected fs error
+            logger.warning("Could not remove stale native build lock {}: {}", lock_path, exc)
+            continue
+        logger.warning(
+            "Reclaimed stale native build lock {} (orphaned by an interrupted build)",
+            lock_path,
+        )
+
+
+@contextlib.contextmanager
+def _native_build_guard(build_dir: Path) -> Iterator[None]:
+    """Serialize native builds across processes and reclaim stale locks.
+
+    Uses a non-blocking ``flock`` on a sentinel file. The kernel drops an
+    ``flock`` automatically when its holder dies, so a build killed by Ctrl-C or
+    a crash never wedges the next run. If the guard is already held by a *live*
+    process we fail fast with :class:`NativeBuildBusyError` instead of blocking.
+
+    On non-POSIX platforms (no ``fcntl``) this is a no-op and the build relies on
+    PyTorch's own locking, preserving prior behavior.
+    """
+
+    if fcntl is None:  # pragma: no cover - non-POSIX platforms
+        yield
+        return
+
+    guard_path = build_dir / _BUILD_GUARD_NAME
+    guard = open(guard_path, "w")  # noqa: SIM115 - lifetime managed below
+    try:
+        try:
+            fcntl.flock(guard.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except OSError as exc:
+            if exc.errno in (errno.EWOULDBLOCK, errno.EAGAIN):
+                message = (
+                    "Another process is compiling the OmniDreams single-view "
+                    "native extension. Wait for it to finish, then relaunch."
+                )
+                logger.warning(message)
+                raise NativeBuildBusyError(message) from exc
+            raise
+        # We hold the exclusive guard: any leftover torch/ninja lock is stale.
+        _reclaim_stale_locks(build_dir)
+        yield
+    finally:
+        try:
+            fcntl.flock(guard.fileno(), fcntl.LOCK_UN)
+        finally:
+            guard.close()
+
+
 def load_extension(
     build_root: Path | str | None = None,
     *,
@@ -354,7 +476,11 @@ def load_extension(
             )
             extension_build_dir.mkdir(parents=True, exist_ok=True)
 
-            with _scoped_torch_max_jobs(max_jobs), _scoped_cuda_arch_list():
+            with (
+                _native_build_guard(extension_build_dir),
+                _scoped_torch_max_jobs(max_jobs),
+                _scoped_cuda_arch_list(),
+            ):
                 _extension = load_torch_extension(
                     name=extension_name,
                     sources=[str(source) for source in _extension_sources()],

integrations/omnidreams/tests/test_omnidreams_singleview_native.py

@@ -152,8 +152,12 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType:
     monkeypatch.setattr(native, "validate_thirdparty", lambda: thirdparty_info)
     monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension)
     monkeypatch.setattr(native.os, "cpu_count", lambda: 48)
+    # Pin RAM well above the CPU/hard cap so the memory throttle does not bind
+    # and the assertion below is independent of the CI runner's memory.
+    monkeypatch.setattr(native, "_total_ram_gb", lambda: 1024.0)
     monkeypatch.setattr(native, "_python_package_dir", lambda package: None)
     monkeypatch.delenv("MAX_JOBS", raising=False)
+    monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_NATIVE_MEM_PER_JOB_GB", raising=False)
     monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False)
     monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False)
 
@@ -344,6 +348,118 @@ def fake_load_torch_extension(**_: object) -> ModuleType:
     assert native.extension_load_error() is None
 
 
+@pytest.mark.ci_cpu
+def test_resolved_max_jobs_throttles_by_memory(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("MAX_JOBS", raising=False)
+    monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_NATIVE_MEM_PER_JOB_GB", raising=False)
+    monkeypatch.setattr(native.os, "cpu_count", lambda: 48)
+    # 16 GiB / 8 GiB-per-job -> 2 jobs, below the CPU and hard caps.
+    monkeypatch.setattr(native, "_total_ram_gb", lambda: 16.0)
+
+    assert native._resolved_max_jobs(None) == "2"
+
+
+@pytest.mark.ci_cpu
+def test_resolved_max_jobs_memory_env_override(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("MAX_JOBS", raising=False)
+    monkeypatch.setattr(native.os, "cpu_count", lambda: 48)
+    monkeypatch.setattr(native, "_total_ram_gb", lambda: 64.0)
+    monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_NATIVE_MEM_PER_JOB_GB", "32")
+
+    # 64 GiB / 32 GiB-per-job -> 2 jobs.
+    assert native._resolved_max_jobs(None) == "2"
+
+
+@pytest.mark.ci_cpu
+def test_resolved_max_jobs_unknown_memory_uses_cpu_cap(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("MAX_JOBS", raising=False)
+    monkeypatch.setattr(native.os, "cpu_count", lambda: 4)
+    monkeypatch.setattr(native, "_total_ram_gb", lambda: None)
+
+    assert native._resolved_max_jobs(None) == "4"
+
+
+@pytest.mark.ci_cpu
+def test_resolved_max_jobs_explicit_value_bypasses_memory(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(native, "_total_ram_gb", lambda: 8.0)
+
+    assert native._resolved_max_jobs(6) == "6"
+
+
+@pytest.mark.ci_cpu
+def test_native_build_guard_reclaims_stale_locks(tmp_path: Path) -> None:
+    if native.fcntl is None:
+        pytest.skip("requires POSIX fcntl")
+
+    build_dir = tmp_path / "torch_extensions" / "ext"
+    build_dir.mkdir(parents=True)
+    (build_dir / "lock").write_text("")
+    (build_dir / ".ninja_lock").write_text("")
+
+    with native._native_build_guard(build_dir):
+        # While holding the guard, the orphaned locks are reclaimed.
+        assert not (build_dir / "lock").exists()
+        assert not (build_dir / ".ninja_lock").exists()
+
+
+@pytest.mark.ci_cpu
+def test_native_build_guard_fails_fast_when_held(tmp_path: Path) -> None:
+    if native.fcntl is None:
+        pytest.skip("requires POSIX fcntl")
+
+    build_dir = tmp_path / "torch_extensions" / "ext"
+    build_dir.mkdir(parents=True)
+
+    with native._native_build_guard(build_dir):
+        with pytest.raises(native.NativeBuildBusyError):
+            with native._native_build_guard(build_dir):
+                pass
+
+
+@pytest.mark.ci_cpu
+def test_load_extension_reclaims_stale_lock_before_building(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    import torch.utils.cpp_extension as cpp_extension
+
+    build_root = tmp_path / "native-build"
+    thirdparty_info = _fake_thirdparty_info(tmp_path)
+    seen: dict[str, bool] = {}
+
+    def fake_load_torch_extension(**kwargs: object) -> ModuleType:
+        build_dir = Path(str(kwargs["build_directory"]))
+        # By the time torch's builder runs, the stale lock has been cleared, so
+        # its FileBaton acquires immediately instead of polling forever.
+        seen["lock_present"] = (build_dir / "lock").exists()
+        return _fake_extension_module()
+
+    monkeypatch.setattr(native, "_extension", None)
+    monkeypatch.setattr(native, "_extension_load_error", None)
+    monkeypatch.setattr(native, "validate_thirdparty", lambda: thirdparty_info)
+    monkeypatch.setattr(native, "_python_package_dir", lambda package: None)
+    monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension)
+
+    extension_name = native._extension_name(thirdparty_info)
+    build_dir = build_root / "torch_extensions" / extension_name
+    build_dir.mkdir(parents=True)
+    (build_dir / "lock").write_text("")  # orphan from an interrupted build
+
+    extension = native.load_extension(build_root=build_root)
+
+    assert extension is not None
+    assert seen["lock_present"] is False
+    assert native.extension_load_error() is None
+
+
 @pytest.mark.ci_cpu
 def test_native_build_wraps_sync_setup_errors(monkeypatch: pytest.MonkeyPatch) -> None:
     helper = native._native_build()