VLA/wall_x/_vendor/harrix/serving/launch_serving.py

#!/usr/bin/env python3
"""
Server script for Wall-X model.

This script serves a Wall-X model using a websocket server, allowing
clients to connect and get action predictions from observations.

"""

import dataclasses
import enum
import inspect
import logging
import os
import socket
import sys
import yaml
import traceback
import tyro

from wall_x._vendor.harrix.serving.websocket_policy_server import WebsocketPolicyServer


def _server_model_config_to_infer_kwargs(model_config) -> dict:
    from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig

    infer_params = inspect.signature(InferConfig.__init__).parameters
    return {
        k: v
        for k, v in vars(model_config).items()
        if v is not None and k in infer_params
    }


def get_wallx_policy(model_config, image_passing_mode, serialize_actions=True):
    from wall_x._vendor.harrix.serving.policy.wall_x_policy import WallXPolicy
    from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig

    config = InferConfig(**_server_model_config_to_infer_kwargs(model_config))
    return WallXPolicy(config=config, image_passing_mode=image_passing_mode, serialize_actions=serialize_actions)




logger = logging.getLogger(__name__)


class EnvMode(enum.Enum):
    """Supported environments/datasets."""

    X2ROBOT = "x2robot"
    LIBERO = "libero"


@dataclasses.dataclass
class ServerModelConfig:
    """Configuration for loading a Wall-X model."""

    checkpoint_path: str | None = None
    train_config_path: str | None = None
    # robot_host: str = '0.0.0.0'
    # robot_port: int = 33723
    robot_type: str = "desktop"  # ["desktop", "turtle"]
    robot_action_start_ratio: float = (
        0.0  # proportion of action execution to start from
    )
    robot_action_end_ratio: float = 1.0  # proportion of action execution to end at
    robot_action_interpolate_multiplier: int = 10  # action interpolation multiplier
    robot_use_joint_angle_control: bool = (
        False  # use joint angle control (model must predict joints)
    )
    turtle_as_desktop: bool = (
        False  # use turtle platform as desktop with fixed base/head/camera/height
    )
    action_horizon: int = 32  # specify the correct horizon for the model
    action_dim: int | None = None
    model_device: str = "cuda"
    num_inference_timesteps: int = 10
    num_inference_steps: int | None = None
    cfg_scale: float | None = None
    seed: int | None = None
    save_video_dir: str = "./videos"

    # Please specify explicitly if the checkpoint was not trained on the x2robot dataset.
    norm_key: str | None = None
    # Model cameras; None = infer from train config ``data.key_mappings.camera``.
    cam_names: list[str] | None = None
    # Robot camera keys in incoming websocket observations.
    camera_front_key: str = "camera_front"
    camera_left_key: str = "camera_left"
    camera_right_key: str = "camera_right"
    # Serving prompt controls. If the client request has no instruction,
    # default_instruction is used. prompt_template follows train-config semantics.
    default_instruction: str | None = None
    prompt_template: str | None = None
    qwen25_prompt_template: str | None = None
    prompt_priority_order: str | None = None

@dataclasses.dataclass
class Args:
    """Arguments for the serve_wall_x script."""

    # Environment mode (used for default configurations)
    env: EnvMode = EnvMode.X2ROBOT

    # Model configuration. If not provided, uses default config for the environment
    model_config: ServerModelConfig | None = None

    # Default text prompt to use if not provided in observation
    default_prompt: str | None = None

    # Port to serve the policy on
    port: int = 43007

    # Host to bind the server to
    host: str = "0.0.0.0"

    # Enable debug logging
    debug: bool = False

    # Image passing mode
    image_passing_mode: str = "base64"  # ["numpy", "base64"]

    # Model type
    model_type: str = "wallx"  # OSS supports qwen2.5 Wall-X only

    # Serialize actions via robot_preprocessor (True for robot control, False for raw output)
    serialize_actions: bool = True

    # -- Dynamic batching -----------------------------------------
    # Set max_batch_size to enable dynamic batching. None = single mode.
    max_batch_size: int | None = None
    max_wait_time_ms: float = 0
    max_queue_size: int = 100
    timeout_ms: float = 30000

    # -- Engine flags ---------------------------------------------
    enable_experimental_engine: bool = False
    enable_cuda_graph: bool = False


