Merge branch 'main' into codex-hwr-domain-timeout

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/design/feature/host_weight_runtime.md

@@ -313,8 +313,15 @@ One process per exact identity owns a build; other workers wait and then acquire
 leases for the published artifact. Publication is invisible until all payloads
 and metadata are validated, hashed, fsynced, and atomically renamed.
 
-`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not
-cancel synchronous validation, a producer that has already started, or atomic
+`coordination_timeout_seconds` bounds domain-initialization and lookup/build
+lock acquisition. Store construction and each later resolution or publication
+operation have separate budgets from the same wait policy, rather than one
+end-to-end startup deadline. A domain-init timeout follows the retryable domain
+failure policy below. After contention ends, a fresh runtime construction can
+retry initialization; the timed-out runtime retains its original failure.
+
+The coordination budget does not cancel filesystem I/O, synchronous validation,
+a producer that has already started, or atomic
 publication. A hung in-process producer therefore blocks its owning process and
 must be handled by external process supervision. Enforceable producer
 cancellation requires a future process-isolated producer contract.

docs/design/module/host_weight_runtime.md

@@ -136,8 +136,16 @@ so only when it is marked retryable; unsupported capabilities, identity
 collisions, producer failures, and publication failures remain visible as
 startup failures.
 
-`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it
-is not a hard wall-clock deadline for in-process validation, production, or
+`coordination_timeout_seconds` bounds acquisition of the domain-initialization
+lock and lookup/build locks. Store construction and each later resolution or
+publication operation receive separate coordination budgets from the same
+`WaitPolicy`; they do not share an end-to-end startup deadline. A domain-init
+lock timeout is a retryable domain failure: preferred mode may fall back and
+required mode fails. Retrying initialization requires constructing a fresh
+runtime after contention ends; a waiter never cancels or unlocks the owner.
+
+The coordination budget is not a hard wall-clock deadline for filesystem I/O,
+in-process validation, production, or
 atomic publication. Once this V1 implementation becomes the producer, the
 synchronous producer and publication run to completion. A hung producer blocks
 its owning process and requires external process supervision; enforceable

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -4,10 +4,8 @@
 import asyncio
 import multiprocessing as mp
 import queue
-import signal
 import threading
 import time
-import weakref
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock
 
@@ -1332,7 +1330,6 @@ class TestExecutorShutdownCleaner:
         class FakeProcess:
             def __init__(self, name):
                 self.name = name
-                self.pid = 123
                 self.alive = True
                 self.terminated = False
                 self.join_timeouts = []
@@ -1361,132 +1358,6 @@ class TestExecutorShutdownCleaner:
         assert first.terminated and second.terminated
         assert not first.is_alive() and not second.is_alive()
 
