[Interactive Drive] Add physics via Physx- #423

.gitignore

@@ -4,6 +4,8 @@ compile_commands.json
 # Visual Studio Code configs.
 .vscode/
 
+.pytest-tmp*
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

THIRD-PARTY-NOTICES

@@ -33,6 +33,7 @@ huggingface-hub        Apache-2.0   https://github.com/huggingface/huggingface_h
 loguru                 MIT          https://github.com/Delgan/loguru
 numpy                  BSD-3-Clause https://github.com/numpy/numpy
 nvidia-ml-py           BSD-3-Clause https://pypi.org/project/nvidia-ml-py/
+psutil                 BSD-3-Clause https://github.com/giampaolo/psutil
 safetensors            Apache-2.0   https://github.com/huggingface/safetensors
 torch                  BSD-3-Clause https://pytorch.org
 torchvision            BSD-3-Clause https://github.com/pytorch/vision
@@ -60,6 +61,11 @@ opencv-python-headless Apache-2.0   https://github.com/opencv/opencv-python
 grpcio, grpcio-tools   Apache-2.0   https://github.com/grpc/grpc
 shapely                BSD-3-Clause https://github.com/shapely/shapely
 ludus-renderer         Apache-2.0   (NVIDIA, internal distribution)
+cmake                  BSD-3-Clause https://github.com/Kitware/CMake
+NVIDIA PhysX 5.9.0     BSD-3-Clause https://github.com/NVIDIA-Omniverse/PhysX
+    ``ludus-renderer`` downloads the pinned PhysX source archive at first use,
+    verifies its SHA-256, and builds it into a platform-local cache. The source
+    and resulting native module are not redistributed in this repository.
 
 Files in
     integrations/omnidreams/omnidreams/conditioning/world_scenario/

docs/source/models/omnidreams.rst

@@ -185,6 +185,20 @@ Run the demo and stream to your browser:
 Then open ``http://<server-ip>:8080/`` in any browser on the same
 network and pick a scene from the picker in the bottom-right.
 
+Collision physics, the vehicle speed limit, and the collision visual effect are
+disabled by default. Add ``--game-mode`` to enable the speed limit and
+collisions with scene actors and static map geometry, along with the collision
+visual flare:
+
+.. code-block:: bash
+
+   uv run --package flashdreams-omnidreams interactive-drive \
+       --stream-mjpeg :8080 \
+       --game-mode
+
+Combine ``--game-mode`` with ``--disable-visual-flare`` to retain collision
+physics without the full-screen collision effect.
+
 .. note::
 
    **The first launch is slow.** The first time you start the demo, the world

flashdreams/flashdreams/infra/compile.py

@@ -17,13 +17,22 @@
 
 from __future__ import annotations
 
-from typing import Literal, TypeVar, cast
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Literal, TypeVar, cast
 
 import torch
 import torch.nn as nn
 
 M = TypeVar("M", bound=nn.Module)
 
