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))