-    @pytest.mark.parametrize("cooperative", [False, True])
-    def test_real_workers_exit_and_are_reaped(self, monkeypatch, cooperative):
-        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
-
-        ctx = mp.get_context("fork")
-        ready, stop = ctx.Event(), ctx.Event()
-
-        def run_worker():
-            signal.signal(signal.SIGTERM, signal.SIG_IGN)
-            ready.set()
-            stop.wait()
-
-        dead = ctx.Process(target=lambda: None)
-        worker = ctx.Process(target=run_worker)
-        monkeypatch.setattr(executor_module, "_WORKER_SHUTDOWN_GRACE_S", 0.2)
-        monkeypatch.setattr(executor_module, "_WORKER_TERMINATE_GRACE_S", 0.1)
-        monkeypatch.setattr(executor_module, "_WORKER_KILL_GRACE_S", 5.0, raising=False)
-        dead.start()
-        worker.start()
-        try:
-            dead.join(5)
-            assert not dead.is_alive()
-            assert ready.wait(5), "worker did not install its signal handler"
-            mq = Mock()
-            if cooperative:
-                mq.enqueue.side_effect = lambda *args, **kwargs: stop.set()
-            else:
-                mq.enqueue.side_effect = OSError("shutdown queue unavailable")
-            cleaner = executor_module._ExecutorShutdownCleaner(mq, 2, [dead, worker])
-
-            cleaner()
-
-            assert not worker.is_alive()
-            assert worker.exitcode == (0 if cooperative else -signal.SIGKILL)
-            assert cleaner.processes == []
-            assert cleaner.broadcast_mq is None
-            cleaner()
-        finally:
-            for proc in (dead, worker):
-                if proc.is_alive():
-                    proc.kill()
-                proc.join(5)
-                proc.close()
-
-    @pytest.mark.parametrize("failed_action", ["terminate", "kill"])
-    def test_kill_phase_shares_deadline_and_continues_after_os_errors(self, monkeypatch, failed_action):
-        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
-
-        first, second = Mock(pid=101), Mock(pid=102)
-        first.name, second.name = "first", "second"
-        for proc in (first, second):
-            proc.is_alive.side_effect = [True, True, True, False]
-        getattr(first, failed_action).side_effect = OSError("signal failed")
-        first.join.side_effect = [OSError("join failed"), None, None]
-        monotonic = Mock(side_effect=[100, 100, 110, 120, 120, 124, 130, 130, 134])
-        monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
-        cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
-
-        cleaner()
-
-        first.kill.assert_called_once()
-        second.kill.assert_called_once()
-        assert [call.args[0] for call in first.join.call_args_list] == [15, 5, 5]
-        assert [call.args[0] for call in second.join.call_args_list] == [5, 1, 1]
-        assert cleaner.processes == []
-
-    def test_executor_retains_survivor_for_explicit_shutdown_retry(self, monkeypatch):
-        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
-
-        survivor = Mock(pid=123)
-        survivor.name = "surviving-worker"
-        survivor.is_alive.return_value = True
-        cleaner = executor_module._ExecutorShutdownCleaner(processes=[survivor])
-        executor, _, _ = _make_executor()
-        executor._shutdown_cleaner = cleaner
-        executor._processes = [survivor]
-        executor._finalizer = weakref.finalize(executor, cleaner)
-        executor._pump_stop = threading.Event()
-        executor._futures_lock = threading.RLock()
-        executor._rpc_futures = {}
-        executor._output_futures = {}
-        executor._batch_split_map = {}
-        log_error = Mock()
-        monkeypatch.setattr(executor_module.logger, "error", log_error)
-
-        executor.shutdown()
-
-        assert not executor._finalizer.alive
-        assert executor._shutdown_cleaner is cleaner
-        assert executor._processes == [survivor]
-        log_error.assert_called_once()
-        assert log_error.call_args.args[1] == [("surviving-worker", 123)]
-        survivor.is_alive.return_value = False
-
-        executor.shutdown()
-        executor.shutdown()
-
-        assert executor._shutdown_cleaner is None
-        assert executor._processes == []
-
-    def test_concurrent_cleanup_does_not_repeat_process_operations(self):
-        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
-
-        joining, release = threading.Event(), threading.Event()
-        proc = Mock(pid=123)
-        proc.is_alive.return_value = True
-
-        def join(timeout):
-            joining.set()
-            assert release.wait(5)
-            proc.is_alive.return_value = False
-
-        proc.join.side_effect = join
-        cleaner = executor_module._ExecutorShutdownCleaner(processes=[proc])
-        thread = threading.Thread(target=cleaner)
-        thread.start()
-        try:
-            assert joining.wait(5)
-            cleaner()
-            proc.join.assert_called_once()
-        finally:
-            release.set()
-            thread.join(5)
-        assert not thread.is_alive()
-        assert cleaner.processes == []
-
 
 # ───────── monitor thread & death sentinel integration tests ─────────
 
@@ -1535,7 +1406,7 @@ class TestMultiprocExecutorWorkerMonitor:
         executor._result_mq = None
         executor._shutdown_cleaner = None
         # Use a no-op so shutdown() doesn't crash on None resources.
