Upload UR5 full fine-tuned checkpoint: pi05_pour_full configs/serve_policy_from_checkpoint.py

configs/serve_policy_from_checkpoint.py

@@ -0,0 +1,335 @@
+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))