[RFC]: diffusion engine clean up

Motivation.

Optimize DiffusionEngine concurrency, telemetry, and CPU utilization

Description: The DiffusionEngine currently has a few architectural bottlenecks that severely limit concurrency, waste host CPU cycles, and report inaccurate telemetry. This issue tracks the refactoring needed to make the engine truly asynchronous, improve memory handling during the warmup phase, and fix bugs in our metric reporting.

Problems Identified:

1. CPU Starvation (Busy-Wait): The add_req_and_wait_for_response method uses an unthrottled while True: loop to poll the scheduler, which pins a CPU core at 100% and starves other threads.
2. Sequential Bottleneck: _rpc_lock wraps the entire request lifecycle (scheduling + execution). This prevents any other requests from being queued while the GPU is processing a task, defeating the purpose of the scheduler.
3. Telemetry Bugs: The metrics dictionary in step() swaps the names for execution time and total time, duplicates the preprocess_time_ms key, and mixes telemetry times with request parameters.
4. Redundant Computation: supports_audio_output is evaluated inside a loop for every prompt, and _dummy_run expensively allocates and slices random numpy arrays instead of generating precisely what is needed.
5. Readability: Deeply nested audio-slicing logic in the multiple-request handling block makes the step() method difficult to maintain.

Proposed Change.

• Fix Concurrency & Locking:
  • Introduce a sleep (e.g., time.sleep(0.001)) or an asyncio.Event / threading.Condition in the add_req_and_wait_for_response polling loop.
  • Scope _rpc_lock to cover only the scheduler state updates, releasing it before calling self.executor.add_req(req).
• Fix Telemetry & Metrics:
  • Correct the logical bug where diffusion_engine_exec_time_ms and diffusion_engine_total_time_ms are swapped.
  • Standardize metric keys (e.g., prefix times with time_ and parameters with param_).
  • Remove the redundant preprocessing_time_ms append block.
• Optimize Performance & Memory:
  • Evaluate supports_audio_output(self.od_config.model_class_name) once outside the prompt loop.
  • Change _dummy_run to generate an exact-sized array (e.g., np.random.randn(audio_sr * 2)) or use np.zeros if the model allows it.
  • Evaluate using .to("cpu", non_blocking=True) instead of .cpu() for offloading, if post-processing allows.
• Refactor Code Quality:
  • Extract the complex audio slicing logic in step() into a private helper function (e.g., _extract_audio_slice).

Expected Impact:

• Vastly improved CPU utilization (no more busy-waiting).
• Ability to queue multiple requests asynchronously.
• Accurate, dashboard-ready telemetry.
• Cleaner, more maintainable code in the core engine step() function.

Feedback Period.

No response

CC List.

@ZJY0516 @SamitHuang @wtomin @asukaqaq-s

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.