[Core] Add request-scoped producer state and fail processor exceptions

docs/design/feature/request_stage_state.md

@@ -0,0 +1,55 @@
+# Request-scoped producer state
+
+Custom incremental producers invoked by `OmniChunkTransferAdapter` can store
+typed state in a namespace under the adapter's existing `request_payload`
+container. This is a core implementation candidate for RFC #6453; model
+migrations and model-specific segment resets are separate work.
+
+```python
+from vllm_omni.distributed.omni_connectors.transfer_adapter.request_state import (
+    StageRequestIdentity,
+)
+
+identity = StageRequestIdentity(request.request_id, request.external_req_id)
+state = transfer_manager.request_state.get_or_create(
+    identity, namespace="producer.stage_edge", factory=ModelState
+)
+```
+
+The accessor is available during a queued custom producer invocation. Both IDs
+must come from the supplied request. A namespace returns the same state until
+whole-request release. Factories construct model state without connector I/O;
+they run outside the sender lock and may be called concurrently, so they must
+not perform external side effects.
+
+The runtime owns the request container. Producers may update their namespace
+but must not remove the container. Existing processors can continue using raw
+payload entries, including tensor-valued entries. A request cannot mix legacy
+raw entries and namespaced state: access fails instead of overwriting model
+data.
+
+| Outcome | State lifetime |
+| --- | --- |
+| Producer buffers data and returns no payload | Retained |
+| Nonterminal send succeeds | Retained |
+| Resumable segment completes | Retained |
+| Whole-request terminal send succeeds | Released by canonical cleanup |
+| Processor raises, or connector fails/raises | Retained until failure finalization |
+| Cancellation or timeout finalizes the request | Released by canonical cleanup |
+| Duplicate cleanup | Harmless |
+| External ID is reused after release | Fresh namespaces |
+
+Generation validation uses the existing sender token and sender lock. A
+retained accessor cannot recreate state after cancellation or release, even
+when a new request reuses the external ID. Factories that overlap cleanup also
+fail before committing their result. This does not make a retained model state
+object immutable: callers must not mutate it after their invocation finishes.
+
+Processor exceptions are recorded against the scheduler's internal request ID.
+They do not produce a successful empty terminal payload. Existing scheduler
+failure finalization releases receiver and sender state; this API adds no retry
+policy, timeout mechanism, or process-global request map.
+
+The CPU contract tests include a synthetic producer through the production
+save queue, sender, connector outcomes, failure collection, and finalization.
+They do not establish model correctness or a latency improvement.

docs/design/index.md

@@ -22,6 +22,7 @@ implementation contract; it is not, by itself, a general support claim.
 - [Disaggregated Inference](feature/disaggregated_inference.md)
 - [Host Weight Runtime](feature/host_weight_runtime.md)
 - [Async Chunk](feature/async_chunk.md)
