configs/serve_policy_from_checkpoint.py

import dataclasses
import json
import logging
import pathlib
import socket
from typing import Any

import torch
import tyro

from openpi.policies import policy as _policy
from openpi.policies import policy_config as _policy_config
from openpi.serving import websocket_policy_server
from openpi.shared import normalize as _normalize
from openpi.training import config as _config


def _torch_load_metadata(metadata_path: pathlib.Path) -> dict[str, Any]:
    try:
        return torch.load(metadata_path, map_location="cpu", weights_only=False)
    except TypeError:
        return torch.load(metadata_path, map_location="cpu")


def load_train_config_from_checkpoint(
    checkpoint_dir: pathlib.Path,
    *,
    config_name: str | None = None,
    lora_rank: int | None = None,
    lora_alpha: float | None = None,
) -> _config.TrainConfig:
    """Load a TrainConfig and restore checkpoint-time model overrides when available."""

    checkpoint_dir = checkpoint_dir.resolve()
    metadata_path = checkpoint_dir / "metadata.pt"
    metadata_config: dict[str, Any] | None = None

    if metadata_path.exists():
        metadata = _torch_load_metadata(metadata_path)
        metadata_config = metadata.get("config")
        if not isinstance(metadata_config, dict):
            metadata_config = None

    resolved_config_name = config_name
    if resolved_config_name is None and metadata_config is not None:
        resolved_config_name = metadata_config.get("name")
    if resolved_config_name is None:
        raise ValueError(
            "Could not infer config name. Pass --config-name, or use a checkpoint with metadata.pt."
        )

    train_config = _config.get_config(resolved_config_name)

    if metadata_config is not None:
        model_config = metadata_config.get("model")
        if isinstance(model_config, dict):
            for key, value in model_config.items():
                if hasattr(train_config.model, key):
                    object.__setattr__(train_config.model, key, value)

    if lora_rank is not None:
        object.__setattr__(train_config.model, "lora_rank", lora_rank)
    if lora_alpha is not None:
        object.__setattr__(train_config.model, "lora_alpha", lora_alpha)

    return train_config


@dataclasses.dataclass
class Args:
    checkpoint_dir: pathlib.Path
    config_name: str | None = None
    default_prompt: str | None = None
    port: int = 8000
    pytorch_device: str | None = None
    num_denoise_steps: int = 10
    model_label: str | None = None
    reference_metrics_json: pathlib.Path | None = None
    norm_stats_path: pathlib.Path | None = None
    lora_rank: int | None = None
    lora_alpha: float | None = None
    record: bool = False


def main(args: Args) -> None:
    checkpoint_metadata = _load_checkpoint_metadata(args.checkpoint_dir)
    train_config = load_train_config_from_checkpoint(
        args.checkpoint_dir,
        config_name=args.config_name,
        lora_rank=args.lora_rank,
        lora_alpha=args.lora_alpha,
    )
    logging.info("Serving config: %s", train_config.name)
    logging.info("Serving checkpoint: %s", args.checkpoint_dir)

    norm_stats = _load_norm_stats_override(args.norm_stats_path)
    policy = _policy_config.create_trained_policy(
        train_config,
        args.checkpoint_dir,
        default_prompt=args.default_prompt,
        sample_kwargs={"num_steps": args.num_denoise_steps},
        norm_stats=norm_stats,
        pytorch_device=args.pytorch_device,
    )
    _log_policy_efficiency_context(policy, args.checkpoint_dir)
    policy_metadata = dict(policy.metadata or {})
    policy_metadata.update(
        _server_metadata(
            policy,
            args,
            train_config.name,
            checkpoint_metadata,
        )
    )
    logging.info("Server metadata: %s", policy_metadata)

    if args.record:
        policy = _policy.PolicyRecorder(policy, "policy_records")

    hostname = socket.gethostname()
    local_ip = socket.gethostbyname(hostname)
    logging.info("Creating server (host: %s, ip: %s, port: %d)", hostname, local_ip, args.port)

    server = websocket_policy_server.WebsocketPolicyServer(
        policy=policy,
        host="0.0.0.0",
        port=args.port,
        metadata=policy_metadata,
    )
    server.serve_forever()