+_INDUCTOR_CACHE_SUBDIR = "torchinductor"
+"""FlashDreams cache subdirectory for coupled FX-graph and Triton artifacts."""
+
+_TRITON_BUNDLE_PATCH_MARKER = "_flashdreams_complete_static_bundles"
+"""Class marker preventing duplicate installation of the Triton bundle repair."""
+
 CompileMode = Literal[
     "default",
     "reduce-overhead",
@@ -41,6 +50,81 @@
 """
 
 
+def _configure_inductor_cache() -> None:
+    """Place Inductor artifacts in the persistent FlashDreams cache by default."""
+    from torch._inductor.runtime.cache_dir_utils import default_cache_dir
+
+    cache_root = Path(
+        os.path.expanduser(
+            os.environ.get("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams")
+        )
+    )
+    configured_cache = os.environ.get("TORCHINDUCTOR_CACHE_DIR")
+    if configured_cache is not None and configured_cache != default_cache_dir():
+        return
+
+    os.environ["TORCHINDUCTOR_CACHE_DIR"] = str(cache_root / _INDUCTOR_CACHE_SUBDIR)
+
+
+def _add_static_autotuner_hashes_to_winners(
+    bundler: Any,
+    *,
+    to_path_key: Callable[[str], str],
+) -> None:
+    """Retain every cubin referenced by serialized static autotuners.
+
+    PyTorch's Triton bundler filters artifacts whenever any kernel records an
+    autotuning winner. The serialized static autotuners still retain all their
+    compile results, so filtering those cubins produces an internally incomplete
+    FX-graph cache entry. Add their hashes to the retained set before collection.
+
+    Args:
+        bundler: Active private ``TritonBundler`` class.
+        to_path_key: Convert a raw Triton hash to its cache-directory key.
+    """
+    winners = getattr(bundler, "_winners", None)
+    static_autotuners = getattr(bundler, "_static_autotuners", None)
+    if not winners or not static_autotuners:
+        return
+
+    for entry in static_autotuners:
+        for result in getattr(entry.kernel, "compile_results", ()):
+            kernel_hash = getattr(getattr(result, "kernel", None), "hash", None)
+            if isinstance(kernel_hash, str):
+                winners.add(to_path_key(kernel_hash))
+
+
+def _patch_triton_bundle_collection() -> None:
+    """Make PyTorch FX-graph bundles self-contained for static Triton launchers."""
+    try:
+        from torch._inductor.runtime.triton_heuristics import (
+            triton_hash_to_path_key,
+        )
+        from torch._inductor.triton_bundler import TritonBundler
+    except ImportError:
+        return
+
+    if getattr(TritonBundler, _TRITON_BUNDLE_PATCH_MARKER, False):
+        return
+
+    original_collect = getattr(TritonBundler, "collect", None)
+    if original_collect is None:
+        return
+
+    @classmethod
+    def collect_with_complete_static_bundles(
+        cls: Any, *args: Any, **kwargs: Any
+    ) -> Any:
+        _add_static_autotuner_hashes_to_winners(
+            cls,
+            to_path_key=triton_hash_to_path_key,
+        )
+        return original_collect(*args, **kwargs)
+
+    setattr(TritonBundler, "collect", collect_with_complete_static_bundles)
+    setattr(TritonBundler, _TRITON_BUNDLE_PATCH_MARKER, True)
+
+
 def compile_module(
     module: M,
     *,
@@ -60,4 +144,6 @@ def compile_module(
         The compiled module, statically typed as the same ``M`` so attribute
         access on the wrapped module continues to type-check at call sites.
     """
+    _configure_inductor_cache()
+    _patch_triton_bundle_collection()
     return cast(M, torch.compile(module, mode=mode))

flashdreams/pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     # non-subclassable TypeAliasType in NumPy 2.5 (notably on Python 3.13).
     "numpy>=1.24,<2.5",
     "nvidia-ml-py>=12.0",
+    "psutil>=7.0",
     "safetensors>=0.4",
     "tqdm>=4.60",
     "transformers>=5.0,<6",

flashdreams/tests/test_benchmark_harness.py

@@ -19,18 +19,15 @@
 
 import json
 import os
-import subprocess
 import sys
 import time
 import types
 from pathlib import Path
-from typing import cast
 
 import numpy as np
 import pytest
 
 from tools.benchmarks import cli as benchmark_cli
-from tools.benchmarks import harness as benchmark_harness
 from tools.benchmarks import pai_bench_profile
 from tools.benchmarks import quality as benchmark_quality
 from tools.benchmarks.harness import run_benchmark_suite
@@ -684,7 +681,6 @@ def test_run_benchmark_suite_emits_progress_heartbeat(tmp_path: Path) -> None:
     )
 
 
-@pytest.mark.skipif(os.name != "posix", reason="requires POSIX process groups")
 def test_scenario_timeout_terminates_descendant_processes(tmp_path: Path) -> None:
     child_marker = tmp_path / "child_survived.txt"
     spawned_marker = tmp_path / "child_spawned.txt"
@@ -735,142 +731,6 @@ def test_scenario_timeout_terminates_descendant_processes(tmp_path: Path) -> Non
     assert not child_marker.exists()
 
 
