From fc0d9d3cd4858aef7d931d6beb419a54d5e12b2b Mon Sep 17 00:00:00 2001
From: Abhishek Enaguthi <enaguthiabhishek@gmail.com>
Date: Thu, 9 Jul 2026 04:22:22 -0700
Subject: [PATCH] Fix JAX inference timing to include GPU sync.

policy_timing.infer_ms previously stopped the timer at async JAX dispatch,
under-reporting GPU work by ~30ms vs the profiler. Sync actions with
jax.block_until_ready on the JAX path before recording infer_ms.

Fixes #983
---
 src/openpi/policies/policy.py      |  5 ++-
 src/openpi/policies/policy_test.py | 61 ++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/src/openpi/policies/policy.py b/src/openpi/policies/policy.py
index b9b708bdca..08d8c8f261 100644
--- a/src/openpi/policies/policy.py
+++ b/src/openpi/policies/policy.py
@@ -89,9 +89,12 @@ def infer(self, obs: dict, *, noise: np.ndarray | None = None) -> dict:  # type:
 
         observation = _model.Observation.from_dict(inputs)
         start_time = time.monotonic()
+        actions = self._sample_actions(sample_rng_or_pytorch_device, observation, **sample_kwargs)
+        if not self._is_pytorch_model:
+            actions = jax.block_until_ready(actions)
         outputs = {
             "state": inputs["state"],
-            "actions": self._sample_actions(sample_rng_or_pytorch_device, observation, **sample_kwargs),
+            "actions": actions,
         }
         model_time = time.monotonic() - start_time
         if self._is_pytorch_model:
diff --git a/src/openpi/policies/policy_test.py b/src/openpi/policies/policy_test.py
index 5808e5274a..eb53ccbe53 100644
--- a/src/openpi/policies/policy_test.py
+++ b/src/openpi/policies/policy_test.py
@@ -1,11 +1,72 @@
 from openpi_client import action_chunk_broker
+import jax
+import jax.numpy as jnp
+import numpy as np
 import pytest
+import time
 
+from openpi import transforms as _transforms
 from openpi.policies import aloha_policy
+from openpi.policies import policy as _policy
 from openpi.policies import policy_config as _policy_config
 from openpi.training import config as _config
 
 
+def test_jax_async_dispatch_under_reports_without_sync():
+    """Document #983 bug class: CPU timers stop before jitted work completes."""
+
+    @jax.jit
+    def slow_matmul():
+        x = jnp.ones((1024, 1024))
+        return x @ x.T
+
+    start = time.monotonic()
+    result = slow_matmul()
+    cpu_only_ms = (time.monotonic() - start) * 1000
+
+    start = time.monotonic()
+    result = slow_matmul()
+    jax.block_until_ready(result)
+    synced_ms = (time.monotonic() - start) * 1000
+
+    assert synced_ms >= cpu_only_ms
+
+
+def test_jax_policy_infer_syncs_actions_before_timing(monkeypatch):
+    calls: list[object] = []
+    real_block_until_ready = jax.block_until_ready
+
+    def spy_block_until_ready(value):
+        calls.append(value)
+        return real_block_until_ready(value)
+
+    monkeypatch.setattr(jax, "block_until_ready", spy_block_until_ready)
+
+    policy = object.__new__(_policy.Policy)
+    policy._is_pytorch_model = False
+    policy._input_transform = _transforms.compose([])
+    policy._output_transform = _transforms.compose([])
+    policy._sample_kwargs = {}
+    policy._rng = jax.random.key(0)
+
+    @jax.jit
+    def fake_sample(_rng, _observation, **_kwargs):
+        return jnp.ones((1, 4, 14))
+
+    policy._sample_actions = fake_sample
+
+    obs = {
+        "state": np.ones((14,), dtype=np.float32),
+        "image": {"cam_high": np.zeros((224, 224, 3), dtype=np.uint8)},
+        "image_mask": {"cam_high": np.array(True)},
+    }
+    result = policy.infer(obs)
+
+    assert calls, "expected jax.block_until_ready to run on the JAX inference path"
+    assert "policy_timing" in result
+    assert result["policy_timing"]["infer_ms"] >= 0
+
+
 @pytest.mark.manual
 def test_infer():
     config = _config.get_config("pi0_aloha_sim")