# Default model configurations for each environment
DEFAULT_CONFIGS: dict[EnvMode, ServerModelConfig] = {
    EnvMode.X2ROBOT: ServerModelConfig(
        checkpoint_path=None,
        train_config_path=None,
        robot_action_start_ratio=0.0,
        robot_action_end_ratio=1.0,
        robot_action_interpolate_multiplier=10,
        robot_use_joint_angle_control=False,
        turtle_as_desktop=False,
        action_horizon=32,
        action_dim=None,
        model_device="cuda",
        num_inference_timesteps=10,
    ),
    EnvMode.LIBERO: ServerModelConfig(
        checkpoint_path=None,
        train_config_path=None,
        robot_type="desktop",
        robot_action_start_ratio=0.0,
        robot_action_end_ratio=1.0,
        robot_action_interpolate_multiplier=1,
        action_horizon=10,
        action_dim=None,
        model_device="cuda",
        num_inference_timesteps=10,
        cam_names=None,
    ),
}


def get_model_config(args: Args) -> ServerModelConfig:
    """Get model configuration from args or defaults."""
    if args.model_config is not None:
        return args.model_config

    if config := DEFAULT_CONFIGS.get(args.env):
        logger.info(f"Using default configuration for {args.env.value}")
        return config

    raise ValueError(
        f"No default configuration for {args.env.value}. "
        f"Please provide --model-config with model_path and action_tokenizer_path."
    )


def create_policy(args: Args):
    """Create a policy from the given arguments."""
    model_config = get_model_config(args)
    if args.model_type != "wallx":
        raise ValueError(
            f"Unsupported model type: {args.model_type!r}. "
            "The public package only supports model_type='wallx'."
        )
    policy = get_wallx_policy(
        model_config, args.image_passing_mode, args.serialize_actions
    )
    return policy


def main(args: Args) -> None:
    """Main function to start the Wall-X model server."""
    log_level = logging.DEBUG if args.debug else logging.INFO
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # Set engine environment variables
    if args.enable_experimental_engine:
        os.environ["ENABLE_EXPERIMENTAL_INFERENCE_ENGINE"] = "true"
        logger.info("ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=true")

    if args.enable_cuda_graph:
        os.environ["ENABLE_CUDA_GRAPH"] = "true"
        logger.info("ENABLE_CUDA_GRAPH=true")

    logger.info("Starting model server")
    logger.info(f"Model type: {args.model_type}")
    logger.info(f"Environment: {args.env.value}")
    logger.info(f"Port: {args.port}")
    logger.info(f"Host: {args.host}")
    logger.info(f"Serialize actions: {args.serialize_actions}")

    # Create policy
    try:
        policy = create_policy(args)
    except Exception as e:
        logger.error(f"Failed to create policy: {e}")
        logger.error(traceback.format_exc())
        sys.exit(1)

    # Get policy metadata
    policy_metadata = policy.metadata
    policy_metadata["env"] = args.env.value

    # Get network info
    hostname = socket.gethostname()
    try:
        local_ip = socket.gethostbyname(hostname)
    except Exception:
        local_ip = "unknown"

    logger.info(f"Server hostname: {hostname}")
    logger.info(f"Server IP: {local_ip}")
    logger.info(f"Server will be available at: ws://{args.host}:{args.port}")
    logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz")

    batching_str = (
        f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms"
        if args.max_batch_size
        else "disabled"
    )
    logger.info(f"Batching: {batching_str}")

    # Create and start server
    server = WebsocketPolicyServer(
        policy=policy,
        host=args.host,
        port=args.port,
        metadata=policy_metadata,
        max_batch_size=args.max_batch_size,
        max_wait_time_ms=args.max_wait_time_ms,
        max_queue_size=args.max_queue_size,
        timeout_ms=args.timeout_ms,
    )

    logger.info("Starting server...")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("Server stopped by user")
    except Exception as e:
        logger.error(f"Server error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main(tyro.cli(Args))