-        executor._finalizer = weakref.finalize(executor, lambda: None)
+        executor._finalizer = lambda: None
         # ------------------------------------------------------------------
         # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates
         # over them, so they need to exist even when constructed via __new__.
@@ -1570,7 +1441,7 @@ class TestMultiprocExecutorWorkerMonitor:
         executor._broadcast_mq = None
         executor._result_mq = None
         executor._shutdown_cleaner = None
-        executor._finalizer = weakref.finalize(executor, lambda: None)
+        executor._finalizer = lambda: None
 
         proc = _make_short_lived_process()
         executor._processes = [proc]

tests/host_weight_runtime/test_runtime_resolution.py

@@ -4,6 +4,10 @@
 
 from __future__ import annotations
 
+import multiprocessing as mp
+import threading
+import time
+from multiprocessing.connection import Connection
 from pathlib import Path
 from typing import cast
 
@@ -15,6 +19,7 @@ from vllm_omni.host_weight_runtime import (
     CanonicalJson,
     ComponentIdentity,
     FailureCode,
+    HostWeightError,
     HostWeightFailure,
     HostWeightLease,
     HostWeightLeaseCarrier,
@@ -40,12 +45,14 @@ from vllm_omni.host_weight_runtime import (
     StoreStatus,
     TensorWriteSpec,
     ValidationLevel,
+    WaitPolicy,
     WeightArtifactIdentity,
     WeightProductionSpec,
     WeightRepresentation,
     WeightSourceIdentity,
 )
 from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -518,6 +525,107 @@ def test_retryable_store_initialization_failure_obeys_runtime_mode(
     assert resolution.report.attempts[0].failure.retryable
 
 
+def _initialize_contended_runtime(config: HostWeightRuntimeConfig, output: Connection) -> None:
+    try:
+        output.send("starting")
+        started = time.monotonic()
+        runtime = HostWeightRuntime.from_config(config)
+        resolution = runtime.resolve(_identity())
+        try:
+            output.send((time.monotonic() - started, resolution.report))
+        finally:
+            if resolution.lease is not None:
+                resolution.lease.close()
+    finally:
+        output.close()
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [
+        (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK),
+        (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+    ],
+)
+def test_domain_lock_timeout_and_fresh_runtime_recovery(
+    tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome
+) -> None:
+    config = HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path), wait=WaitPolicy(0.1))
+    initial = HostWeightRuntime.from_config(config)
+    built = initial.resolve(_identity(), producer=CountingProducer(_identity()))
+    assert built.lease is not None
+    built.lease.close()
+    metadata = {name: (tmp_path / name).read_bytes() for name in ("domain.json", "domain-policy.json")}
+    lock_path = tmp_path / "locks" / "domain-init.lock"
+    ctx = mp.get_context("spawn")
+    reader, writer = ctx.Pipe(duplex=False)
+    process = ctx.Process(target=_initialize_contended_runtime, args=(config, writer))
+
+    with FileLock(lock_path, exclusive=True, deadline=None):
+        process.start()
+        writer.close()
+        try:
+            # Import/bootstrap is outside the short coordination budget.
+            assert reader.poll(60), "waiter did not start"
+            assert reader.recv() == "starting"
+            assert reader.poll(5), "domain initialization did not respect the coordination timeout"
+            elapsed, report = reader.recv()
+            assert 0.1 <= elapsed < 5
+            assert report.outcome is expected
+            failure = report.attempts[0].failure
+            assert failure is not None
+            assert failure.stage is ResolutionStage.DOMAIN
+            assert failure.code is FailureCode.DOMAIN_UNAVAILABLE
+            assert failure.retryable
+            assert "timed out waiting for" in failure.message
+            assert "domain-init.lock" in failure.message
+            assert lock_is_active(lock_path), "waiter must not unlock the owner"
+            process.join(5)
+            assert process.exitcode == 0
+        finally:
+            if process.is_alive():
+                process.kill()
+            process.join(5)
+            process.close()
+            reader.close()
+
+    assert {name: (tmp_path / name).read_bytes() for name in metadata} == metadata
+    recovered = HostWeightRuntime.from_config(config).resolve(_identity())
+    assert recovered.report.outcome is ResolutionOutcome.LOCAL_HIT
+    assert recovered.lease is not None
+    recovered.lease.close()
+
+
+def test_direct_store_waits_for_domain_lock_released_before_deadline(tmp_path: Path) -> None:
+    config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path))
+    assert config.domain is not None
+    initial = HostWeightRuntime.from_config(config)
+    assert initial.store is not None
+    lock_path = tmp_path / "locks" / "domain-init.lock"
+    with FileLock(lock_path, exclusive=True, deadline=None) as owner:
+        release = threading.Timer(0.1, owner.close)
+        release.start()
+        try:
+            store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(5))
+        finally:
+            release.join(5)
+    assert store.root == tmp_path
+
+
+def test_direct_store_domain_timeout_is_typed(tmp_path: Path) -> None:
+    config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path))
+    assert config.domain is not None
+    HostWeightRuntime.from_config(config)
+    lock_path = tmp_path / "locks" / "domain-init.lock"
+    with FileLock(lock_path, exclusive=True, deadline=None):
+        with pytest.raises(HostWeightError) as caught:
+            FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(0.01))
+        assert caught.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+        assert caught.value.failure.retryable
+        assert lock_is_active(lock_path)
+    assert not lock_is_active(lock_path)
+
+
 @pytest.mark.parametrize(
     ("stage", "code"),
     [

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 from __future__ import annotations
 
 import concurrent.futures
@@ -13,7 +10,7 @@ import threading
 import time
 import weakref
 from collections.abc import Callable
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from multiprocessing.synchronize import Event
 from typing import TYPE_CHECKING, Any, cast
 
@@ -45,7 +42,6 @@ _DEQUEUE_TIMEOUT_S = 5.0
 _DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0))
 _WORKER_SHUTDOWN_GRACE_S = 15.0
 _WORKER_TERMINATE_GRACE_S = 5.0