-def test_windows_popen_uses_new_process_group(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    monkeypatch.setattr(benchmark_harness.os, "name", "nt")
-    monkeypatch.setattr(
-        benchmark_harness.subprocess,
-        "CREATE_NEW_PROCESS_GROUP",
-        512,
-        raising=False,
-    )
-
-    assert benchmark_harness._process_group_popen_kwargs() == {"creationflags": 512}
-
-
-def test_windows_timeout_uses_taskkill_process_tree(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    class FakeProcess:
-        pid = 1234
-        returncode: int | None = None
-        killed = False
-
-        def poll(self) -> int | None:
-            return self.returncode
-
-        def wait(self, timeout: float | None = None) -> int:
-            if self.returncode is None:
-                raise subprocess.TimeoutExpired("fake", timeout or 0.0)
-            return self.returncode
-
-        def kill(self) -> None:
-            self.killed = True
-            self.returncode = -9
-
-    process = FakeProcess()
-    commands: list[list[str]] = []
-
-    def fake_run(
-        command: list[str],
-        **_: object,
-    ) -> subprocess.CompletedProcess[str]:
-        commands.append(command)
-        process.returncode = 1
-        return subprocess.CompletedProcess(command, 0)
-
-    monkeypatch.setattr(benchmark_harness.os, "name", "nt")
-    monkeypatch.setattr(benchmark_harness.subprocess, "run", fake_run)
-
-    benchmark_harness._terminate_process(cast(subprocess.Popen[str], process))
-
-    assert commands == [["taskkill", "/F", "/T", "/PID", "1234"]]
-    assert process.killed is False
-
-
-def test_windows_timeout_uses_powershell_when_taskkill_fails(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    class FakeProcess:
-        pid = 1234
-        returncode: int | None = None
-        killed = False
-
-        def poll(self) -> int | None:
-            return self.returncode
-
-        def wait(self, timeout: float | None = None) -> int:
-            if self.returncode is None:
-                raise subprocess.TimeoutExpired("fake", timeout or 0.0)
-            return self.returncode
-
-        def kill(self) -> None:
-            self.killed = True
-            self.returncode = -9
-
-    process = FakeProcess()
-    commands: list[list[str]] = []
-
-    def fake_run(
-        command: list[str],
-        **_: object,
-    ) -> subprocess.CompletedProcess[str]:
-        commands.append(command)
-        if command[0] == "taskkill":
-            return subprocess.CompletedProcess(command, 1)
-        process.returncode = 1
-        return subprocess.CompletedProcess(command, 0)
-
-    monkeypatch.setattr(benchmark_harness.os, "name", "nt")
-    monkeypatch.setattr(benchmark_harness.subprocess, "run", fake_run)
-
-    benchmark_harness._terminate_process(cast(subprocess.Popen[str], process))
-
-    assert commands[0] == ["taskkill", "/F", "/T", "/PID", "1234"]
-    assert commands[1][0] == "powershell.exe"
-    assert "-Command" in commands[1]
-    assert commands[1][-1] == "1234"
-    assert process.killed is False
-
-
-def test_windows_timeout_cleanup_failure_is_explicit(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    class FakeProcess:
-        pid = 1234
-        returncode: int | None = None
-        killed = False
-
-        def poll(self) -> int | None:
-            return self.returncode
-
-        def wait(self, timeout: float | None = None) -> int:
-            if self.returncode is None:
-                raise subprocess.TimeoutExpired("fake", timeout or 0.0)
-            return self.returncode
-
-        def kill(self) -> None:
-            self.killed = True
-            self.returncode = -9
-
-    process = FakeProcess()
-
-    def fake_run(
-        command: list[str],
-        **_: object,
-    ) -> subprocess.CompletedProcess[str]:
-        return subprocess.CompletedProcess(command, 1)
-
-    monkeypatch.setattr(benchmark_harness.os, "name", "nt")
-    monkeypatch.setattr(benchmark_harness.subprocess, "run", fake_run)
-
-    with pytest.raises(RuntimeError, match="Failed to terminate Windows process tree"):
-        benchmark_harness._terminate_process(cast(subprocess.Popen[str], process))
-
-    assert process.killed is True
-
-
 def test_quality_command_can_report_skipped_status(tmp_path: Path) -> None:
     script = (
         "from pathlib import Path; import sys; "