Accelerate testing. Add jax cache helper.

Author: amacati

crazyflow/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from functools import partial
+from pathlib import Path
 from typing import TypeVar
 
 import jax
@@ -99,3 +100,17 @@ def _add_marker_to_scene(self: BaseRender, marker: dict):
         self.scn.ngeom += 1
 
     BaseRender._add_marker_to_scene = _add_marker_to_scene
+
+
+def enable_cache(
+    cache_path: Path = Path("/tmp/jax_cache"),
+    min_entry_size_bytes: int = -1,
+    min_compile_time_secs: int = 0,
+    enable_xla_caches: bool = False,
+):
+    """Enable JAX cache."""
+    jax.config.update("jax_compilation_cache_dir", str(cache_path))
+    jax.config.update("jax_persistent_cache_min_entry_size_bytes", min_entry_size_bytes)
+    jax.config.update("jax_persistent_cache_min_compile_time_secs", min_compile_time_secs)
+    if enable_xla_caches:
+        jax.config.update("jax_persistent_cache_enable_xla_caches", "all")

examples/cache.py

@@ -0,0 +1,50 @@
+"""This example demonstrates persistent JAX function caching with the Sim class.
+
+When JAX functions decorated with @jit are first called, JAX traces and compiles them to XLA. This
+compilation is expensive but only needs to happen once. Subsequent calls reuse the compiled version.
+However, the cache is not persistent between Python sessions.
+
+The Sim class uses many jitted functions internally, particularly in the step() method which
+compiles a chain of physics and control functions. On the first step() call, the entire chain is
+compiled.
+
+After the Python session ends, the cached functions get lost. However, we can enable a persistent
+cache that is used when the function we are jit-compiling has been compiled before, which
+significantly speeds up the jit compile time at the first step.
+
+By enabling caching:
+1. The first run compiles and caches all jitted functions persistently
+2. The second run loads the cached compiled functions instead of recompiling
+3. This gives a significant speedup in initialization and first step times
+
+The cache persists between Python sessions, so compilation only needs to happen once on a machine,
+or when the cache directory is deleted.
+"""
+
+import time
+from pathlib import Path
+
+from crazyflow.sim import Sim
+from crazyflow.utils import enable_cache
+
+
+def main():
+    cache_dir = Path("/tmp/jax_cache_test")
+    if use_cache := cache_dir.exists():
+        print("Cache directory exists. This run will be fast.")
+        print("\nTo run without cache, delete the directory.")
+    else:
+        print("Cache directory does not exist. This run will be slow.")
+        print("\nTo run with cache, run this script again.")
+    enable_cache(cache_path=cache_dir)
+    t0 = time.perf_counter()
+    sim = Sim()
+    t1 = time.perf_counter()
+    sim.step()
+    t2 = time.perf_counter()
+    prefix = "Using cache: " if use_cache else "Not using cache: "
+    print(f"{prefix}\n Init: {t1 - t0:.3f}s\n Step: {t2 - t1:.3f}s")
+
+
+if __name__ == "__main__":
+    main()

examples/spiral.py

@@ -13,7 +13,7 @@ def control(start_xy: np.ndarray, t: float) -> np.ndarray:
 
 
 def main():
-    sim = Sim(n_drones=4, control=Control.state, integrator="rk4", physics="mujoco")
+    sim = Sim(n_drones=4, control=Control.state, integrator="rk4", physics="analytical")
     sim.reset()
     duration = 5.0
     fps = 60

tests/conftest.py

@@ -0,0 +1,7 @@
+import jax
+
+jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
+jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+# Do not enable XLA caches, crashes PyTest
+# jax.config.update("jax_persistent_cache_enable_xla_caches", "all")

tests/integration/test_examples.py

@@ -6,15 +6,18 @@
 import pytest
 
 EXAMPLES_DIR = Path(__file__).resolve().parent.parent.parent / "examples"