def _load_norm_stats_override(path: pathlib.Path | None):
    if path is None:
        return None
    path = pathlib.Path(path).expanduser().resolve()
    if path.is_file():
        norm_stats = _normalize.deserialize_json(path.read_text(encoding="utf-8"))
        logging.info("Loaded explicit norm stats from %s", path)
        return norm_stats
    norm_stats = _normalize.load(path)
    logging.info("Loaded explicit norm stats from %s", path / "norm_stats.json")
    return norm_stats


def _load_checkpoint_metadata(checkpoint_dir: pathlib.Path) -> dict[str, Any]:
    metadata_path = pathlib.Path(checkpoint_dir) / "metadata.pt"
    if not metadata_path.exists():
        return {}
    return _torch_load_metadata(metadata_path)


def _server_metadata(
    policy: _policy.Policy,
    args: Args,
    config_name: str,
    checkpoint_metadata: dict[str, Any],
) -> dict[str, Any]:
    global_step = checkpoint_metadata.get("global_step")
    if global_step is None:
        try:
            global_step = int(pathlib.Path(args.checkpoint_dir).name)
        except ValueError:
            global_step = None

    reference = _load_reference_metrics(args.reference_metrics_json)
    metadata = {
        "model_label": args.model_label or config_name,
        "config_name": config_name,
        "checkpoint_dir": str(pathlib.Path(args.checkpoint_dir).resolve()),
        "global_step": global_step,
        "num_denoise_steps": args.num_denoise_steps,
        "pytorch_device": args.pytorch_device,
        "reference_metrics_json": str(args.reference_metrics_json) if args.reference_metrics_json else None,
        "norm_stats_path": str(args.norm_stats_path.resolve()) if args.norm_stats_path else None,
        "reference_inference": reference,
        "model": _collect_model_metadata(policy, args.checkpoint_dir, checkpoint_metadata),
        "training_efficiency": _collect_training_efficiency(checkpoint_metadata),
    }
    return metadata


def _load_reference_metrics(path: pathlib.Path | None) -> dict[str, Any]:
    if path is None:
        return {}
    path = pathlib.Path(path)
    if not path.exists():
        logging.warning("Reference metrics JSON does not exist: %s", path)
        return {}
    metrics = json.loads(path.read_text(encoding="utf-8"))
    latency = metrics.get("latency") or {}
    frequency = metrics.get("frequency") or {}
    model_infer = latency.get("model_infer_ms") or {}
    wall = latency.get("wall_ms") or {}
    flop_profile = metrics.get("flop_profile") or {}
    return {
        "source": str(path),
        "warmup_runs": _number_or_none(metrics.get("warmup_runs")),
        "timed_runs": _number_or_none(metrics.get("timed_runs")),
        "num_denoise_steps": _number_or_none(metrics.get("num_denoise_steps")),
        "action_horizon": _number_or_none(metrics.get("action_horizon")),
        "actions_shape": _jsonable_value(metrics.get("actions_shape")),
        "model_infer_ms": _jsonable_value(model_infer),
        "wall_ms": _jsonable_value(wall),
        "model_infer_per_action_ms": _jsonable_value(latency.get("model_infer_per_action_ms") or {}),
        "wall_per_action_ms": _jsonable_value(latency.get("wall_per_action_ms") or {}),
        "latency_decomposition": _jsonable_value(metrics.get("latency_decomposition") or {}),
        "approx_gflops": _number_or_none(metrics.get("approx_gflops")),
        "flop_profile": {
            key: _jsonable_value(value)
            for key, value in flop_profile.items()
        },
        "profiler_gflops": _number_or_none(flop_profile.get("profiler_gflops")),
        "manual_conv2d_gflops": _number_or_none(flop_profile.get("manual_conv2d_gflops")),
        "manual_conv2d_gflops_added": _number_or_none(flop_profile.get("manual_conv2d_gflops_added")),
        "profiler_conv2d_gflops": _number_or_none(flop_profile.get("profiler_conv2d_gflops")),
        "model_infer_ms_median": _number_or_none(model_infer.get("median")),
        "model_infer_ms_p95": _number_or_none(model_infer.get("p95")),
        "wall_ms_median": _number_or_none(wall.get("median")),
        "model_query_hz": _number_or_none(frequency.get("model_query_hz")),
        "model_action_generation_hz": _number_or_none(frequency.get("model_action_generation_hz")),
        "wall_query_hz": _number_or_none(frequency.get("wall_query_hz")),
        "wall_action_generation_hz": _number_or_none(frequency.get("wall_action_generation_hz")),
    }


