sgl-project--sglang/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
import gc
import logging
import multiprocessing as mp
import os
import tempfile
import time
from contextlib import ExitStack
from dataclasses import dataclass, field
from typing import Any, Callable, List, Union

import numpy as np
import torch
from setproctitle import setproctitle

from sglang.multimodal_gen import envs
from sglang.multimodal_gen.runtime.distributed import (
    get_sp_group,
    get_tp_rank,
    get_tp_world_size,
    maybe_init_distributed_environment_and_model_parallel,
    model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
    get_cfg_group,
    get_classifier_free_guidance_rank,
    get_classifier_free_guidance_world_size,
    get_ring_parallel_rank,
    get_ring_parallel_world_size,
    get_tp_group,
    get_ulysses_parallel_rank,
    get_ulysses_parallel_world_size,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
    materialize_output_sample,
    post_process_sample,
    save_outputs,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    configure_layerwise_offload_modules,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.memory_occupation_controller import (
    MemoryOccupationController,
)
from sglang.multimodal_gen.runtime.pipelines_core import (
    ComposedPipelineBase,
    LoRAPipeline,
    Req,
    build_pipeline,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.post_training.gpu_worker_post_training_mixin import (
    GPUWorkerPostTrainingMixin,
)
from sglang.multimodal_gen.runtime.realtime.session import (
    RealtimeSessionCache,
)
from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch
from sglang.multimodal_gen.runtime.utils.logging_utils import (
    configure_logger,
    globally_suppress_loggers,
    init_logger,
)
from sglang.multimodal_gen.runtime.utils.perf_logger import (
    PerformanceLogger,
    capture_memory_snapshot,
)
from sglang.multimodal_gen.runtime.utils.realtime_video import (
    RAW_RGB_CONTENT_TYPE,
    build_raw_rgb_frame_batches,
)
from sglang.multimodal_gen.runtime.utils.trace_wrapper import (
    DiffStage,
    init_diffusion_tracing,
    trace_slice,
)
from sglang.multimodal_gen.utils import kill_itself_when_parent_died
from sglang.srt.utils.network import NetworkAddress

logger = init_logger(__name__)

OFFLOAD_DISABLE_RECOMMENDATION_ORDER = (
    "vae",
    "image_encoder",
    "text_encoder",
    "text_encoder_2",
    "transformer",
)


@dataclass
class _ExpandedOutputParts:
    tensor_outputs: list[torch.Tensor] = field(default_factory=list)
    list_outputs: list[Any] = field(default_factory=list)
    tensor_audio: list[torch.Tensor] = field(default_factory=list)
    trajectory_latents: list[torch.Tensor] = field(default_factory=list)
    noise_preds: list[torch.Tensor] = field(default_factory=list)
    output_file_paths: list[str] = field(default_factory=list)
    metrics_list: list[Any] = field(default_factory=list)
    trajectory_decoded_parts: list[list[torch.Tensor]] | None = None


class GPUWorker(GPUWorkerPostTrainingMixin):
    """
    A worker that executes the model on a single GPU.
    """

    def __init__(
        self,
        local_rank: int,
        rank: int,
        master_port: int,
        server_args: ServerArgs,
    ):
        self.local_rank = local_rank
        self.rank = rank
        self.master_port = master_port
        # FIXME: should we use tcp as distribute init method?
        self.server_args = server_args
        self.pipeline: ComposedPipelineBase = None

        self.init_device_and_model()
        self.sp_group = get_sp_group()
        self.sp_cpu_group = self.sp_group.cpu_group
        self.tp_group = get_tp_group()
        self.tp_cpu_group = self.tp_group.cpu_group

        self.cfg_group = get_cfg_group()
        self.cfg_cpu_group = self.cfg_group.cpu_group
        self._realtime_sessions = RealtimeSessionCache(max_sessions=1)
        self.memory_occupation: MemoryOccupationController | None = None

    def release_realtime_session(self, session_id: str) -> OutputBatch:
        """release the session of a realtime connection"""
        if not session_id:
            return OutputBatch(
                output={
                    "released": False,
                    "session_id": session_id,
                    "reason": "empty_session_id",
                }
            )

        released = self._realtime_sessions.release(session_id)
        if released:
            if torch.cuda.is_initialized():
                torch.cuda.empty_cache()
        return OutputBatch(output={"released": released, "session_id": session_id})

    def _configure_persistent_torch_compile_cache(self) -> None:
        """Persist torch.compile's Inductor/Triton cache across restarts"""
        compile_cache_root = os.path.join(
            envs.SGLANG_DIFFUSION_CACHE_ROOT, "torch_compile_cache"
        )
        tmp_root = tempfile.gettempdir()
        for env_name, sub in (
            ("TORCHINDUCTOR_CACHE_DIR", "inductor"),
            ("TRITON_CACHE_DIR", "triton"),
        ):
            current = os.environ.get(env_name)
            if current and not current.startswith(tmp_root):
                # Respect an explicit, non-ephemeral user-provided cache dir.
                continue
            cache_path = os.path.join(compile_cache_root, sub)
            try:
                os.makedirs(cache_path, exist_ok=True)
            except OSError as e:
                logger.warning(
                    "Could not create torch.compile cache dir %s: %s", cache_path, e
                )
                continue
            os.environ[env_name] = cache_path
        logger.info(
            "torch.compile cache: TORCHINDUCTOR_CACHE_DIR=%s TRITON_CACHE_DIR=%s",
            os.environ.get("TORCHINDUCTOR_CACHE_DIR"),
            os.environ.get("TRITON_CACHE_DIR"),
        )

    def is_sleeping(self) -> bool:
        return self.memory_occupation.is_sleeping() if self.memory_occupation else False

    def _get_memory_occupation(self) -> MemoryOccupationController:
        if self.memory_occupation is None:
            self.memory_occupation = MemoryOccupationController(
                pipeline=self.pipeline,
                rank=self.rank,
                use_fsdp_inference=self.server_args.use_fsdp_inference,
            )
        return self.memory_occupation

    def init_device_and_model(self) -> None:
        """Initialize the device and load the model."""
        torch.get_device_module().set_device(self.local_rank)
        # Set environment variables for distributed initialization
        os.environ["MASTER_ADDR"] = "localhost"
        os.environ["MASTER_PORT"] = str(self.master_port)
        os.environ["LOCAL_RANK"] = str(self.local_rank)
        os.environ["RANK"] = str(self.rank)
        os.environ["WORLD_SIZE"] = str(self.server_args.num_gpus)
        self._configure_persistent_torch_compile_cache()
        # initialize the distributed environment
        maybe_init_distributed_environment_and_model_parallel(
            tp_size=self.server_args.tp_size,
            cfg_degree=self.server_args.cfg_parallel_degree or 1,
            ulysses_degree=self.server_args.ulysses_degree,
            ring_degree=self.server_args.ring_degree,
            sp_size=self.server_args.sp_degree,
            dp_size=self.server_args.dp_size,
            distributed_init_method=NetworkAddress(
                "127.0.0.1", self.master_port
            ).to_tcp(),
            dist_timeout=self.server_args.dist_timeout,
        )

        # set proc title
        if model_parallel_is_initialized():
            suffix = ""
            if get_tp_world_size() != 1:
                tp_rank = get_tp_rank()
                suffix += f"_TP{tp_rank}"
            if get_ulysses_parallel_world_size() != 1:
                u_rank = get_ulysses_parallel_rank()
                suffix += f"_U{u_rank}"
            if get_ring_parallel_world_size() != 1:
                r_rank = get_ring_parallel_rank()
                suffix += f"_R{r_rank}"
            if get_classifier_free_guidance_world_size() != 1:
                c_rank = get_classifier_free_guidance_rank()
                suffix += f"_C{c_rank}"
            setproctitle(f"sgl_diffusion::scheduler{suffix}")
        else:
            setproctitle(f"sgl_diffusion::scheduler_{self.local_rank}")

        self.pipeline = build_pipeline(self.server_args)

        # apply layerwise offload after lora is applied while building LoRAPipeline
        # otherwise empty offloaded weights could fail lora converting
        if self.server_args.layerwise_offload_components:
            configure_layerwise_offload_modules(
                self.pipeline.modules,
                self.server_args,
                component_names=self.server_args.layerwise_offload_components,
                warn_missing=(
                    self.server_args.is_arg_explicitly_set(
                        "layerwise_offload_components"
                    )
                    or self.server_args.is_arg_explicitly_set("dit_layerwise_offload")
                ),
            )

        logger.info(
            f"Worker {self.rank}: Initialized device, model, and distributed environment."
        )

    def do_mem_analysis(self, output_batch: OutputBatch):
        final_snapshot = capture_memory_snapshot()
        if output_batch.metrics:
            output_batch.metrics.record_memory_snapshot("mem_analysis", final_snapshot)

        # for details on max_memory_reserved: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html
        peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
        peak_allocated_bytes = torch.get_device_module().max_memory_allocated()

        output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
        peak_reserved_gb = peak_reserved_bytes / (1024**3)
        peak_allocated_gb = peak_allocated_bytes / (1024**3)

        remaining_gpu_mem_gb = (
            current_platform.get_device_total_memory() / (1024**3) - peak_reserved_gb
        )
        can_stay_resident = self.get_can_stay_resident_components(remaining_gpu_mem_gb)
        suggested_args_str = self._format_offload_disable_suggestions(can_stay_resident)

        pool_overhead_gb = peak_reserved_gb - peak_allocated_gb

        logger.debug(
            f"Peak GPU memory: {peak_reserved_gb:.2f} GB, "
            f"Peak allocated: {peak_allocated_gb:.2f} GB, "
            f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb / peak_reserved_gb * 100:.1f}%), "
            f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
            f"Components that could stay resident (based on the last request workload): {can_stay_resident}. "
            f"Related offload server args to disable: {suggested_args_str}"
        )

    def _format_offload_disable_suggestions(self, components: List[str]) -> str:
        component_set = set(components)
        suggestions = []
        seen_args = set()

        for component in OFFLOAD_DISABLE_RECOMMENDATION_ORDER:
            if component not in component_set:
                continue

            arg = None
            if component == "vae":
                arg = "--vae-cpu-offload"
            elif component == "image_encoder":
                arg = "--image-encoder-cpu-offload"
            elif component in ("text_encoder", "text_encoder_2"):
                arg = "--text-encoder-cpu-offload"
            elif component == "transformer":
                if self.server_args.is_dit_layerwise_offload_selected:
                    arg = "--dit-layerwise-offload"
                elif self.server_args.dit_cpu_offload:
                    arg = "--dit-cpu-offload"

            if arg is not None and arg not in seen_args:
                suggestions.append(arg)
                seen_args.add(arg)

        return ", ".join(suggestions) if suggestions else "None"

    def execute_forward(
        self, batch: List[Req], return_req: bool = False
    ) -> OutputBatch | Req:
        """
        Execute a forward pass.

        Args:
            batch: List of requests to process.
            return_req: If True, return the raw Req instead of OutputBatch.
                Used by disaggregated pipelines to access intermediate tensors.
        """
        assert self.pipeline is not None
        if len(batch) > 1:
            if return_req:
                raise ValueError(
                    "Grouped execute_forward does not support return_req=True"
                )
            # grouped reqs currently come only from expanded num_outputs_per_prompt
            self._validate_group_forward_reqs(batch)
            return self._execute_forward_batch(batch)

        req = batch[0]
        return self._execute_forward_common(
            req,
            forward_fn=lambda: self.pipeline.forward(req, self.server_args),
            log_reqs=[req],
            return_req=return_req,
            save_output_paths=lambda output_batch: self._save_output_paths(
                req, output_batch
            ),
            error_context=f"request {req.request_id}",
        )

    def _execute_forward_batch(self, batch: list[Req]) -> OutputBatch | Req:
        """Execute expanded multi-output requests as one grouped forward."""
        # TODO: support early return or mix-stage execution for reqs in a group
        assert self.pipeline is not None
        req = batch[0]
        return self._execute_forward_common(
            req,
            forward_fn=lambda: self._forward_group(batch),
            log_reqs=batch,
            return_req=False,
            save_output_paths=lambda output_batch: self._save_group_output_paths(
                batch, output_batch
            ),
            error_context=f"grouped request {req.request_id}",
        )

    def _execute_forward_common(
        self,
        req: Req,
        *,
        forward_fn: Callable[[], Req | OutputBatch],
        log_reqs: list[Req],
        return_req: bool,
        save_output_paths: Callable[[OutputBatch], None],
        error_context: str,
    ) -> OutputBatch | Req:
        """
        Args:
            forward_fn: the actual forward function for reqs
        """
        output_batch = None
        try:
            if self.rank == 0 and not current_platform.is_cpu():
                torch.get_device_module().reset_peak_memory_stats()

            start_time = time.monotonic()
            self._realtime_sessions.attach(req)

            # capture memory baseline for each req in grouped forward on rank-0
            request_metrics = [
                item.metrics for item in log_reqs if item.metrics is not None
            ]
            if self.rank == 0 and request_metrics and not current_platform.is_cpu():
                baseline_snapshot = capture_memory_snapshot()
                for metrics in request_metrics:
                    metrics.record_memory_snapshot("before_forward", baseline_snapshot)

            for item in log_reqs:
                item.log(server_args=self.server_args)
            with ExitStack() as stack:
                for item in log_reqs:
                    stack.enter_context(
                        trace_slice(item.trace_ctx, DiffStage.GPU_FORWARD)
                    )
                result = forward_fn()

            # disagg roles return raw Req so callers can keep and transfer intermediate tensors
            # before converting it to OutputBatch
            if return_req and isinstance(result, Req):
                return result

            output_batch = self._to_output_batch(result)
            self._record_output_peak_memory(output_batch)

            output_metrics = self._iter_output_metrics(output_batch)
            if self.rank == 0 and output_metrics and not current_platform.is_cpu():
                peak_snapshot = capture_memory_snapshot()
                for metrics in output_metrics:
                    metrics.record_memory_snapshot("after_forward", peak_snapshot)

            if (
                self.rank == 0
                and not req.suppress_logs
                and not current_platform.is_cpu()
                and logger.isEnabledFor(logging.DEBUG)
            ):
                self.do_mem_analysis(output_batch)

            duration_ms = (time.monotonic() - start_time) * 1000
            for metrics in output_metrics:
                metrics.total_duration_ms = duration_ms

            self._materialize_output_transport(output_batch, req, save_output_paths)

            if (
                torch.cuda.is_initialized()
                and output_batch.output is None
                and not req.return_raw_frames
            ):
                torch.cuda.empty_cache()

            if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING:
                if not req.is_warmup:
                    PerformanceLogger.log_request_summary(metrics=output_batch.metrics)

            # dump per-request perf report to the server-mode file path.
            if (
                req.perf_dump_path is not None
                and not req.is_warmup
                and output_batch.metrics is not None
            ):
                PerformanceLogger.dump_benchmark_report(
                    file_path=req.perf_dump_path,
                    metrics=output_batch.metrics,
                    meta={"model": self.server_args.model_path},
                    tag="server_perf_dump",
                )
        except Exception as e:
            logger.error(
                f"Error executing {error_context}: {e}",
                exc_info=True,
            )
            if isinstance(e, _oom_exceptions()):
                logger.warning(OOM_MSG)
            if output_batch is None:
                output_batch = OutputBatch()
            output_batch.error = f"Error executing {error_context}: {e}"
            self._record_output_peak_memory(output_batch)
            # clean cache if OOM
            if torch.cuda.is_initialized():
                torch.cuda.empty_cache()
        return output_batch

    def _materialize_output_transport(
        self,
        output_batch: OutputBatch,
        req: Req,
        save_output_paths: Callable[[OutputBatch], None],
    ) -> None:
        if req.return_raw_frames:
            self._materialize_raw_frame_transport(output_batch, req)
        elif req.save_output and req.return_file_paths_only:
            self._materialize_file_path_transport(output_batch, save_output_paths)
        elif getattr(req, "return_frames", False):
            self._materialize_frame_outputs_for_return(output_batch, req)

    def _materialize_raw_frame_transport(
        self, output_batch: OutputBatch, req: Req
    ) -> None:
        if self.rank != 0:
            return
        if output_batch.output is not None:
            output_batch.raw_frame_content_type = RAW_RGB_CONTENT_TYPE
            (
                output_batch.raw_frame_batches,
                output_batch.raw_frame_metadata,
            ) = build_raw_rgb_frame_batches(
                output_batch.output,
                req,
                output_batch,
                post_process_sample,
            )
            output_batch.output = None
        output_batch.audio = None
        output_batch.audio_sample_rate = None

    def _materialize_file_path_transport(
        self,
        output_batch: OutputBatch,
        save_output_paths: Callable[[OutputBatch], None],
    ) -> None:
        if self.rank == 0:
            save_output_paths(output_batch)
        output_batch.output = None
        output_batch.audio = None
        output_batch.audio_sample_rate = None

    def _materialize_frame_outputs_for_return(
        self, output_batch: OutputBatch, req: Req
    ) -> None:
        """materialize the output from tensor to numpy frames for faster serialization"""
        if (
            self.rank != 0
            or output_batch.output is None
            or not getattr(req, "return_frames", False)
        ):
            return

        if (
            os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
            and torch.cuda.is_initialized()
        ):
            torch.cuda.synchronize()
        start_time = time.perf_counter()
        output_batch.output = [
            self._materialize_frame_output(output, output_batch, req)
            for output in output_batch.output
        ]
        if output_batch.metrics is not None:
            if (
                os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
                and torch.cuda.is_initialized()
            ):
                torch.cuda.synchronize()
            output_batch.metrics.record_stage(
                "GPUWorker.frame_materialize_for_return",
                time.perf_counter() - start_time,
            )

    @staticmethod
    def _materialize_frame_output(
        output: Any, output_batch: OutputBatch, req: Req
    ) -> np.ndarray:
        if (
            isinstance(output, torch.Tensor)
            and not req.enable_frame_interpolation
            and not req.enable_upscaling
        ):
            if output.dim() == 3:
                output = output.unsqueeze(1)
            output = (output * 255).clamp(0, 255).to(torch.uint8)
            return output.permute(1, 2, 3, 0).cpu().numpy()

        if (
            isinstance(output, np.ndarray)
            and output.dtype == np.uint8
            and output.ndim == 4
            and output.shape[-1] in (1, 3, 4)
        ):
            return output

        materialized = materialize_output_sample(
            output,
            req.data_type,
            req.fps,
            enable_frame_interpolation=req.enable_frame_interpolation,
            frame_interpolation_exp=req.frame_interpolation_exp,
            frame_interpolation_scale=req.frame_interpolation_scale,
            frame_interpolation_model_path=req.frame_interpolation_model_path,
            enable_upscaling=req.enable_upscaling,
            upscaling_model_path=req.upscaling_model_path,
            upscaling_scale=req.upscaling_scale,
        )
        return np.asarray(materialized.frames)

    def _record_output_peak_memory(self, output_batch: OutputBatch) -> None:
        if self.rank != 0 or current_platform.is_cpu():
            return
        peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
        output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)

    def _forward_group(self, batch: list[Req]) -> OutputBatch:
        assert self.pipeline is not None
        results = self.pipeline.forward_batch(batch, self.server_args)
        output_batches = [self._to_output_batch(result) for result in results]
        return self._merge_expanded_output_batches(output_batches)

    def _save_output_paths(self, req: Req, output_batch: OutputBatch) -> None:
        """save outputs to files"""
        if self.rank != 0 or output_batch.output is None:
            return

        dynamic_output_paths = None
        if req.extra:
            dynamic_output_paths = req.extra.get("dynamic_batch_output_paths")
        if dynamic_output_paths is not None and (
            len(dynamic_output_paths) != len(output_batch.output)
        ):
            logger.warning(
                "dynamic_batch_output_paths length mismatch (got=%d, expected=%d). "
                "Falling back to merged request output file naming.",
                len(dynamic_output_paths),
                len(output_batch.output),
            )
            dynamic_output_paths = None

        if dynamic_output_paths is not None:

            def build_output_path(idx: int) -> str:
                return dynamic_output_paths[idx]

        else:
            num_outputs = len(output_batch.output)

            def build_output_path(idx: int) -> str:
                return req.output_file_path(num_outputs, idx)

        output_batch.output_file_paths = save_outputs(
            output_batch.output,
            req.data_type,
            req.fps,
            True,
            build_output_path,
            audio=output_batch.audio,
            audio_sample_rate=output_batch.audio_sample_rate,
            output_compression=req.output_compression,
            enable_frame_interpolation=req.enable_frame_interpolation,
            frame_interpolation_exp=req.frame_interpolation_exp,
            frame_interpolation_scale=req.frame_interpolation_scale,
            frame_interpolation_model_path=req.frame_interpolation_model_path,
            enable_upscaling=req.enable_upscaling,
            upscaling_model_path=req.upscaling_model_path,
            upscaling_scale=req.upscaling_scale,
        )

    def _save_group_output_paths(
        self,
        reqs: list[Req],
        output_batch: OutputBatch,
    ) -> None:
        if self.rank != 0 or output_batch.output is None:
            return
        if len(output_batch.output) != len(reqs):
            raise RuntimeError(
                f"Expected {len(reqs)} grouped outputs, got {len(output_batch.output)}"
            )

        first_req = reqs[0]
        output_batch.output_file_paths = save_outputs(
            output_batch.output,
            first_req.data_type,
            first_req.fps,
            True,
            lambda idx: reqs[idx].output_file_path(1, 0),
            audio=output_batch.audio,
            audio_sample_rate=output_batch.audio_sample_rate,
            output_compression=first_req.output_compression,
            enable_frame_interpolation=first_req.enable_frame_interpolation,
            frame_interpolation_exp=first_req.frame_interpolation_exp,
            frame_interpolation_scale=first_req.frame_interpolation_scale,
            frame_interpolation_model_path=first_req.frame_interpolation_model_path,
            enable_upscaling=first_req.enable_upscaling,
            upscaling_model_path=first_req.upscaling_model_path,
            upscaling_scale=first_req.upscaling_scale,
        )

    @staticmethod
    def _validate_group_forward_reqs(reqs: list[Req]) -> None:
        """Validate fields that the grouped output/save path treats as shared."""
        first_req = reqs[0]
        shared_output_fields = (
            "save_output",
            "return_frames",
            "return_file_paths_only",
            "data_type",
            "fps",
            "output_compression",
            "enable_frame_interpolation",
            "frame_interpolation_exp",
            "frame_interpolation_scale",
            "frame_interpolation_model_path",
            "enable_upscaling",
            "upscaling_model_path",
            "upscaling_scale",
        )
        for req in reqs[1:]:
            mismatched = [
                field
                for field in shared_output_fields
                if getattr(req, field, None) != getattr(first_req, field, None)
            ]
            if mismatched:
                raise ValueError(
                    "Grouped execute_forward requires matching output settings; "
                    f"mismatched fields: {mismatched}"
                )

    @staticmethod
    def _iter_output_metrics(output_batch: OutputBatch):
        """Return all metrics objects carried by an output batch."""
        if output_batch.metrics_list is not None:
            return [
                metrics for metrics in output_batch.metrics_list if metrics is not None
            ]
        if output_batch.metrics is not None:
            return [output_batch.metrics]
        return []

    @staticmethod
    def _to_output_batch(result: Req | OutputBatch) -> OutputBatch:
        if isinstance(result, Req):
            return GPUWorker._req_to_output_batch(result)
        return result

    @staticmethod
    def _req_to_output_batch(result: Req) -> OutputBatch:
        return OutputBatch(
            output=result.output,
            audio=getattr(result, "audio", None),
            audio_sample_rate=getattr(result, "audio_sample_rate", None),
            metrics=result.metrics,
            trajectory_timesteps=getattr(result, "trajectory_timesteps", None),
            trajectory_latents=getattr(result, "trajectory_latents", None),
            rollout_trajectory_data=getattr(result, "rollout_trajectory_data", None),
            noise_pred=getattr(result, "noise_pred", None),
            trajectory_decoded=getattr(result, "trajectory_decoded", None),
        )

    @staticmethod
    def _merge_expanded_output_batches(
        output_batches: list[OutputBatch],
    ) -> OutputBatch:
        """Merge per-output batches produced by grouped execution."""
        merged = OutputBatch()
        parts = _ExpandedOutputParts()

        for output_batch in output_batches:
            GPUWorker._merge_expanded_singletons(merged, output_batch)
            GPUWorker._collect_expanded_parts(parts, output_batch)

        GPUWorker._finalize_expanded_parts(
            merged,
            parts,
            audio_sample_rate=output_batches[0].audio_sample_rate,
        )

        return merged

    @staticmethod
    def _merge_expanded_singletons(
        merged: OutputBatch, output_batch: OutputBatch
    ) -> None:
        if output_batch.error is not None and merged.error is None:
            merged.error = output_batch.error
        merged.peak_memory_mb = max(merged.peak_memory_mb, output_batch.peak_memory_mb)
        if (
            merged.trajectory_timesteps is None
            and output_batch.trajectory_timesteps is not None
        ):
            merged.trajectory_timesteps = output_batch.trajectory_timesteps
        if (
            merged.rollout_trajectory_data is None
            and output_batch.rollout_trajectory_data is not None
        ):
            merged.rollout_trajectory_data = output_batch.rollout_trajectory_data

    @staticmethod
    def _collect_expanded_parts(
        parts: _ExpandedOutputParts, output_batch: OutputBatch
    ) -> None:
        """Collect expanded outputs"""
        parts.metrics_list.append(output_batch.metrics)
        if output_batch.output_file_paths:
            parts.output_file_paths.extend(output_batch.output_file_paths)
        if isinstance(output_batch.output, torch.Tensor):
            parts.tensor_outputs.append(output_batch.output)
        elif output_batch.output is not None:
            parts.list_outputs.extend(output_batch.output)
        if isinstance(output_batch.audio, torch.Tensor):
            parts.tensor_audio.append(output_batch.audio)
        if isinstance(output_batch.trajectory_latents, torch.Tensor):
            parts.trajectory_latents.append(output_batch.trajectory_latents)
        if isinstance(output_batch.noise_pred, torch.Tensor):
            parts.noise_preds.append(output_batch.noise_pred)
        if output_batch.trajectory_decoded:
            GPUWorker._collect_trajectory_decoded(
                parts, output_batch.trajectory_decoded
            )

    @staticmethod
    def _collect_trajectory_decoded(
        parts: _ExpandedOutputParts, trajectory_decoded: list[torch.Tensor]
    ) -> None:
        if parts.trajectory_decoded_parts is None:
            parts.trajectory_decoded_parts = [[] for _ in trajectory_decoded]
        for index, decoded in enumerate(trajectory_decoded):
            parts.trajectory_decoded_parts[index].append(decoded)

    @staticmethod
    def _finalize_expanded_parts(
        merged: OutputBatch,
        parts: _ExpandedOutputParts,
        *,
        audio_sample_rate: int | None,
    ) -> None:
        """
        merge batched output
        """
        if parts.output_file_paths:
            merged.output_file_paths = parts.output_file_paths
        if any(metrics is not None for metrics in parts.metrics_list):
            merged.metrics_list = parts.metrics_list
            merged.metrics = next(
                metrics for metrics in parts.metrics_list if metrics is not None
            )
        if parts.tensor_outputs:
            merged.output = torch.cat(parts.tensor_outputs, dim=0)
        elif parts.list_outputs:
            merged.output = parts.list_outputs
        if parts.tensor_audio:
            merged.audio = torch.cat(parts.tensor_audio, dim=0)
            merged.audio_sample_rate = audio_sample_rate
        if parts.trajectory_latents:
            merged.trajectory_latents = torch.cat(parts.trajectory_latents, dim=0)
        if parts.noise_preds:
            merged.noise_pred = torch.cat(parts.noise_preds, dim=0)
        if parts.trajectory_decoded_parts:
            merged.trajectory_decoded = [
                torch.cat(decoded_step, dim=0)
                for decoded_step in parts.trajectory_decoded_parts
            ]

    def get_can_stay_resident_components(
        self, remaining_gpu_mem_gb: float
    ) -> List[str]:
        """
        Calculate which components can stay resident on GPU without being offloaded.
        """
        can_stay_resident = []
        if not self.pipeline:
            return can_stay_resident

        # Map memory_usage keys to server_args offload flags.
        # If the flag is False, the component is already resident, so we do not suggest it.
        # If the flag is True, it is currently offloaded, so it is a candidate to stay resident.
        offload_flags = {
            "transformer": self.server_args.dit_cpu_offload
            or self.server_args.is_dit_layerwise_offload_selected,
            "vae": self.server_args.vae_cpu_offload,
            "text_encoder": self.server_args.text_encoder_cpu_offload,
            "text_encoder_2": self.server_args.text_encoder_cpu_offload,
            "image_encoder": self.server_args.image_encoder_cpu_offload,
        }

        for name in OFFLOAD_DISABLE_RECOMMENDATION_ORDER:
            # Only consider components that are currently configured to be offloaded
            is_offload_configured = offload_flags.get(name, False)
            if not is_offload_configured:
                continue

            usage = self.pipeline.memory_usages.get(name)
            if usage is None:
                continue

            if usage <= remaining_gpu_mem_gb:
                can_stay_resident.append(name)
                remaining_gpu_mem_gb -= usage

        return can_stay_resident

    def set_lora(
        self,
        lora_nickname: Union[str, List[str]],
        lora_path: Union[str, None, List[Union[str, None]]] = None,
        target: Union[str, List[str]] = "all",
        strength: Union[float, List[float]] = 1.0,
        merge_mode: str | None = None,
    ) -> OutputBatch:
        """
        Set the LoRA adapter(s) for the pipeline.
        Supports both single LoRA (backward compatible) and multiple LoRA adapters.

        Args:
            lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings.
            lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None.
            target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
            strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
            merge_mode: Optional per-request LoRA merge mode.
        """
        if not isinstance(self.pipeline, LoRAPipeline):
            return OutputBatch(error="Lora is not enabled")
        self.pipeline.set_lora(
            lora_nickname, lora_path, target, strength, merge_mode=merge_mode
        )
        return OutputBatch()

    def merge_lora_weights(
        self, target: str = "all", strength: float = 1.0
    ) -> OutputBatch:
        """
        Merge LoRA weights.

        Args:
            target: Which transformer(s) to merge.
            strength: LoRA strength for merge, default 1.0.
        """
        if not isinstance(self.pipeline, LoRAPipeline):
            return OutputBatch(error="Lora is not enabled")
        self.pipeline.merge_lora_weights(target, strength)
        return OutputBatch()

    def unmerge_lora_weights(self, target: str = "all") -> OutputBatch:
        """
        Unmerge LoRA weights.

        Args:
            target: Which transformer(s) to unmerge.
        """
        if not isinstance(self.pipeline, LoRAPipeline):
            return OutputBatch(error="Lora is not enabled")
        self.pipeline.unmerge_lora_weights(target)
        return OutputBatch()

    def list_loras(self) -> OutputBatch:
        """
        List loaded LoRA adapters and current application status per module.
        """
        from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import (
            LoRAPipeline,
        )

        if not isinstance(self.pipeline, LoRAPipeline):
            return OutputBatch(error="Lora is not enabled")
        status = self.pipeline.get_lora_status()
        return OutputBatch(output=status)

    def release_memory_occupation(self) -> dict:
        return self._get_memory_occupation().release_memory_occupation()

    def resume_memory_occupation(self) -> dict:
        if self.memory_occupation is None:
            return {
                "success": True,
                "sleeping": False,
                "message": "already awake",
            }
        return self.memory_occupation.resume_memory_occupation()


OOM_MSG = """
OOM detected. Possible solutions:
  - If the OOM occurs during loading:
    1. Check available memory on every selected GPU, not only total capacity.
       In multi-GPU runs, the least-free selected GPU is the bottleneck.
    2. For single-GPU deployment, use `--performance-mode memory`, component CPU offload,
       or `--dit-layerwise-offload` for supported Wan/MOVA DiTs.
    3. For multi-GPU deployment, keep the default `--performance-mode auto` or set
       `--use-fsdp-inference true` to shard DiT weights with FSDP. FSDP is not a
       single-GPU substitute for CPU offload.
  - If the OOM occurs during runtime:
    1. Reduce resolution, `--num-frames`, or batch size.
    2. Use `--performance-mode memory` for lower memory usage.
    3. Enable SP/Ulysses/Ring for sequence-heavy workloads in multi-GPU setups.
    4. Use FSDP, with CFG parallelism when supported, for validated multi-GPU workloads.
    5. Use a lower-memory attention backend or quantization when available.
  Or, open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose
"""


def _oom_exceptions():
    # torch.OutOfMemoryError exists only in some PyTorch builds
    types = [torch.cuda.OutOfMemoryError]
    if hasattr(torch, "OutOfMemoryError"):
        types.append(torch.OutOfMemoryError)
    return tuple(types)