-example_scripts = list(sorted(EXAMPLES_DIR.glob("*.py")))
+# Make example_scripts a list of strings instead of Path objects so that pytest can use it in its
+# automatic printouts. We convert the elements back to Paths in the test function.
+example_scripts = [str(p) for p in sorted(EXAMPLES_DIR.glob("*.py"))]
 
 
-@pytest.mark.parametrize("example_script", example_scripts)
+@pytest.mark.parametrize("example_script", [str(p) for p in example_scripts])
 @pytest.mark.timeout(60)
 @pytest.mark.integration
-def test_example_main(example_script: Path):
+def test_example_main(example_script: str):
     """Dynamically import and execute the main function from an example script."""
     # Add the examples directory to sys.path to resolve imports
+    example_script = Path(example_script)
     sys.path.insert(0, str(EXAMPLES_DIR))
 
     # Dynamically import the module

tests/integration/test_interfaces.py

@@ -1,3 +1,4 @@
+import jax
 import numpy as np
 import pytest
 from scipy.spatial.transform import Rotation as R
@@ -31,14 +32,15 @@ def test_state_interface(physics: Physics):
 def test_attitude_interface(physics: Physics):
     sim = Sim(physics=physics, control=Control.attitude)
     target_pos = np.array([0.0, 0.0, 1.0])
+    jit_state2attitude = jax.jit(state2attitude)
 
     i_error = np.zeros((1, 1, 3))
 
     for _ in range(int(2 * sim.control_freq)):  # Run simulation for 2 seconds
         pos, vel, quat = sim.data.states.pos, sim.data.states.vel, sim.data.states.quat
         des_pos = np.array([[[0, 0, 1.0]]])
         dt = 1 / sim.data.controls.attitude_freq
-        cmd, i_error = state2attitude(
+        cmd, i_error = jit_state2attitude(
             pos, vel, quat, des_pos, np.zeros((1, 1, 3)), np.zeros((1, 1, 1)), i_error, dt
         )
         sim.attitude_control(cmd)

tests/unit/test_utils.py

@@ -0,0 +1,49 @@
+import jax
+import pytest
+
+from crazyflow.utils import enable_cache
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("enable_xla", [True, False])
+def test_enable_cache(enable_xla: bool):
+    """Test that enable_cache correctly sets JAX cache configuration."""
+    # Store original config values
+    orig_cache_dir = jax.config.values.get("jax_compilation_cache_dir", None)
+    orig_min_size = jax.config.values.get("jax_persistent_cache_min_entry_size_bytes", None)
+    orig_min_time = jax.config.values.get("jax_persistent_cache_min_compile_time_secs", None)
+    orig_xla = jax.config.values.get("jax_persistent_cache_enable_xla_caches", None)
+
+    try:
+        cache_path = "/tmp/jax_cache"
+        min_size = 1000
+        min_time = 2
+
+        enable_cache(
+            cache_path=cache_path,
+            min_entry_size_bytes=min_size,
+            min_compile_time_secs=min_time,
+            enable_xla_caches=enable_xla,
+        )
+
+        assert cache_path == jax.config.jax_compilation_cache_dir, "Cache path not set correctly"
+        assert (
+            min_size == jax.config.jax_persistent_cache_min_entry_size_bytes
+        ), "Min size not set correctly"
+        assert (
+            min_time == jax.config.jax_persistent_cache_min_compile_time_secs
+        ), "Min time not set correctly"
+        expected_xla = "all" if enable_xla else orig_xla
+        assert (
+            expected_xla == jax.config.jax_persistent_cache_enable_xla_caches
+        ), "XLA caches not set correctly"
+
+    finally:
+        if orig_cache_dir is not None:
+            jax.config.update("jax_compilation_cache_dir", orig_cache_dir)
+        if orig_min_size is not None:
+            jax.config.update("jax_persistent_cache_min_entry_size_bytes", orig_min_size)
+        if orig_min_time is not None:
+            jax.config.update("jax_persistent_cache_min_compile_time_secs", orig_min_time)
+        if orig_xla is not None:
+            jax.config.update("jax_persistent_cache_enable_xla_caches", orig_xla)