sgl-project--sglang/python/sglang/multimodal_gen/configs/sample/vla.py

# SPDX-License-Identifier: Apache-2.0

import argparse
import dataclasses
import os
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from sglang.multimodal_gen.configs.sample.sampling_params import (
    DataType,
    _sanitize_filename,
)
from sglang.multimodal_gen.utils import StoreBoolean, expand_path_fields

if TYPE_CHECKING:
    from sglang.multimodal_gen.runtime.server_args import ServerArgs


@dataclass
class VLASamplingParams:
    """Sampling parameters for VLA/action-generation policies."""

    data_type: DataType = DataType.ACTION
    request_id: str | None = field(default=None, metadata={"batch_sig_exclude": True})
    prompt: str | list[str] | None = field(
        default="", metadata={"batch_sig_exclude": True}
    )
    num_outputs_per_prompt: int = 1
    seed: int | list[int] = field(default=42, metadata={"batch_sig_exclude": True})
    generator_device: str | None = None
    num_inference_steps: int = 10

    output_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
    output_file_name: str | None = field(
        default=None, metadata={"batch_sig_exclude": True}
    )
    save_output: bool = False
    return_file_paths_only: bool = False

    profile: bool = field(default=False, metadata={"batch_sig_exclude": True})
    num_profiled_timesteps: int = field(default=5, metadata={"batch_sig_exclude": True})
    profile_all_stages: bool = field(
        default=False, metadata={"batch_sig_exclude": True}
    )
    debug: bool = field(default=False, metadata={"batch_sig_exclude": True})
    perf_dump_path: str | None = field(
        default=None, metadata={"batch_sig_exclude": True}
    )
    suppress_logs: bool = field(default=False, metadata={"batch_sig_exclude": True})

    enable_sequence_shard: bool | None = None
    max_sequence_length: int | None = None
    no_override_protected_fields: bool = field(
        default=False, metadata={"batch_sig_exclude": True}
    )

    def __post_init__(self) -> None:
        self.data_type = DataType.ACTION
        self._validate()

        env_steps = os.environ.get("SGLANG_TEST_NUM_INFERENCE_STEPS")
        if env_steps is not None and self.num_inference_steps is not None:
            self.num_inference_steps = int(env_steps)

    def build_request_extra(self) -> dict[str, Any]:
        extra = {}
        diffusers_kwargs = getattr(self, "diffusers_kwargs", None)
        if diffusers_kwargs:
            extra["diffusers_kwargs"] = diffusers_kwargs
        explicit_fields = getattr(self, "_explicit_fields", None)
        if explicit_fields is not None:
            extra["explicit_fields"] = sorted(explicit_fields)
        return extra

    def apply_request_extra(self, req: Any) -> None:
        req.extra.update(self.build_request_extra())

    def _validate(self):
        if (
            not isinstance(self.num_outputs_per_prompt, int)
            or self.num_outputs_per_prompt <= 0
        ):
            raise ValueError(
                "num_outputs_per_prompt must be a positive int, "
                f"got {self.num_outputs_per_prompt!r}"
            )

        if isinstance(self.seed, list):
            if not self.seed:
                raise ValueError("seed list must not be empty")
            for seed in self.seed:
                if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0:
                    raise ValueError(
                        f"seed list must contain non-negative ints, got {self.seed!r}"
                    )
        elif (
            isinstance(self.seed, bool)
            or not isinstance(self.seed, int)
            or self.seed < 0
        ):
            raise ValueError(
                f"seed must be a non-negative int or list of ints, got {self.seed!r}"
            )

        if (
            not isinstance(self.num_inference_steps, int)
            or self.num_inference_steps <= 0
        ):
            raise ValueError(
                "num_inference_steps must be a positive int, "
                f"got {self.num_inference_steps!r}"
            )

        if self.generator_device not in (None, "cuda", "musa", "cpu"):
            raise ValueError(
                "generator_device must be one of None, 'cuda', 'musa', or 'cpu', "
                f"got {self.generator_device!r}"
            )

    def _validate_with_pipeline_config(self, pipeline_config):
        if not pipeline_config.task_type.is_action_gen():
            raise ValueError(
                f"VLASamplingParams requires an ACTION pipeline, got {pipeline_config.task_type.name}"
            )

    def _adjust(self, server_args: "ServerArgs"):
        expand_path_fields(self)
        self.data_type = DataType.ACTION
        self.return_file_paths_only = False
        if self.output_path is None and server_args.output_path is not None:
            self.output_path = server_args.output_path
        if self.output_path is None:
            self.save_output = False
        if self.save_output and not server_args.comfyui_mode:
            self._set_output_file_name()

    def _set_output_file_ext(self):
        if self.output_file_name and not self.output_file_name.endswith(".json"):
            self.output_file_name = f"{self.output_file_name}.json"

    def _set_output_file_name(self):
        if self.output_file_name is None:
            self.output_file_name = "vla_action"
        self.output_file_name = _sanitize_filename(self.output_file_name)
        self._set_output_file_ext()

    def output_file_path(self):
        if self.output_path is None or self.output_file_name is None:
            return None
        return os.path.join(self.output_path, self.output_file_name)

    def _merge_with_user_params(
        self,
        user_params: "VLASamplingParams",
        explicit_fields: set[str] | None = None,
    ):
        if user_params is None:
            return

        predefined_fields = set(type(self).__annotations__.keys())
        allow_override_protected = not user_params.no_override_protected_fields
        for field_info in dataclasses.fields(user_params):
            field_name = field_info.name
            user_value = getattr(user_params, field_name)
            if field_info.default is not dataclasses.MISSING:
                default_class_value = field_info.default
            elif field_info.default_factory is not dataclasses.MISSING:
                default_class_value = field_info.default_factory()
            else:
                default_class_value = dataclasses.MISSING

            if explicit_fields is not None:
                is_user_modified = field_name in explicit_fields
            else:
                is_user_modified = user_value != default_class_value
            is_protected_field = field_name in predefined_fields
            if is_user_modified and (
                allow_override_protected or not is_protected_field
            ):
                setattr(self, field_name, user_value)

        if explicit_fields is not None:
            self._explicit_fields = set(explicit_fields)
        self.__post_init__()

    @staticmethod
    def add_cli_args(parser: Any) -> Any:
        def add_argument(*name_or_flags, **kwargs):
            kwargs.setdefault("default", argparse.SUPPRESS)
            return parser.add_argument(*name_or_flags, **kwargs)

        add_argument(
            "--prompt",
            type=str,
            nargs="+",
            help="Language instruction(s) for the VLA policy.",
        )
        add_argument(
            "--num-inference-steps",
            type=int,
            help="Number of action denoising steps.",
        )
        add_argument(
            "--num-outputs-per-prompt",
            type=int,
            help="Number of candidate actions to generate per observation.",
        )
        add_argument(
            "--seed",
            type=int,
            nargs="+",
            help="Random seed for action noise generation.",
        )
        add_argument(
            "--generator-device",
            type=str,
            choices=["cuda", "musa", "cpu"],
            help="Device for random generator. Default: use the model-specific setting.",
        )
        add_argument(
            "--profile",
            action="store_true",
            help="Enable torch profiler for action denoising.",
        )
        add_argument(
            "--num-profiled-timesteps",
            type=int,
            help="Number of denoising timesteps to profile after warmup.",
        )
        add_argument(
            "--profile-all-stages",
            action="store_true",
            dest="profile_all_stages",
            help="Used with --profile, profile all pipeline stages.",
        )
        add_argument("--debug", action="store_true")
        add_argument(
            "--enable-sequence-shard",
            action=StoreBoolean,
            help="Enable sequence dimension shard with sequence parallelism.",
        )
        add_argument(
            "--max-sequence-length",
            type=int,
            help="Maximum prefix sequence length.",
        )
        add_argument(
            "--no-override-protected-fields",
            action="store_true",
            help="If set, disallow user params to override subclass-defined fields.",
        )
        return parser

    @classmethod
    def get_cli_args(cls, args: argparse.Namespace):
        sampling_params_fields = {attr.name for attr in dataclasses.fields(cls)}
        args_attrs = set(vars(args).keys())
        attrs = sampling_params_fields & args_attrs
        cli_args = {
            attr: getattr(args, attr)
            for attr in attrs
            if hasattr(args, attr) and getattr(args, attr) is not None
        }
        if isinstance(cli_args.get("seed"), list) and len(cli_args["seed"]) == 1:
            cli_args["seed"] = cli_args["seed"][0]
        return cli_args

    def output_size_str(self) -> str:
        return "action"

    def seconds(self) -> float:
        return 0.0