-_WORKER_KILL_GRACE_S = 5.0
 _RESULT_PUMP_JOIN_TIMEOUT_S = 2.0
 
 
@@ -84,61 +80,33 @@ class _ExecutorShutdownCleaner:
     broadcast_mq: MessageQueue | None = None
     num_workers: int = 0
     processes: list[mp.Process] | None = None
-    _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
 
     def __call__(self) -> None:
         """Clean up background resources."""
-        # The worker monitor and an explicit shutdown may race. A retry must
-        # not signal/join the same Process objects while cleanup is in flight.
-        if not self._lock.acquire(blocking=False):
-            return
-        try:
-            self._cleanup()
-        finally:
-            self._lock.release()
-
-    def _cleanup(self) -> None:
         if self.broadcast_mq is not None:
             try:
                 for _ in range(self.num_workers):
                     self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0)
 
+                self.broadcast_mq = None
             except Exception as exc:
                 logger.warning("Failed to send shutdown signal: %s", exc)
-            finally:
-                self.broadcast_mq = None
 
         if self.processes:
+            join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S
+            for proc in self.processes:
+                if not proc.is_alive():
+                    continue
+                proc.join(max(0.0, join_deadline - time.monotonic()))
+
             alive = [proc for proc in self.processes if proc.is_alive()]
-            for action, grace in (
-                (None, _WORKER_SHUTDOWN_GRACE_S),
-                ("terminate", _WORKER_TERMINATE_GRACE_S),
-                ("kill", _WORKER_KILL_GRACE_S),
-            ):
-                if not alive:
-                    break
-                if action is not None:
-                    for proc in alive:
-                        try:
-                            logger.warning("Calling %s on diffusion worker %s (pid=%s)", action, proc.name, proc.pid)
-                            getattr(proc, action)()
-                        except OSError:
-                            logger.exception("Failed to %s diffusion worker %s (pid=%s)", action, proc.name, proc.pid)
+            for proc in alive:
+                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
+                proc.terminate()
 