def run_scheduler_process(
    local_rank: int,
    rank: int,
    master_port: int,
    server_args: ServerArgs,
    pipe_writer: mp.connection.Connection,
    # For all workers: pipe to receive tasks from rank 0
    task_pipe_r: mp.connection.Connection,
    # For slave workers: pipe to send results back to rank 0
    result_pipe_w: mp.connection.Connection | None,
    # For rank 0 worker only: pipes to send tasks to slaves
    task_pipes_to_slaves: list[mp.connection.Connection] | None = None,
    # For rank 0 worker only: pipes to receive results from slaves
    result_pipes_from_slaves: list[mp.connection.Connection] | None = None,
) -> None:
    """
    The entry point for the worker process.
    Rank 0 acts as the master, handling ZMQ requests and coordinating slaves.
    Ranks > 0 act as slaves, waiting for tasks from the master.
    """
    kill_itself_when_parent_died()
    configure_logger(server_args)
    globally_suppress_loggers()
    if current_platform.is_cuda():
        set_cuda_arch()
    elif current_platform.is_musa():
        set_musa_arch()

    init_diffusion_tracing(server_args, f"DiffWorker_rank{rank}")

    port_args = PortArgs.from_server_args(server_args)

    # start the scheduler event loop
    assert task_pipes_to_slaves is not None
    assert result_pipes_from_slaves is not None
    from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler

    try:
        scheduler = Scheduler(
            server_args,
            gpu_id=rank,
            port_args=port_args,
            task_pipes_to_slaves=task_pipes_to_slaves,
            result_pipes_from_slaves=result_pipes_from_slaves,
            local_rank=local_rank,
        )
        logger.info(f"Worker {rank}: Scheduler loop started.")
        pipe_writer.send(
            {
                "status": "ready",
            }
        )
        scheduler.event_loop()
    except _oom_exceptions() as _e:
        logger.warning(OOM_MSG)
        raise
    finally:
        # Clean up resources to speed up shutdown
        if "scheduler" in locals():
            del scheduler
        gc.collect()
        if torch.cuda.is_initialized():
            torch.cuda.empty_cache()
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            torch.distributed.destroy_process_group()
        logger.info(f"Worker {rank}: Shutdown complete.")