diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 615e4067e..3b3f8e77c 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -242,6 +242,18 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable logging of diffusion pipeline stats.", ) + parser.add_argument( + "--init-timeout", + type=int, + default=600, + help="Timeout for initializing a single stage in seconds (default: 600s)", + ) + parser.add_argument( + "--stage-init-timeout", + type=int, + default=600, + help="Timeout for initializing a single stage in seconds (default: 600s)", + ) parser.add_argument( "--use-system-prompt", type=str, @@ -346,6 +358,8 @@ def main(): "mode": "text-to-image", "log_stats": args.log_stats, "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler, + "init_timeout": args.init_timeout, + "stage_init_timeout": args.stage_init_timeout, **lora_args, **quant_kwargs, } diff --git a/tests/conftest.py b/tests/conftest.py index 27833fe28..18a0ee57d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -75,6 +75,7 @@ class OmniServerParams(NamedTuple): use_omni: bool = True use_stage_cli: bool = False init_timeout: int | None = None + stage_init_timeout: int | None = None # None defers to the server's own default (300 s) def assert_image_diffusion_response( @@ -1768,8 +1769,8 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st ) server_args = params.server_args or [] - if params.use_omni: - server_args = ["--stage-init-timeout", "120", *server_args] + if params.use_omni and params.stage_init_timeout is not None: + server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)] if params.init_timeout is not None: server_args = [*server_args, "--init-timeout", str(params.init_timeout)] if params.use_stage_cli: @@ -3257,7 +3258,7 @@ def omni_runner(request, model_prefix): with _omni_server_lock: model, stage_config_path = request.param model = model_prefix + model - with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner: + with OmniRunner(model, seed=42, stage_configs_path=stage_config_path) as runner: print("OmniRunner started successfully") yield runner print("OmniRunner stopping...") diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index c625239e5..c566c2e0a 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -23,6 +23,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" CONFIG_FILE_PATH = str(Path(__file__).parent.parent / "tests" / "test.json") BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +STAGE_INIT_TIMEOUT = 600 STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs" @@ -43,7 +44,7 @@ def omni_server(request): print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args = ["--stage-init-timeout", "120", "--init-timeout", "900"] + server_args = ["--stage-init-timeout", str(STAGE_INIT_TIMEOUT), "--init-timeout", "900"] if stage_config_path: server_args = ["--stage-configs-path", stage_config_path] + server_args with OmniServer(model, server_args) as server: diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py index 0a81b0207..062750b3c 100644 --- a/tests/e2e/accuracy/conftest.py +++ b/tests/e2e/accuracy/conftest.py @@ -114,8 +114,8 @@ class AccuracyServerConfig: params = self.generate_params model = self.model_prefix + params.model server_args = params.server_args or [] - if params.use_omni: - server_args = ["--stage-init-timeout", "120", *server_args] + if params.use_omni and params.stage_init_timeout is not None: + server_args = ["--stage-init-timeout", str(params.stage_init_timeout), *server_args] with OmniServer( model, server_args, @@ -226,6 +226,7 @@ def _build_accuracy_server_config( server_args=generate_server_args, env_dict={"CUDA_VISIBLE_DEVICES": shared_gpu}, use_omni=True, + stage_init_timeout=300, ), judge_params=OmniServerParams( model=judge_model, diff --git a/tests/e2e/online_serving/test_bagel_online.py b/tests/e2e/online_serving/test_bagel_online.py index ca24f5f81..a3f999f13 100644 --- a/tests/e2e/online_serving/test_bagel_online.py +++ b/tests/e2e/online_serving/test_bagel_online.py @@ -47,7 +47,7 @@ test_params = [ OmniServerParams( model=MODEL, stage_config_path=STAGE_CONFIGS_PATH, - server_args=["--stage-init-timeout", "300"], + stage_init_timeout=300, ), ] diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py index 4648c424f..39b6dc8e2 100644 --- a/tests/e2e/online_serving/test_dynin_omni_expansion.py +++ b/tests/e2e/online_serving/test_dynin_omni_expansion.py @@ -30,7 +30,7 @@ T2I_PROMPT = "A high quality detailed living room interior photo." T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving." I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors." -TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG)] +TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, stage_init_timeout=600)] _STAGE_COUNT = 3 _I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False} diff --git a/tests/e2e/online_serving/test_qwen3_tts_batch.py b/tests/e2e/online_serving/test_qwen3_tts_batch.py index d0d633661..1a453afb7 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_batch.py +++ b/tests/e2e/online_serving/test_qwen3_tts_batch.py @@ -30,6 +30,7 @@ from tests.conftest import ( from tests.utils import hardware_test MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" +STAGE_INIT_TIMEOUT_S = 120 def get_stage_config(name: str = "qwen3_tts.yaml"): @@ -47,7 +48,7 @@ def omni_server(): "--stage-configs-path", stage_config_path, "--stage-init-timeout", - "120", + str(STAGE_INIT_TIMEOUT_S), "--trust-remote-code", "--enforce-eager", "--disable-log-stats", @@ -337,7 +338,7 @@ def omni_server_batch2(): "--stage-configs-path", config_path, "--stage-init-timeout", - "120", + str(STAGE_INIT_TIMEOUT_S), "--trust-remote-code", "--enforce-eager", "--disable-log-stats", diff --git a/tests/e2e/online_serving/test_qwen3_tts_speaker_embedding.py b/tests/e2e/online_serving/test_qwen3_tts_speaker_embedding.py index 64e13e155..8c1c86081 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_speaker_embedding.py +++ b/tests/e2e/online_serving/test_qwen3_tts_speaker_embedding.py @@ -23,6 +23,7 @@ from tests.utils import hardware_test MODEL_BASE = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" MODEL_BASE_1_7B = "Qwen/Qwen3-TTS-12Hz-1.7B-Base" +STAGE_INIT_TIMEOUT_S = 120 # A synthetic 1024-dim speaker embedding (all 0.1 — not a real voice, but # exercises the full code path through the talker's _build_prompt_embeds). @@ -47,7 +48,7 @@ def _server_args(): "--stage-configs-path", get_stage_config(), "--stage-init-timeout", - "120", + str(STAGE_INIT_TIMEOUT_S), "--trust-remote-code", "--enforce-eager", "--disable-log-stats", diff --git a/tests/e2e/online_serving/test_qwen3_tts_websocket.py b/tests/e2e/online_serving/test_qwen3_tts_websocket.py index df0514601..849d1c115 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_websocket.py +++ b/tests/e2e/online_serving/test_qwen3_tts_websocket.py @@ -19,6 +19,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" +STAGE_INIT_TIMEOUT_S = 120 def get_stage_config() -> str: @@ -37,7 +38,7 @@ def omni_server(): "--stage-configs-path", stage_config_path, "--stage-init-timeout", - "120", + str(STAGE_INIT_TIMEOUT_S), "--trust-remote-code", "--enforce-eager", "--disable-log-stats", diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 31d3ed775..002e8226f 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1,5 +1,6 @@ import importlib import os +import threading import types import pytest @@ -86,6 +87,144 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke os.environ[env_var] = old_env +def test_initialize_stages_passes_stage_init_timeout_to_diffusion_handshake(monkeypatch): + """Regression test for stage_init_timeout passing to complete_diffusion_handshake + in the diffusion stage path. + """ + import vllm_omni.diffusion.data as diffusion_data_mod + import vllm_omni.diffusion.stage_diffusion_client as client_mod + import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.platforms import current_omni_platform + + engine = object.__new__(AsyncOmniEngine) + engine.model = "dummy-model" + engine.config_path = "dummy-config" + engine.num_stages = 1 + engine.async_chunk = False + engine.diffusion_batch_size = 1 + engine.single_stage_mode = False + engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion", engine_args={})] + + metadata = types.SimpleNamespace( + stage_id=0, + stage_type="diffusion", + runtime_cfg={"devices": "0"}, + prompt_expand_func=None, + final_output=True, + final_output_type="image", + default_sampling_params=None, + custom_process_input_func=None, + engine_input_source=None, + cfg_kv_collect_func=None, + ) + + captured_timeout = None + device_env_var = current_omni_platform.device_control_env_var + prev_device_env = os.environ.get(device_env_var) + os.environ[device_env_var] = "0" + + monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) + monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) + monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata) + monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) + monkeypatch.setattr( + engine_mod, + "finalize_initialized_stages", + lambda stage_clients, _input_processor: ( + stage_clients, + [types.SimpleNamespace()], + [{"final_output_type": "image"}], + ), + ) + monkeypatch.setattr( + diffusion_data_mod.OmniDiffusionConfig, + "from_kwargs", + classmethod(lambda cls, **kwargs: types.SimpleNamespace(parallel_config=types.SimpleNamespace(world_size=1))), + ) + monkeypatch.setattr( + client_mod, + "spawn_diffusion_proc", + lambda model, od_cfg: (object(), "ipc://handshake", "ipc://request", "ipc://response"), + ) + + def _capture_handshake_timeout(proc, handshake_address, handshake_timeout): + nonlocal captured_timeout + captured_timeout = handshake_timeout + + monkeypatch.setattr(client_mod, "complete_diffusion_handshake", _capture_handshake_timeout) + monkeypatch.setattr( + client_mod.zmq, + "Context", + lambda: types.SimpleNamespace(socket=lambda _: types.SimpleNamespace(connect=lambda _: None)), + ) + + try: + engine._initialize_stages(stage_init_timeout=302) + finally: + if prev_device_env is None: + os.environ.pop(device_env_var, None) + else: + os.environ[device_env_var] = prev_device_env + + assert captured_timeout == 302 + + +def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch): + """Regression test for stage_init_timeout reaching complete_stage_handshake + in the LLM stage path. + """ + import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.platforms import current_omni_platform + + engine = object.__new__(AsyncOmniEngine) + engine.model = "dummy-model" + engine.single_stage_mode = False + engine._omni_master_server = None + + metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}) + fake_vllm_config = types.SimpleNamespace() + fake_addresses = types.SimpleNamespace() + fake_proc = types.SimpleNamespace() + + captured_timeout = None + + device_env_var = current_omni_platform.device_control_env_var + prev_device_env = os.environ.get(device_env_var) + os.environ[device_env_var] = "0" + + monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) + monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object)) + monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: []) + monkeypatch.setattr( + engine_mod, + "spawn_stage_core", + lambda **_: (fake_addresses, fake_proc, "ipc://handshake"), + ) + + def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout): + nonlocal captured_timeout + captured_timeout = handshake_timeout + + monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout) + + try: + engine._launch_llm_stage( + stage_cfg=types.SimpleNamespace(engine_args={}), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=302, + llm_stage_launch_lock=threading.Lock(), + ) + finally: + if prev_device_env is None: + os.environ.pop(device_env_var, None) + else: + os.environ[device_env_var] = prev_device_env + + assert captured_timeout == 302 + + def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): """Regression test for GLM-Image t2i preprocessing path. diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 627a98395..2c5bf6cc7 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -1459,6 +1459,7 @@ class TestLaunchLlmStageSingleStageMode: fake_proc = Mock() fake_handshake_address = "ipc:///tmp/fake-handshake" + stage_init_timeout = 60 with ( patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"), @@ -1484,7 +1485,7 @@ class TestLaunchLlmStageSingleStageMode: stage_cfg=_make_stage_cfg(0), metadata=metadata, stage_connector_spec={}, - stage_init_timeout=60, + stage_init_timeout=stage_init_timeout, llm_stage_launch_lock=threading.Lock(), ) @@ -1498,6 +1499,7 @@ class TestLaunchLlmStageSingleStageMode: fake_handshake_address, fake_addresses, fake_vllm_config, + stage_init_timeout, ) mock_omni.assert_not_called() assert isinstance(result, StartedLlmStage) diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index cd7159b68..7e740dc89 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -50,11 +50,12 @@ class StageDiffusionClient: model: str, od_config: OmniDiffusionConfig, metadata: StageMetadata, + stage_init_timeout: int, batch_size: int = 1, ) -> None: # Spawn StageDiffusionProc subprocess and wait for READY. proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config) - complete_diffusion_handshake(proc, handshake_address) + complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) self._initialize_client(metadata, request_address, response_address, proc=proc, batch_size=batch_size) @classmethod diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 2bba41925..cef697630 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -39,8 +39,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -_HANDSHAKE_POLL_TIMEOUT_S = 600 - class StageDiffusionProc: """Subprocess entry point for diffusion inference. @@ -619,13 +617,14 @@ def spawn_diffusion_proc( def complete_diffusion_handshake( proc: BaseProcess, handshake_address: str, + handshake_timeout: int, ) -> None: """Wait for the diffusion subprocess to signal READY. On failure the process is terminated before re-raising. """ try: - _perform_diffusion_handshake(proc, handshake_address) + _perform_diffusion_handshake(proc, handshake_address, handshake_timeout) except Exception: shutdown([proc]) raise @@ -634,6 +633,7 @@ def complete_diffusion_handshake( def _perform_diffusion_handshake( proc: BaseProcess, handshake_address: str, + handshake_timeout: int, ) -> None: """Run the handshake with the diffusion subprocess.""" with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: @@ -641,11 +641,15 @@ def _perform_diffusion_handshake( poller.register(handshake_socket, zmq.POLLIN) poller.register(proc.sentinel, zmq.POLLIN) - timeout_ms = _HANDSHAKE_POLL_TIMEOUT_S * 1000 + timeout_ms = handshake_timeout * 1000 while True: events = dict(poller.poll(timeout=timeout_ms)) if not events: - raise TimeoutError("Timed out waiting for READY from StageDiffusionProc") + raise TimeoutError( + f"Timed out waiting for READY from StageDiffusionProc after {handshake_timeout}s. " + f"This typically indicates model loading or warmup is taking too long. " + f"Consider increasing `stage_init_timeout` for large models." + ) if handshake_socket in events: identity, raw = handshake_socket.recv_multipart() msg = msgspec.msgpack.decode(raw) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 7dc5db0ac..1e92780b6 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -430,7 +430,9 @@ class AsyncOmniEngine: else: assert proc is not None assert handshake_address is not None - complete_stage_handshake(proc, handshake_address, addresses, vllm_config) + complete_stage_handshake( + proc, handshake_address, addresses, vllm_config, stage_init_timeout + ) logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id) finally: if previous_visible_devices is None: @@ -759,6 +761,7 @@ class AsyncOmniEngine: self.model, stage_cfg, metadata, + stage_init_timeout=stage_init_timeout, batch_size=self.diffusion_batch_size, ) logger.info( diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 05d8f107c..689378a79 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -37,8 +37,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -_HANDSHAKE_POLL_TIMEOUT_S = 600 - class StageEngineCoreProc(EngineCoreProc): """Stage-specific engine core process for vLLM-Omni. @@ -145,13 +143,14 @@ def complete_stage_handshake( handshake_address: str, addresses: EngineZmqAddresses, vllm_config: VllmConfig, + handshake_timeout: int, ) -> None: """Perform the HELLO/INIT/READY handshake with an already-spawned proc. On failure the process is terminated before re-raising. """ try: - _perform_handshake(proc, handshake_address, addresses, vllm_config) + _perform_handshake(proc, handshake_address, addresses, vllm_config, handshake_timeout) except Exception: shutdown([proc]) raise @@ -162,6 +161,7 @@ def _perform_handshake( handshake_address: str, addresses: EngineZmqAddresses, vllm_config: VllmConfig, + handshake_timeout: int, ) -> None: """Run the HELLO / INIT / READY handshake with the subprocess.""" with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: @@ -169,7 +169,7 @@ def _perform_handshake( poller.register(handshake_socket, zmq.POLLIN) poller.register(proc.sentinel, zmq.POLLIN) - identity, msg = _recv(poller, handshake_socket, proc, "HELLO") + identity, msg = _recv(poller, handshake_socket, proc, "HELLO", handshake_timeout) if msg.get("status") != "HELLO": raise RuntimeError(f"Expected HELLO, got: {msg}") @@ -179,7 +179,7 @@ def _perform_handshake( ) handshake_socket.send_multipart([identity, msgspec.msgpack.encode(init_payload)]) - identity, msg = _recv(poller, handshake_socket, proc, "READY") + identity, msg = _recv(poller, handshake_socket, proc, "READY", handshake_timeout) if msg.get("status") != "READY": raise RuntimeError(f"Expected READY, got: {msg}") num_gpu_blocks = msg.get("num_gpu_blocks") @@ -192,13 +192,18 @@ def _recv( handshake_socket: zmq.Socket, proc: BaseProcess, expected: str, + timeout_s: int = 600, ) -> tuple[bytes, dict]: """Wait for one handshake message; raise if the process dies first.""" - timeout_ms = _HANDSHAKE_POLL_TIMEOUT_S * 1000 + timeout_ms = timeout_s * 1000 while True: events = dict(poller.poll(timeout=timeout_ms)) if not events: - raise TimeoutError(f"Timed out waiting for {expected} from StageEngineCoreProc") + raise TimeoutError( + f"Timed out waiting for {expected} from StageEngineCoreProc after {timeout_s}s. " + f"This typically indicates model loading or initialization is taking too long. " + f"Consider increasing `stage_init_timeout` for large models." + ) if handshake_socket in events: identity, raw = handshake_socket.recv_multipart() return identity, msgspec.msgpack.decode(raw) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 272df14f8..158b4c547 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -335,7 +335,7 @@ def build_vllm_config( def acquire_device_locks( stage_id: int, engine_args_dict: dict[str, Any], - stage_init_timeout: int = 300, + stage_init_timeout: int, ) -> list[int]: """Acquire exclusive file locks on devices needed by this stage. @@ -528,6 +528,7 @@ def initialize_diffusion_stage( model: str, stage_cfg: Any, metadata: StageMetadata, + stage_init_timeout: int, batch_size: int = 1, ) -> Any: """Build a diffusion stage client. @@ -536,6 +537,7 @@ def initialize_diffusion_stage( model: Model name or path. stage_cfg: Stage configuration. metadata: Extracted stage metadata. + stage_init_timeout: Timeout in seconds for stage initialization handshake batch_size: Maximum number of requests to batch together in the diffusion engine. Passed through to ``StageDiffusionClient`` and ultimately to ``AsyncOmni``. @@ -543,7 +545,9 @@ def initialize_diffusion_stage( from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient od_config = build_diffusion_config(model, stage_cfg, metadata) - return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size) + return StageDiffusionClient( + model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size + ) def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None: