diff --git a/docs/zh/design_lingbot_condition_prefetch.md b/docs/zh/design_lingbot_condition_prefetch.md new file mode 100644 index 0000000..5a058c9 --- /dev/null +++ b/docs/zh/design_lingbot_condition_prefetch.md @@ -0,0 +1,266 @@ +# LingBot-World-Fast Condition 预取流水设计 + +> 状态：当前实现设计说明。本文描述内部调度策略，不新增用户配置、环境变量或服务协议。 + +## 1. 最终状态：完整运行时序 + +```mermaid +sequenceDiagram + autonumber + participant C as 控制客户端 + participant S as LingBot Service + participant P as Pipeline Runtime + participant O as Streaming Orchestrator + participant E as VAE Encode Actor + participant D as DiT Actor + participant V as VAE Decode Actor + + C->>S: 建立流式 session + S->>P: _create_initialized_session(config) + P->>P: encode prompt、准备图像、计算 latent/KV 几何 + P->>E: initialize encode cache(image, cache_handle) + E-->>P: encode cache ready + P->>V: initialize decode cache(cache_handle) + V-->>P: decode cache ready + P->>D: initialize denoise/KV/noise cache(cache_handle) + D-->>P: denoise cache ready + P-->>S: initialized runtime + + S->>O: create_session(runtime, final_sequence_id) + O-)E: enqueue condition[0] + O-)E: enqueue condition[1] + O-->>S: streaming session ready + + loop 每个 chunk i + par 与 control 无关的 condition 路径 + E-->>O: condition[i] ready + and 真实交互控制路径 + C->>S: control[i] + S->>P: resolve + validate control[i] + P-->>S: control tensor[i] + S->>O: try_submit control[i] + Note over S,O: control[i] 被接受时记录 latency anchor + end + + O->>O: 按 session/sequence join condition[i] + control[i] + + alt World-KV 已命中 latent[i] + O->>D: advance noise cursor + D-->>O: cached latent[i] + else 正常生成 + O->>D: denoise(condition[i], control[i], session caches) + D-->>O: latent[i] + end + + par Post-denoise condition refill + opt i+2 未越界且预取窗口有容量 + O-)E: enqueue condition[i+2] + end + and 当前 chunk decode + O->>V: decode latent[i] + V-->>O: frames[i] + end + + O->>O: 按 sequence ID 有序提交输出与 scheduler metrics + O-->>S: frames[i] + S-->>C: chunk[i] + applied controls + target facts + end + + C->>S: stop / disconnect / session complete + S->>O: close_session(drain) + O->>V: release decode cache + V-->>O: released + O->>D: release denoise/KV/noise cache + D-->>O: released + O->>E: release encode cache + E-->>O: released + O-->>S: session state released +``` + +图中的 `par` 表示两条路径之间没有数据依赖，不承诺它们在物理设备上同时执行。condition 可能早于 control 完成， +也可能在 control 到达后才完成；DiT 始终等待同一 sequence ID 的两项输入都就绪。 + +最终结构中有三条必须分开的路径： + +- `condition`：只依赖 session 图像、VAE causal state 和 chunk 位置，可以在 control 之前计算； +- `control`：来自真实交互输入，必须按 chunk 严格有序，不能预测或跨 chunk 复用； +- `latent/frames`：只有同一 sequence ID 的 condition 与 control join 后才能产生。 + +VAE Encode、DiT 和 VAE Decode 分别由唯一 actor 管理。它们即使放在同一张 GPU 上也可以重叠；调度器不会根据 +device placement 隐式创建 resource group。session 结束时，cache 通过 owning actor 按 +`decode → denoise → encode` 的逆拓扑顺序释放。 + +## 2. 最终状态的准入与状态边界 + +每个 session 维护两个单调递增 cursor： + +- `next_condition_index`：下一条尚未提交的 condition 序号； +- `next_control_index`：下一条允许接收的 control 序号。 + +预取窗口满足： + +```text +0 <= next_condition_index - next_control_index <= 2 +``` + +control 准入必须同时满足： + +1. `chunk_index == next_control_index`； +2. 对应 condition 已提交，或当前仍可补交该 condition； +3. control edge 仍有容量。 + +只有 scheduler 接受 control 后，control cursor 才递增。重复、跳号和乱序 control 直接失败。 +`_ensure_next_condition_locked()` 只在预取被 backpressure 暂时耗尽时补齐当前 control 所需的一个 condition； +正常窗口补充由当前 chunk denoise 完成后的 refill 触发。 + +预取深度 `2` 是内部常量，不是用户配置。condition、control、latent 和 output edge 也都具有显式的 per-session +容量，因此长 session 不会形成 duration-sized tensor 列表。Condition actor 自身仍然串行执行；两级预取表示 +最多允许一个任务执行、另一个任务或结果停留在有界路径中，不表示同一个 VAE worker 同时运行两个 encode。 + +## 3. 本次 Diff：优化前后时序对比 + +### 3.1 优化前 + +```mermaid +sequenceDiagram + autonumber + participant C as 控制客户端 + participant S as LingBot Service + participant O as Streaming Orchestrator + participant E as VAE Encode Actor + participant D as DiT Actor + participant V as VAE Decode Actor + + C->>S: first control[0] + Note over S: 收到首个 control 后才初始化 + S->>S: create runtime + session + + loop 每个 chunk i + S->>O: atomic push(encode_request[i], control[i]) + O->>E: encode condition[i] + E-->>O: condition[i] + O->>D: condition[i] + control[i] + D-->>O: latent[i] + O->>V: decode latent[i] + V-->>O: frames[i] + O-->>S: frames[i] + S-->>C: chunk[i] + end + + Note over C,V: control 等待、condition encode、denoise、decode 更容易形成串行链 +``` + +优化前，`encode_request` 与 `control` 必须在一次原子 ingress 中提交。condition 明明与当前 control 无关，却必须 +等待 control 到达；首个 control 之后还要承担 runtime/session 初始化，直接拉长第一条交互关键路径。 + +### 3.2 优化后 + +```mermaid +sequenceDiagram + autonumber + participant C as 控制客户端 + participant S as LingBot Service + participant O as Streaming Orchestrator + participant E as VAE Encode Actor + participant D as DiT Actor + participant V as VAE Decode Actor + + S->>S: create runtime + session + S->>O: create_session(runtime) + O-)E: prefetch condition[0] + O-)E: prefetch condition[1] + + par Condition 可以提前完成 + E-->>O: condition[0] + and 等待真实控制 + C->>S: control[0] + S->>O: submit control[0] + Note over S,O: 从 control acceptance 开始计时 + end + + O->>D: join(condition[0], control[0]) + D-->>O: latent[0] + + par 后置补充未来 condition + O-)E: prefetch condition[2] + and 解码当前 latent + O->>V: decode latent[0] + V-->>O: frames[0] + end + + O-->>S: ordered frames[0] + S-->>C: chunk[0] +``` + +### 3.3 变化分析 + +| 维度 | 优化前 | 优化后 | +|---|---|---| +| 首个 control | 先等待 control，再初始化 runtime/session | 先初始化并预取，再等待 control | +| Condition ingress | 与 control 原子提交 | 与 control 解耦，最多领先 2 个 chunk | +| DiT 启动条件 | encode 在 control 后开始，随后 join | condition 可提前就绪，control 到达后即可 join | +| 后续 refill | 下一 control 到达后才触发 encode | 当前 denoise 完成后补充窗口 | +| 主要重叠 | 重叠机会少，容易形成串行链 | 未来 VAE encode 主要与当前 VAE decode 重叠 | +| Control 顺序 | 依赖原子 ingress 隐式维持 | `next_control_index` 显式拒绝重复、跳号和乱序 | +| 延迟起点 | 第一条任意 ingress | `latency_anchor_artifact="control"` | +| 内存边界 | edge 容量有界 | 保持有界，并增加固定深度 2 的 condition lookahead | + +Post-denoise refill 是调度偏好，不是“VAE encode 永不与 DiT 重叠”的硬互斥保证。如果 control 已到达而匹配的 +condition 尚未提交，活性保护仍可补交该 condition，以免 session 停滞。 + +## 4. 指标语义 + +Condition 预取可能早于 control 数秒。如果仍使用第一条 ingress 的时间作为起点，scheduler 会把用户尚未发送 +control 的等待时间错误计入 control-to-output latency。 + +`StreamingPipelineSpec.latency_anchor_artifact="control"` 规定： + +- condition-only ingress 不写入 `ingress_accepted_at`； +- control 被接受时才记录该 sequence 的 ingress 时间； +- 对应输出发出后，scheduler 才计算 control-to-output latency。 + +该字段默认是 `None`，其他 pipeline 保持“第一条任意 ingress”为起点的旧行为。它只修正 scheduler 指标语义， +不会合并 target chunk compute、客户端 delivery latency 或 AIPerf warmup/聚合职责。 + +## 5. 收益与代价 + +以下数据来自固定 4×H100、`chunk_size=3`、SageAttention SM90 的累计 checkpoint；它们不是严格的单开关 A/B： + +| Checkpoint | Chunk mean | 相对上一阶段 | 相对初始基线 | +|---|---:|---:|---:| +| 4 卡 SageAttention 基线 | 1.800984s | - | - | +| 首轮 condition/cache 版本 | 1.695177s | -5.9% | -5.9% | +| Condition 与 control 解耦预取 | 1.579448s | -6.8% | -12.3% | +| Post-denoise refill | 1.449961s | -8.2% | -19.5% | + +按相邻累计 checkpoint 看，post-denoise refill 是当时最大的单项下降；按正确性依赖看，应把 runtime 前置、 +condition/control 解耦、两级预取、后置 refill、control 顺序和延迟锚点作为一个整体。 + +主要代价：连接建立后即使用户一直不发送 control，也会提前占用 runtime、三类 session cache 和最多两个 +condition slot。断连、停止、stage failure 或预取失败时，仍由原有 session cleanup 释放这些状态。 + +本设计不改变扩散步数、模型权重、dtype、attention 实现、VAE 数值路径或输出顺序，属于无损调度优化。 + +## 6. 实现映射 + +| 设计职责 | 实现位置 | +|---|---| +| 首个 control 前创建 runtime/session | `telefuser/pipelines/lingbot_world_fast/service.py::_run_actor_worker_loop` | +| Session cursor、两级预取和 control 准入 | `telefuser/pipelines/lingbot_world_fast/streaming.py` | +| Post-denoise refill | `LingBotWorldFastStreamingRuntime::_denoise` | +| Control latency anchor | `telefuser/orchestrator/streaming_pipeline_orchestrator.py` | +| 时序、顺序与指标回归 | `tests/unit/pipelines/lingbot_world_fast/`、`tests/unit/orchestrator/` | + +## 7. 验证要求 + +功能测试至少覆盖： + +- session 创建后、control 到达前已提交两个 condition； +- denoise 完成后触发 refill，并补入下一 condition； +- 重复和乱序 control 被拒绝； +- condition-only ingress 不启动 control latency timer； +- 单 chunk、多 session、World-KV hit、stage failure 和 cleanup failure 行为不回退。 + +性能复测必须固定模型、分辨率、`chunk_size`、扩散 schedule、attention 实现、GPU 数量、placement 和 AIPerf +warmup 口径。正式结论至少比较多轮 mean、P90、P99、加权 compute FPS，并同时检查 GPU/VAE/DiT 时序，避免 +把单轮调度抖动误判为稳定收益。 diff --git a/docs/zh/stream_scheduler.md b/docs/zh/stream_scheduler.md index f0aaa23..155b694 100644 --- a/docs/zh/stream_scheduler.md +++ b/docs/zh/stream_scheduler.md @@ -102,3 +102,6 @@ LingBot 的 `vae_encode_config` 和 `vae_decode_config` 是两个独立且完整 - session state 必须隔离，并通过 owning actor 释放。 - 只为真实且明确的部署约束声明 resource group。 - 应验证 session 交错、backpressure、取消、actor failure 和 cleanup failure。 + +LingBot-World-Fast 的完整 condition/control 流水、最终时序和本次优化前后对比见 +[LingBot-World-Fast Condition 预取流水设计](design_lingbot_condition_prefetch.md)。 diff --git a/examples/run_examples.py b/examples/run_examples.py index 85aa064..a4ef087 100644 --- a/examples/run_examples.py +++ b/examples/run_examples.py @@ -1032,6 +1032,19 @@ def _emit_result(data: dict) -> None: print(f"{_RESULT_MARKER}{json.dumps(data)}", flush=True) +def _close_pipeline(pipeline: object | None) -> None: + """Release pipeline-owned workers without masking the regression result.""" + if pipeline is None: + return + close = getattr(pipeline, "close", None) + if not callable(close): + return + try: + close() + except Exception as exc: + print(f"Warning: failed to close pipeline: {exc}", file=sys.stderr, flush=True) + + def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | None) -> None: """Subprocess entry point: load, run, save one pipeline. @@ -1129,7 +1142,8 @@ def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | No "peak_gpu_memory_mb": round(gpu_mem_peak, 2), } ) - del pipeline + _close_pipeline(pipeline) + pipeline = None gc.collect() sys.exit(1) @@ -1156,7 +1170,8 @@ def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | No "peak_gpu_memory_mb": round(gpu_mem_peak, 2), } ) - pipeline = None # noqa: F841 + _close_pipeline(pipeline) + pipeline = None gc.collect() sys.exit(1) @@ -1198,7 +1213,8 @@ def _run_single(pipeline_key: str, config_path: str | None, output_dir: str | No } ) - del pipeline # noqa: F821 + _close_pipeline(pipeline) + pipeline = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() diff --git a/telefuser/orchestrator/streaming_pipeline_orchestrator.py b/telefuser/orchestrator/streaming_pipeline_orchestrator.py index 019f6db..e593454 100644 --- a/telefuser/orchestrator/streaming_pipeline_orchestrator.py +++ b/telefuser/orchestrator/streaming_pipeline_orchestrator.py @@ -141,10 +141,13 @@ class StreamingPipelineSpec: output_artifacts: frozenset[str] = frozenset() output_capacity_per_session: int = 1 resource_groups: tuple[StreamingResourceGroupSpec, ...] = () + latency_anchor_artifact: str | None = None def __post_init__(self) -> None: if self.output_capacity_per_session < 1: raise ValueError("output_capacity_per_session must be at least one") + if self.latency_anchor_artifact == "": + raise ValueError("latency_anchor_artifact must be non-empty when provided") @dataclass(frozen=True) @@ -561,7 +564,8 @@ def try_push_inputs(self, session_id: str, sequence_id: int, inputs: Mapping[str for edge in edges ): return False - runtime.ingress_accepted_at.setdefault(sequence_id, time.monotonic()) + if self.spec.latency_anchor_artifact is None or self.spec.latency_anchor_artifact in inputs: + runtime.ingress_accepted_at.setdefault(sequence_id, time.monotonic()) for artifact, value in inputs.items(): runtime.artifacts[(sequence_id, artifact)] = _ArtifactSlot( value, {edge.target_stage for edge in edges_by_artifact[artifact]} @@ -1033,6 +1037,11 @@ def _validate(self) -> None: for artifact in self.spec.output_artifacts: if artifact not in declared_producers: raise ValueError(f"Output artifact {artifact!r} has no producer") + if ( + self.spec.latency_anchor_artifact is not None + and self.spec.latency_anchor_artifact not in external_artifacts + ): + raise ValueError(f"Latency anchor artifact {self.spec.latency_anchor_artifact!r} is not an external input") stage_positions = {stage.stage_id: index for index, stage in enumerate(self.spec.stages)} roots = deque(stage.stage_id for stage in self.spec.stages if indegree[stage.stage_id] == 0) diff --git a/telefuser/pipelines/lingbot_world_fast/service.py b/telefuser/pipelines/lingbot_world_fast/service.py index 85d301c..8bb7dac 100644 --- a/telefuser/pipelines/lingbot_world_fast/service.py +++ b/telefuser/pipelines/lingbot_world_fast/service.py @@ -895,9 +895,6 @@ def _run_actor_worker_loop( emit_status: Callable[..., None], ) -> None: """Drive dynamic control ingress and ordered output through the shared actor graph.""" - first_item = self._next_realtime_control(state, control_context, control_builder, 0, emit_status, block=True) - if first_item is None: - return runtime_measurement = self._start_benchmark_measurement(state) try: runtime = self.pipeline._create_initialized_session(state.config, progress_callback=emit_status) @@ -921,6 +918,9 @@ def _run_actor_worker_loop( runtime=self._runtime_metadata(runtime), **({"measurement": {"name": "runtime_creation", **runtime_facts}} if runtime_facts is not None else {}), ) + first_item = self._next_realtime_control(state, control_context, control_builder, 0, emit_status, block=True) + if first_item is None: + return submitted = 0 controls_by_chunk: dict[int, list[str] | None] = {} diff --git a/telefuser/pipelines/lingbot_world_fast/streaming.py b/telefuser/pipelines/lingbot_world_fast/streaming.py index b7b5587..59ae6d5 100644 --- a/telefuser/pipelines/lingbot_world_fast/streaming.py +++ b/telefuser/pipelines/lingbot_world_fast/streaming.py @@ -33,6 +33,9 @@ from .pipeline import LingBotWorldFastPipeline +_CONDITION_PREFETCH_DEPTH = 2 + + @dataclass(frozen=True) class LingBotWorldFastStreamingSession: """Lightweight identity for one session in the shared streaming runtime.""" @@ -47,6 +50,8 @@ class _LingBotStreamingSessionEntry: runtime: LingBotWorldFastGenerationSession epoch: int progress_callback: Callable[..., None] | None + next_condition_index: int = 0 + next_control_index: int = 0 class LingBotWorldFastStreamingRuntime: @@ -90,6 +95,7 @@ def __init__(self, pipeline: LingBotWorldFastPipeline) -> None: ), output_artifacts=frozenset({"frames"}), output_capacity_per_session=2, + latency_anchor_artifact="control", ) self.orchestrator = StreamingPipelineOrchestrator(spec, actors) @@ -108,16 +114,28 @@ def create_session( if session_id in self._sessions: raise ValueError(f"LingBot streaming session {session_id!r} already exists") epoch = self.orchestrator.create_session(session_id, final_sequence_id=runtime.chunk_count - 1) - self._sessions[session_id] = _LingBotStreamingSessionEntry(runtime, epoch, progress_callback) - return LingBotWorldFastStreamingSession(session_id, epoch, runtime.cache_handle) + entry = _LingBotStreamingSessionEntry(runtime, epoch, progress_callback) + self._sessions[session_id] = entry + session = LingBotWorldFastStreamingSession(session_id, epoch, runtime.cache_handle) + try: + self._prefetch_conditions(session_id, entry) + except BaseException: + self.close_session(session) + raise + return session def can_submit_chunk(self, session: LingBotWorldFastStreamingSession) -> bool: - """Return whether the session can atomically admit another chunk.""" - self._require_session(session) + """Return whether the session can admit its next control chunk.""" + entry = self._require_session(session) try: - if self.orchestrator.status(session.session_id) != StreamingSessionStatus.RUNNING: - return False - return self.orchestrator.can_push_inputs(session.session_id, ("encode_request", "control")) + with self._lock: + if self.orchestrator.status(session.session_id) != StreamingSessionStatus.RUNNING: + return False + if entry.next_control_index >= entry.runtime.chunk_count: + return False + if not self._ensure_next_condition_locked(session.session_id, entry): + return False + return self.orchestrator.can_push_input(session.session_id, "control") except RuntimeError: if self.orchestrator.error(session.session_id) is not None: return False @@ -139,25 +157,88 @@ def try_submit_chunk( chunk_index: int, control: torch.Tensor, ) -> bool: - """Atomically submit one chunk to the shared actor graph.""" + """Submit the next control while condition encoding runs ahead independently.""" entry = self._require_session(session) if chunk_index < 0 or chunk_index >= entry.runtime.chunk_count: raise ValueError("chunk_index exceeds the LingBot session length") try: - return self.orchestrator.try_push_inputs( - session.session_id, - chunk_index, - { - "encode_request": None, - "control": control, - }, - ) + with self._lock: + if chunk_index != entry.next_control_index: + raise ValueError(f"Expected LingBot control chunk {entry.next_control_index}, got {chunk_index}") + if not self._ensure_next_condition_locked(session.session_id, entry): + return False + accepted = self.orchestrator.try_push_inputs( + session.session_id, + chunk_index, + {"control": control}, + ) + if accepted: + entry.next_control_index += 1 + return accepted except RuntimeError: error = self.orchestrator.error(session.session_id) if error is not None: raise RuntimeError("LingBot streaming scheduler failed") from error raise + def _prefetch_conditions( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> None: + with self._lock: + self._prefetch_conditions_locked(session_id, entry) + + def _prefetch_conditions_locked( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> None: + """Keep at most two condition chunks ahead of accepted controls.""" + while ( + entry.next_condition_index < entry.runtime.chunk_count + and entry.next_condition_index - entry.next_control_index < _CONDITION_PREFETCH_DEPTH + ): + if not self.orchestrator.try_push_inputs( + session_id, + entry.next_condition_index, + {"encode_request": None}, + ): + return + entry.next_condition_index += 1 + + def _ensure_next_condition_locked( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> bool: + """Ensure the next control has a matching condition without filling the lookahead window.""" + if entry.next_condition_index > entry.next_control_index: + return True + if entry.next_condition_index >= entry.runtime.chunk_count: + return False + if not self.orchestrator.try_push_inputs( + session_id, + entry.next_condition_index, + {"encode_request": None}, + ): + return False + entry.next_condition_index += 1 + return True + + def _refill_conditions_after_denoise( + self, + session_id: str, + entry: _LingBotStreamingSessionEntry, + ) -> None: + """Overlap future condition encoding with decode instead of current-chunk DiT work.""" + with self._lock: + if self._sessions.get(session_id) is not entry: + return + if self.orchestrator.status(session_id) != StreamingSessionStatus.RUNNING: + return + self._prefetch_conditions_locked(session_id, entry) + def poll_frames(self, session: LingBotWorldFastStreamingSession) -> list[tuple[int, list[Image.Image]]]: """Return decoded frame batches in chunk order.""" self._require_session(session) @@ -365,6 +446,7 @@ def _denoise_kwargs(self, invocation: StreamingStageInvocation) -> dict[str, obj "max_attention_size": runtime.max_attention_size, } + @torch.inference_mode() def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: entry = self._entry_for_invocation(invocation) runtime = entry.runtime @@ -388,6 +470,7 @@ def _denoise(self, invocation: StreamingStageInvocation) -> dict[str, object]: runtime.world_kv_binding.on_chunk_finalized(runtime, index, latent) except Exception as exc: logger.warning(f"world_kv on_chunk_finalized failed at chunk {index}: {exc}") + self._refill_conditions_after_denoise(invocation.key.session_id, entry) return {"latent": latent} def _decode_inputs(self, invocation: StreamingStageInvocation) -> tuple[tuple[object, ...], dict[str, object]]: diff --git a/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py b/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py index 4163cb1..9d0d233 100644 --- a/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py +++ b/tests/unit/orchestrator/test_streaming_pipeline_orchestrator.py @@ -45,6 +45,7 @@ def _wait_for_outputs( def _orchestrator( encode_calls: list[tuple[int, bool, bool]], denoise_calls: list[tuple[int, bool, bool]], + latency_anchor_artifact: str | None = None, ) -> StreamingPipelineOrchestrator: def encode(invocation): encode_calls.append((invocation.key.sequence_id, invocation.is_first, invocation.is_last)) @@ -65,6 +66,7 @@ def denoise(invocation): StreamingEdgeSpec("control", "denoise"), ), output_artifacts=frozenset({"frames"}), + latency_anchor_artifact=latency_anchor_artifact, ) return StreamingPipelineOrchestrator( spec, @@ -127,6 +129,25 @@ def test_session_metrics_report_end_to_end_percentiles_after_output_polling() -> orchestrator.close() +def test_session_metrics_use_the_configured_ingress_latency_anchor() -> None: + orchestrator = _orchestrator([], [], latency_anchor_artifact="control") + try: + orchestrator.create_session("session", final_sequence_id=0) + orchestrator.push_input("session", 0, "image", "prefetched") + assert orchestrator.wait_until_idle("session") + assert orchestrator.session_metrics("session").ingress_accepted_at == () + + control_submitted_after = time.monotonic() + orchestrator.push_input("session", 0, "control", "left") + assert _wait_for_outputs(orchestrator, "session", 1) + metrics = orchestrator.session_metrics("session") + assert len(metrics.ingress_accepted_at) == 1 + assert metrics.ingress_accepted_at[0][1] >= control_submitted_after + assert metrics.control_to_output_latency.count == 1 + finally: + orchestrator.close() + + def test_bounded_condition_edge_applies_backpressure_until_downstream_consumes() -> None: encode_calls: list[tuple[int, bool, bool]] = [] denoise_calls: list[tuple[int, bool, bool]] = [] diff --git a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py index bd1365f..b1e9659 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py @@ -157,6 +157,29 @@ def initialize(*_args: object, **_kwargs: object) -> LingBotWorldFastGenerationS assert state.generation_session is runtime +def test_actor_worker_prefetches_conditions_before_waiting_for_first_control() -> None: + pipeline = MagicMock() + runtime = LingBotWorldFastGenerationSession(config=_state().config, latent_f=1, chunk_size=1, cache_handle=7) + pipeline._create_initialized_session.return_value = runtime + streaming_runtime = MagicMock() + streaming_runtime.create_session.return_value = SimpleNamespace(session_id="actor-session") + pipeline._get_streaming_runtime.return_value = streaming_runtime + service = LingBotWorldFastService(pipeline) + state = _state() + + def stop_after_prefetch(*_args: object, **_kwargs: object) -> None: + assert pipeline._create_initialized_session.called + assert streaming_runtime.create_session.called + state.active = False + return None + + with patch.object(service, "_next_realtime_control", side_effect=stop_after_prefetch): + service._run_actor_worker_loop(state, MagicMock(), MagicMock(), MagicMock()) + + streaming_runtime.create_session.assert_called_once() + streaming_runtime.try_submit_chunk.assert_not_called() + + def test_actor_worker_prefetches_directional_chunks_within_ingress_capacity() -> None: pipeline = MagicMock() runtime = LingBotWorldFastGenerationSession( diff --git a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py index 2e81a34..7cb92b2 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_streaming.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_streaming.py @@ -52,12 +52,22 @@ class _Denoise: def __init__(self, release_order: list[str]): self.release_order = release_order self.advance_calls: list[int] = [] + self.calls: list[dict[str, object]] = [] + self.inference_modes: list[bool] = [] self.release_calls: list[int] = [] self.fail_denoise = False + self.started: threading.Event | None = None + self.release: threading.Event | None = None def denoise_and_update_cache(self, **kwargs): + self.inference_modes.append(torch.is_inference_mode_enabled()) if self.fail_denoise: raise RuntimeError("injected denoise failure") + self.calls.append(kwargs) + if self.started is not None: + self.started.set() + if self.release is not None: + self.release.wait(timeout=1) return kwargs["condition_chunk"] def advance_noise(self, cache_handle: int): @@ -156,11 +166,82 @@ def test_streaming_session_routes_one_chunk_through_three_stages() -> None: assert pipeline.vae_encode_worker.release_calls == [9] assert pipeline.vae_decode_worker.release_calls == [9] assert pipeline.denoise_stage.release_calls == [9] + assert pipeline.denoise_stage.inference_modes == [True] assert pipeline.release_order == ["decode", "denoise", "encode"] assert pipeline.released == [] assert runtime.cache_handle is None +def test_streaming_session_prefetches_two_conditions_ahead_of_controls() -> None: + pipeline = _Pipeline() + runtime = LingBotWorldFastGenerationSession( + config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))), + prompt_emb=torch.tensor([0.0]), + latent_h=1, + latent_w=1, + latent_f=4, + height=8, + width=8, + frame_tokens=1, + chunk_size=1, + max_attention_size=1, + cache_handle=15, + ) + streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) + session = streaming_runtime.create_session(runtime) + denoise_started = threading.Event() + release_denoise = threading.Event() + pipeline.denoise_stage.started = denoise_started + pipeline.denoise_stage.release = release_denoise + try: + assert streaming_runtime.wait_until_idle(session) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1] + assert pipeline.denoise_stage.calls == [] + + streaming_runtime.submit_chunk(session, 0, torch.tensor([4.0])) + assert denoise_started.wait(timeout=1) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1] + release_denoise.set() + assert streaming_runtime.wait_until_idle(session) + assert [call["chunk_index"] for call in pipeline.vae_encode_worker.calls] == [0, 1, 2] + assert len(pipeline.denoise_stage.calls) == 1 + assert streaming_runtime.poll_frames(session)[0][0] == 0 + metrics = streaming_runtime.session_metrics(session) + assert [index for index, _ in metrics.ingress_accepted_at] == [0] + finally: + release_denoise.set() + streaming_runtime.close_session(session) + streaming_runtime.close() + + +def test_streaming_session_rejects_out_of_order_or_duplicate_controls() -> None: + pipeline = _Pipeline() + runtime = LingBotWorldFastGenerationSession( + config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))), + prompt_emb=torch.tensor([0.0]), + latent_h=1, + latent_w=1, + latent_f=2, + height=8, + width=8, + frame_tokens=1, + chunk_size=1, + max_attention_size=1, + cache_handle=16, + ) + streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) + session = streaming_runtime.create_session(runtime) + try: + with pytest.raises(ValueError, match="Expected LingBot control chunk 0, got 1"): + streaming_runtime.try_submit_chunk(session, 1, torch.tensor([4.0])) + assert streaming_runtime.try_submit_chunk(session, 0, torch.tensor([4.0])) + with pytest.raises(ValueError, match="Expected LingBot control chunk 1, got 0"): + streaming_runtime.try_submit_chunk(session, 0, torch.tensor([4.0])) + finally: + streaming_runtime.close_session(session) + streaming_runtime.close() + + def test_streaming_runtime_shares_one_actor_graph_across_sessions() -> None: pipeline = _Pipeline() streaming_runtime = LingBotWorldFastStreamingRuntime(pipeline) diff --git a/tests/unit/test_run_examples.py b/tests/unit/test_run_examples.py new file mode 100644 index 0000000..8c49134 --- /dev/null +++ b/tests/unit/test_run_examples.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import pytest + +from examples.run_examples import _close_pipeline + + +class _ClosablePipeline: + def __init__(self, *, fail: bool = False) -> None: + self.fail = fail + self.close_calls = 0 + + def close(self) -> None: + self.close_calls += 1 + if self.fail: + raise RuntimeError("close failed") + + +def test_close_pipeline_releases_owned_workers() -> None: + pipeline = _ClosablePipeline() + + _close_pipeline(pipeline) + + assert pipeline.close_calls == 1 + + +def test_close_pipeline_does_not_mask_regression_result(capsys: pytest.CaptureFixture[str]) -> None: + pipeline = _ClosablePipeline(fail=True) + + _close_pipeline(pipeline) + + assert pipeline.close_calls == 1 + assert "Warning: failed to close pipeline: close failed" in capsys.readouterr().err