+- [Request-scoped Producer State](feature/request_stage_state.md)
 - [Async Diffusion Output](feature/async_diffusion_output.md)
 - [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
 - [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md)

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -26,11 +26,151 @@
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from vllm_omni.distributed.omni_connectors.transfer_adapter.request_state import (
+    RequestStateUnavailableError,
+    StageRequestIdentity,
+)
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.parametrize("send_success", [True, False])
+def test_namespaced_producer_state_terminal_commit(build_adapter, send_success):
+    adapter, connector = build_adapter()
+    request = _req("internal", RequestStatus.FINISHED_STOPPED, external_req_id="external")
+    connector.put.return_value = (send_success, 1, {})
+    views = []
+
+    def producer(transfer_manager, request, **kwargs):
+        identity = StageRequestIdentity(request.request_id, request.external_req_id)
+        view = transfer_manager.request_state
+        views.append((view, identity))
+        view.get_or_create(identity, "synthetic.codec", list).append(7)
+        view.get_or_create(identity, "synthetic.prompt", dict)["ready"] = True
+        return OmniPayloadStruct(codes=CodesStruct(audio=torch.tensor([7])))
+
+    adapter.custom_process_next_stage_input_func = producer
+    adapter.save_async(None, request)
+    adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    assert connector.put.call_count == 1
+    if send_success:
+        assert "external" not in adapter.request_payload
+    else:
+        assert "external" in adapter.request_payload
+        assert "internal" in adapter.collect_failed_send_request_ids()
+        adapter.finish_requests(["internal"], RequestStatus.FINISHED_ERROR, {"internal": request})
+    assert "external" not in adapter.request_payload
+    view, identity = views[0]
+    with pytest.raises(RequestStateUnavailableError):
+        view.get_or_create(identity, "late", list)
+
+
+def test_namespaced_state_survives_segment_but_not_cancel_or_id_reuse(build_adapter):
+    adapter, connector = build_adapter()
+    views = []
+    states = []
+
+    def producer(transfer_manager, request, **kwargs):
+        identity = StageRequestIdentity(request.request_id, request.external_req_id)
+        view = transfer_manager.request_state
+        state = view.get_or_create(identity, "synthetic.codec", list)
+        state.append(request.request_id)
+        states.append(state)
+        views.append((view, identity))
+        return OmniPayloadStruct(codes=CodesStruct(audio=torch.tensor([1])))
+
+    adapter.custom_process_next_stage_input_func = producer
+    request = _req("old", RequestStatus.RUNNING, external_req_id="external")
+    request.resumable = True
+    request._omni_segment_generation = 0
+    adapter.save_async(None, request, is_segment_finished=True)
+    adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    assert "external" in adapter.request_payload
+    request._omni_segment_generation = 1
+    adapter.save_async(None, request, is_segment_finished=True)
+    adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    assert states[0] is states[1]
+    adapter.finish_requests(["old"], RequestStatus.FINISHED_ABORTED, {"old": request})
+    adapter.cleanup("old", "external")  # Duplicate finalization is harmless.
+    replacement = _req("new", RequestStatus.RUNNING, external_req_id="external")
+    adapter.save_async(None, replacement)
+    while adapter._pending_save_reqs:
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    assert states[-1] == ["new"]
+    old_view, old_identity = views[0]
+    with pytest.raises(RequestStateUnavailableError):
+        old_view.get_or_create(old_identity, "synthetic.codec", list)
+
+
+def test_processor_exception_fails_request_without_empty_terminal_send(build_adapter):
+    adapter, connector = build_adapter()
+    request = _req("internal", RequestStatus.FINISHED_STOPPED, external_req_id="external")
+
+    def producer(transfer_manager, request, **kwargs):
+        identity = StageRequestIdentity(request.request_id, request.external_req_id)
+        transfer_manager.request_state.get_or_create(identity, "synthetic.codec", list).append(1)
+        raise ValueError("invalid codec state")
+
+    adapter.custom_process_next_stage_input_func = producer
+    adapter.save_async(None, request)
+    adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    connector.put.assert_not_called()
+    assert "invalid codec state" in adapter.collect_failed_send_request_ids()["internal"]
+    assert "external" in adapter.request_payload
+    adapter.finish_requests(["internal"], RequestStatus.FINISHED_ERROR, {"internal": request})
+    assert "external" not in adapter.request_payload
+    with pytest.raises(RequestStateUnavailableError):
+        _ = adapter.request_state
+
+
+def test_namespaced_state_retains_buffered_output_and_isolates_finalization(build_adapter):
+    adapter, connector = build_adapter()
+    views = {}
+
+    def producer(transfer_manager, request, **kwargs):
+        identity = StageRequestIdentity(request.request_id, request.external_req_id)
+        view = transfer_manager.request_state
+        views[request.request_id] = (view, identity)
+        view.get_or_create(identity, "synthetic.buffer", list).append(request.request_id)
+        return None
+
+    adapter.custom_process_next_stage_input_func = producer
+    requests = {name: _req(name, RequestStatus.RUNNING) for name in ("first", "second")}
+    for request in requests.values():
+        adapter.save_async(None, request)
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    connector.put.assert_not_called()
+    assert set(adapter.request_payload) == {"first", "second"}
+    adapter.finish_requests(["first"], RequestStatus.FINISHED_ERROR, requests)
+    assert set(adapter.request_payload) == {"second"}
+    view, identity = views["second"]
+    assert view.get_or_create(identity, "synthetic.buffer", list) == ["second"]
+
+
+def test_raising_connector_records_failure_and_keeps_state_until_finalization(build_adapter):
+    adapter, connector = build_adapter()
+    request = _req("internal", RequestStatus.FINISHED_STOPPED, external_req_id="external")
+
+    def producer(transfer_manager, request, **kwargs):
+        identity = StageRequestIdentity(request.request_id, request.external_req_id)
+        transfer_manager.request_state.get_or_create(identity, "synthetic.codec", list).append(7)
+        return OmniPayloadStruct(codes=CodesStruct(audio=torch.tensor([7])))
+
+    def failed_put(**kwargs):
+        adapter.stop_event.set()
+        raise OSError("connector unavailable")
+
+    adapter.custom_process_next_stage_input_func = producer
+    connector.put.side_effect = failed_put
+    adapter.save_async(None, request)
+    adapter.save_loop()
+    assert "connector unavailable" in adapter.collect_failed_send_request_ids()["internal"]
+    assert "external" in adapter.request_payload
+    adapter.finish_requests(["internal"], RequestStatus.FINISHED_ERROR, {"internal": request})
+    assert "external" not in adapter.request_payload
+
+
 @pytest.fixture
 def shm_sender(build_adapter):
     adapter, _ = build_adapter(stage_id=0)

tests/distributed/omni_connectors/test_request_state.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import threading
+from dataclasses import dataclass
+
+import pytest
+
+from vllm_omni.distributed.omni_connectors.transfer_adapter.request_state import (
+    RequestStateAccessor,
+    RequestStateUnavailableError,
+    StageRequestIdentity,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class ProducerState:
+    chunks: int = 0
+
+
+def make_view(payloads, owners, identity, lock=None):
+    owner = owners[identity.external_request_id]
+    return RequestStateAccessor(
+        payloads,
+        identity,
+        owner,
+        lock or threading.Lock(),
+        lambda: owners.get(identity.external_request_id) is owner,
+    )
+
+
+def test_namespaces_and_interleaved_requests_are_isolated():
+    payloads: dict[str, object] = {}
+    owners = {"a": object(), "b": object()}
+    a = StageRequestIdentity("internal-a", "a")
+    b = StageRequestIdentity("internal-b", "b")
+    first = make_view(payloads, owners, a)
+    second = make_view(payloads, owners, b)
+    state = first.get_or_create(a, "producer.codec", ProducerState)
+    state.chunks = 3
+    assert first.get_or_create(a, "producer.codec", ProducerState) is state
+    assert first.get_or_create(a, "producer.prompt", ProducerState).chunks == 0
+    assert second.get_or_create(b, "producer.codec", ProducerState).chunks == 0
+
+
+def test_release_and_id_reuse_fence_retained_accessors():
+    payloads: dict[str, object] = {}
+    owners = {"external": object()}
+    old_id = StageRequestIdentity("old-internal", "external")
+    old = make_view(payloads, owners, old_id)
+    old.get_or_create(old_id, "codec", ProducerState).chunks = 4
+    payloads.pop("external", None)
+    owners.pop("external", None)
+    with pytest.raises(RequestStateUnavailableError):
+        old.get_or_create(old_id, "late", ProducerState)
+    owners["external"] = object()
+    new_id = StageRequestIdentity("new-internal", "external")
+    new = make_view(payloads, owners, new_id)
+    assert new.get_or_create(new_id, "codec", ProducerState).chunks == 0
+    with pytest.raises(RequestStateUnavailableError):
+        old.get_or_create(old_id, "codec", ProducerState)
+    with pytest.raises(RequestStateUnavailableError):
+        new.get_or_create(old_id, "codec", ProducerState)
+
+
+def test_cleanup_during_factory_cannot_resurrect_state():
+    payloads: dict[str, object] = {}
+    owners = {"external": object()}
+    identity = StageRequestIdentity("internal", "external")
+    lock = threading.Lock()
+    view = make_view(payloads, owners, identity, lock)
+
+    def factory():
+        # This would deadlock if the user factory held the sender lock.
+        with lock:
+            owners.pop("external")
+        return ProducerState()
+
+    with pytest.raises(RequestStateUnavailableError):
+        view.get_or_create(identity, "codec", factory)
+    assert payloads == {}
+
+
+def test_factory_failure_leaves_no_partial_container():
+    payloads: dict[str, object] = {}
+    owners = {"external": object()}
+    identity = StageRequestIdentity("internal", "external")
+    view = make_view(payloads, owners, identity)
+
+    def factory():
+        raise ValueError("invalid conditioning")
+
+    with pytest.raises(ValueError, match="invalid conditioning"):
+        view.get_or_create(identity, "codec", factory)
+    assert payloads == {}
+    assert view.get_or_create(identity, "codec", ProducerState).chunks == 0
+
+
+def test_legacy_state_is_not_overwritten():
+    legacy = {"frames": [1, 2]}
+    payloads = {"external": legacy}
+    owners = {"external": object()}
+    identity = StageRequestIdentity("internal", "external")
+    with pytest.raises(RequestStateUnavailableError, match="legacy"):
+        make_view(payloads, owners, identity).get_or_create(identity, "codec", ProducerState)
+    assert payloads["external"] is legacy
+
+
+@pytest.mark.parametrize("namespace", ["", "   "])
+def test_empty_namespace_is_rejected(namespace):
+    identity = StageRequestIdentity("internal", "external")
+    with pytest.raises(ValueError, match="namespace"):
+        make_view({}, {"external": object()}, identity).get_or_create(identity, namespace, ProducerState)

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -25,6 +25,7 @@
 from ..utils.kv_utils import get_local_tp_rank, get_omni_replica_id
 from ..utils.logging import get_connector_logger
 from .base import OmniTransferAdapterBase
+from .request_state import RequestStateAccessor, RequestStateUnavailableError, StageRequestIdentity
 
 logger = get_connector_logger(__name__)
 
@@ -153,6 +154,7 @@ def __init__(self, vllm_config: Any):
         # segment's queued or in-flight work after the same request id resumes.
         self._registered_load_entries: dict[str, _LoadEntry] = {}
         self._sender_tokens: dict[str, _SenderGeneration] = {}
+        self._processor_state_context = threading.local()
         self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs
         (
             self._max_model_len,
@@ -195,7 +197,7 @@ def __init__(self, vllm_config: Any):
         self._adaptive_states: dict[str, Any] = {}
         self.upstream_exhausted_requests: set[str] = set()
         self.segment_finished_requests: set[str] = set()
-        self.request_payload = {}
+        self.request_payload: dict[str, Any] = {}
         self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list)
         self.request_ids_mapping: dict[str, str] = {}
         # Save-thread only: running count of segments reclaimed after finish.
@@ -235,6 +237,14 @@ def __init__(self, vllm_config: Any):
         # OmniSchedulingCoordinator._waiting_since on the full-payload path.
         self._waiting_since: dict[str, float] = {}
 
+    @property
+    def request_state(self) -> RequestStateAccessor:
+        """Namespaced state for the current custom producer invocation."""
+        accessor = getattr(self._processor_state_context, "accessor", None)
+        if accessor is None:
+            raise RequestStateUnavailableError("Request state is only available inside a queued producer invocation")
+        return accessor
+
     @staticmethod
     def _is_truthy_scalar(value: Any) -> bool:
         if isinstance(value, torch.Tensor):
@@ -803,10 +813,22 @@ def _send_single_request_for_generation(
                 }
                 if self._accepts_new_token_ids(processor):
                     processor_kwargs["new_token_ids"] = task.get("new_token_ids", ())
+                if sender_token is not None:
+                    self._processor_state_context.accessor = RequestStateAccessor(
+                        self.request_payload,
+                        StageRequestIdentity(request.request_id, external_req_id),
+                        sender_token,
+                        self._sender_state_lock,
+                        lambda: self._sender_tokens.get(external_req_id) is sender_token and not sender_token.cancelled,
+                    )
                 payload_data = processor(**processor_kwargs)
 
             except Exception as e:
                 logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")
+                self.record_send_failure(request.request_id, f"Stage input processor failed: {e}")
+                return
+            finally:
+                self._processor_state_context.accessor = None
 
         if payload_data is None:
             if not (is_segment_finished or is_finished):