-                deadline = time.monotonic() + grace
-                for proc in alive:
-                    try:
-                        proc.join(max(0.0, deadline - time.monotonic()))
-                    except OSError:
-                        logger.exception("Failed to join diffusion worker %s (pid=%s)", proc.name, proc.pid)
-                alive = [proc for proc in alive if proc.is_alive()]
-
-            self.processes = alive
-            if alive:
-                logger.error(
-                    "Diffusion worker cleanup incomplete after kill: %s; retaining processes for shutdown retry",
-                    [(proc.name, proc.pid) for proc in alive],
-                )
+            terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S
+            for proc in alive:
+                proc.join(max(0.0, terminate_deadline - time.monotonic()))
 
 
 class MultiprocDiffusionExecutor(DiffusionExecutor):
@@ -1038,12 +1006,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
     def shutdown(self) -> None:
         self._closed = True
         self._pump_stop.set()
-        cleaner = self._shutdown_cleaner
         try:
-            if self._finalizer.alive:
-                self._finalizer()
-            elif cleaner is not None:
-                cleaner()
+            self._finalizer()
         finally:
             pump_threads = getattr(self, "_result_pump_threads", [])
             for thread in pump_threads:
@@ -1067,6 +1031,5 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 self._rpc_futures.clear()
                 self._output_futures.clear()
                 self._batch_split_map.clear()
-            self._processes = (cleaner.processes or []) if cleaner is not None else []
-            if not self._processes:
-                self._shutdown_cleaner = None
+            self._shutdown_cleaner = None
+            self._processes = []

vllm_omni/host_weight_runtime/filesystem/store.py

@@ -23,7 +23,7 @@ import regex as re
 import torch
 from safetensors import SafetensorError, safe_open
 
-from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel
+from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel, WaitPolicy
 from ..errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionStage
 from ..identity import CanonicalJson, WeightArtifactIdentity, canonical_json
 from ..inspection import (
@@ -429,6 +429,8 @@ class FilesystemHostWeightStore:
         domain: StorageDomainPolicy,
         capacity: CapacityPolicy,
         integrity: IntegrityPolicy,
+        *,
+        wait: WaitPolicy = WaitPolicy(),
     ) -> None:
         self.domain_policy = domain
         self.capacity_policy = capacity
@@ -464,7 +466,7 @@ class FilesystemHostWeightStore:
                 )
             )
         try:
-            self._initialize_domain()
+            self._initialize_domain(deadline=time.monotonic() + wait.coordination_timeout_seconds)
         except HostWeightError:
             raise
         except OSError as exc:
@@ -485,7 +487,7 @@ class FilesystemHostWeightStore:
                 )
             ) from exc
 
-    def _initialize_domain(self) -> None:
+    def _initialize_domain(self, *, deadline: float) -> None:
         try:
             self.filesystem_type = detect_filesystem_type(self.root)
             detected_storage_class = _storage_class_for_filesystem_type(self.filesystem_type)
@@ -545,7 +547,7 @@ class FilesystemHostWeightStore:
                         )
                     )
 
-        with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=None):
+        with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=deadline):
             domain_path = self.root / _DOMAIN_FILE
             if domain_path.exists():
                 domain_metadata = _read_json_file(domain_path)

vllm_omni/host_weight_runtime/runtime.py

@@ -70,7 +70,7 @@ class HostWeightRuntime:
         from .filesystem import FilesystemHostWeightStore  # noqa: PLC0415
 
         try:
-            store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity)
+            store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=config.wait)
         except HostWeightError as exc:
             return cls(
                 config,