def _collect_model_metadata(
    policy: _policy.Policy,
    checkpoint_dir: pathlib.Path,
    checkpoint_metadata: dict[str, Any],
) -> dict[str, Any]:
    metrics = dict(checkpoint_metadata.get("efficiency_metrics") or {})
    model = getattr(policy, "_model", None)
    if model is not None:
        if hasattr(model, "configure_trainable_parameters"):
            model.configure_trainable_parameters()
        params = list(model.parameters())
        total_params = sum(param.numel() for param in params)
        trainable_params = sum(param.numel() for param in params if param.requires_grad)
        metrics.update(
            {
                "model_total_params": total_params,
                "model_trainable_params": trainable_params,
                "model_trainable_pct": trainable_params / total_params * 100 if total_params else 0.0,
                "model_total_params_b": total_params / 1e9,
                "model_trainable_params_m": trainable_params / 1e6,
            }
        )

    weight_path = pathlib.Path(checkpoint_dir) / "model.safetensors"
    if weight_path.exists():
        checkpoint_size_bytes = weight_path.stat().st_size
        metrics["checkpoint_size_bytes"] = checkpoint_size_bytes
        metrics["checkpoint_size_gb"] = checkpoint_size_bytes / 1024**3

    return {key: _jsonable_scalar(value) for key, value in metrics.items()}


def _collect_training_efficiency(checkpoint_metadata: dict[str, Any]) -> dict[str, Any]:
    metrics = dict(checkpoint_metadata.get("efficiency_metrics") or {})
    for key in (
        "elapsed_training_hours",
        "elapsed_after_first_step_hours",
        "elapsed_run_hours",
        "steps_per_second",
        "steps_per_second_after_first_step",
        "samples_per_second",
        "samples_per_second_after_first_step",
        "cuda_peak_allocated_gb",
        "cuda_peak_reserved_gb",
    ):
        if key in checkpoint_metadata and key not in metrics:
            metrics[key] = checkpoint_metadata[key]
    return {key: _jsonable_scalar(value) for key, value in metrics.items()}


def _jsonable_scalar(value: Any) -> Any:
    number = _number_or_none(value)
    if number is not None:
        return number
    if value is None or isinstance(value, (str, bool)):
        return value
    return str(value)


def _jsonable_value(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(key): _jsonable_value(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_jsonable_value(item) for item in value]
    return _jsonable_scalar(value)


def _number_or_none(value: Any) -> float | int | None:
    try:
        if hasattr(value, "item"):
            value = value.item()
        number = float(value)
    except (TypeError, ValueError):
        return None
    if not torch.isfinite(torch.tensor(number)):
        return None
    return int(number) if number.is_integer() else number


def _log_policy_efficiency_context(policy: _policy.Policy, checkpoint_dir: pathlib.Path) -> None:
    model = getattr(policy, "_model", None)
    if model is not None:
        total_params = sum(param.numel() for param in model.parameters())
        trainable_params = sum(param.numel() for param in model.parameters() if param.requires_grad)
        logging.info(
            "Model parameters: total=%.2fM trainable=%.2fM trainable_pct=%.2f%%",
            total_params / 1e6,
            trainable_params / 1e6,
            trainable_params / total_params * 100 if total_params else 0.0,
        )

    weight_path = pathlib.Path(checkpoint_dir) / "model.safetensors"
    if weight_path.exists():
        logging.info("Checkpoint weight size: %.2fGB", weight_path.stat().st_size / 1024**3)

    if torch.cuda.is_available():
        device = torch.device(getattr(policy, "_pytorch_device", "cuda"))
        if device.type == "cuda":
            torch.cuda.reset_peak_memory_stats(device)
            logging.info(
                "CUDA memory after model load: allocated=%.2fGB reserved=%.2fGB",
                torch.cuda.memory_allocated(device) / 1024**3,
                torch.cuda.memory_reserved(device) / 1024**3,
            )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, force=True)
    main(tyro.cli(Args))