Auto-generate protobuf and Connect RPC files via hatch build hook

.github/workflows/iris-unit-tests.yaml

@@ -31,6 +31,11 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
+
       - name: Install uv
         uses: astral-sh/setup-uv@v7
         with:
@@ -58,6 +63,11 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
+
       - name: Install uv
         uses: astral-sh/setup-uv@v7
         with:

.github/workflows/levanter-tests.yaml

@@ -30,6 +30,10 @@ jobs:
           python-version: ${{ matrix.python-version }}
           enable-cache: true
           working-directory: lib/levanter
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
       - name: Set up Python
         run: uv python install
       - name: Install dependencies
@@ -57,6 +61,10 @@ jobs:
           python-version: ${{ matrix.python-version }}
           enable-cache: true
           working-directory: lib/levanter
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
       - name: Set up Python
         run: uv python install
       - name: Install dependencies
@@ -83,6 +91,10 @@ jobs:
           python-version: ${{ matrix.python-version }}
           enable-cache: true
           working-directory: lib/levanter
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
       - name: Set up Python
         run: uv python install
       - name: Install dependencies
@@ -110,6 +122,10 @@ jobs:
           python-version: ${{ matrix.python-version }}
           enable-cache: true
           working-directory: lib/levanter
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
       - name: Set up Python
         run: uv python install
       - name: Install dependencies
@@ -201,4 +217,22 @@ jobs:
             -v /tmp/uv-cache:/tmp/uv-cache:rw \
             -w /workspace \
             $DOCKER_IMAGE \
-            bash -c "cp -a /workspace-src/. /workspace/ && cd /workspace && timeout --kill-after=5 --signal=TERM 890 uv run --package levanter --frozen --group test --with 'jax[tpu]==$JAX_VERSION' pytest lib/levanter/tests -m 'not entry and not ray and not slow and not torch' --ignore=lib/levanter/tests/test_audio.py --ignore=lib/levanter/tests/test_new_cache.py --ignore=lib/levanter/tests/test_hf_checkpoints.py --ignore=lib/levanter/tests/test_hf_gpt2_serialize.py --ignore=lib/levanter/tests/test_gdn_layer.py -v --tb=short --log-cli-level=WARNING --durations=20"
+            bash -c "\
+              # Install Node.js in userspace if not present (needed for protobuf generation during uv sync)
+              if ! command -v npx >/dev/null 2>&1; then \
+                echo '::group::Installing Node.js in userspace'; \
+                curl -fsSL https://nodejs.org/dist/v22.16.0/node-v22.16.0-linux-x64.tar.xz | tar -xJ -C /tmp && \
+                export PATH=/tmp/node-v22.16.0-linux-x64/bin:\$PATH; \
+                echo '::endgroup::'; \
+              fi && \
+              cp -a /workspace-src/. /workspace/ && cd /workspace && \
+              timeout --kill-after=5 --signal=TERM 890 \
+              uv run --package levanter --frozen --group test --with 'jax[tpu]==$JAX_VERSION' \
+              pytest lib/levanter/tests \
+              -m 'not entry and not ray and not slow and not torch' \
+              --ignore=lib/levanter/tests/test_audio.py \
+              --ignore=lib/levanter/tests/test_new_cache.py \
+              --ignore=lib/levanter/tests/test_hf_checkpoints.py \
+              --ignore=lib/levanter/tests/test_hf_gpt2_serialize.py \
+              --ignore=lib/levanter/tests/test_gdn_layer.py \
+              -v --tb=short --log-cli-level=WARNING --durations=20"

.github/workflows/marin-docs.yaml

@@ -24,6 +24,11 @@ jobs:
           python-version: "3.11"
           enable-cache: true
 
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
+
       - name: Set up Python
         run: uv python install
 

.github/workflows/marin-itest.yaml

@@ -31,6 +31,11 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
+
       - name: Install uv
         uses: astral-sh/setup-uv@v7
         with:

.github/workflows/marin-lint-and-format.yaml

@@ -24,5 +24,10 @@ jobs:
         python-version: "3.11"
         enable-cache: true
 
+    - name: Set up Node.js
+      uses: actions/setup-node@v4
+      with:
+        node-version: '22'
+
     - name: Run pre-commit checks
       run: ./infra/pre-commit.py --all-files

.github/workflows/marin-unit-tests.yaml

@@ -39,7 +39,7 @@ jobs:
           uv sync --package marin --extra cpu --extra dedup --group test --frozen
 
       - name: Set up Node.js ${{ matrix.node-version }}
-        uses: actions/setup-node@v3
+        uses: actions/setup-node@v4
         with:
           node-version: ${{ matrix.node-version }}
 

.github/workflows/zephyr-unit-tests.yaml

@@ -31,6 +31,11 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
+
       - name: Install uv
         uses: astral-sh/setup-uv@v6
         with:

.gitignore

@@ -234,3 +234,8 @@ scr/*
 .agents/tmp/
 .codex
 .entire
+
+# Generated protobuf and Connect RPC files (rebuilt by hatch build hook)
+lib/iris/src/iris/rpc/*_pb2.py
+lib/iris/src/iris/rpc/*_pb2.pyi
+lib/iris/src/iris/rpc/*_connect.py

.readthedocs.yaml

@@ -4,6 +4,7 @@ build:
   os: "ubuntu-24.04"
   tools:
     python: "3.11"
+    nodejs: "22"
   commands:
     - pip install uv
     - uv sync --group docs --package marin

Makefile

@@ -205,5 +205,20 @@ setup_pre_commit:
 	chmod +x $$HOOK_PATH; \
 	echo "Installed git pre-commit hook -> $$HOOK_PATH"
 
-dev_setup: install_uv install_gcloud get_secret_key get_ray_auth_token setup_pre_commit
+install_node:
+	@if command -v node > /dev/null 2>&1; then \
+		echo "Node.js $$(node --version) is already installed."; \
+	elif command -v brew > /dev/null 2>&1; then \
+		echo "Installing Node.js via Homebrew..."; \
+		brew install node; \
+	elif command -v apt-get > /dev/null 2>&1; then \
+		echo "Installing Node.js via apt..."; \
+		curl -fsSL https://deb.nodesource.com/setup_22.x | sudo -E bash - && \
+		sudo apt-get install -y nodejs; \
+	else \
+		echo "Cannot auto-install Node.js. Please install manually: https://nodejs.org/"; \
+		exit 1; \
+	fi
+
+dev_setup: install_uv install_gcloud install_node get_secret_key get_ray_auth_token setup_pre_commit
 	@echo "Dev setup complete."

docker/marin/Dockerfile.tpu-ci

@@ -20,6 +20,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     libssl-dev \
     python3.11 \
     python3.11-dev \
+    && curl -fsSL https://deb.nodesource.com/setup_22.x | bash - \
+    && apt-get install -y --no-install-recommends nodejs \
     && rm -rf /var/lib/apt/lists/*
 
 RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1

infra/pre-commit.py

@@ -528,10 +528,45 @@ def check_notebooks(files: list[pathlib.Path], fix: bool) -> int:
     return _record("Jupyter notebooks", 0)
 
 
+def _ensure_iris_protos() -> None:
+    """Generate iris protobuf files if they are missing and npx is available.
+
+    Pyrefly needs the generated *_pb2.py and *_connect.py files on disk to
+    resolve imports from other modules, even when the generated files themselves
+    are excluded from type-checking via project-excludes.
+    """
+    import shutil
+
+    rpc_dir = ROOT_DIR / "lib" / "iris" / "src" / "iris" / "rpc"
+    # Check if any pb2 file exists already
+    if list(rpc_dir.glob("*_pb2.py")):
+        return
+
+    generate_script = ROOT_DIR / "lib" / "iris" / "scripts" / "generate_protos.py"
+    if not generate_script.exists():
+        return
+
+    if shutil.which("npx") is None:
+        print("  ⚠ Iris protobuf files are missing and npx is not installed; pyrefly may report false errors")
+        return
+
+    print("  Generating iris protobuf files for type checking...")
+    result = subprocess.run(
+        [sys.executable, str(generate_script)],
+        cwd=ROOT_DIR / "lib" / "iris",
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        print(f"  ⚠ Proto generation failed: {result.stderr.strip()}")
+
+
 def check_pyrefly(files: list[pathlib.Path], fix: bool) -> int:
     if not files:
         return 0
 
+    _ensure_iris_protos()
+
     args = ["uvx", "pyrefly@0.42.0", "check", "--baseline", ".pyrefly-baseline.json"]
     result = run_cmd(args)
     output = (result.stdout + result.stderr).strip()

lib/iris/hatch_build.py

@@ -0,0 +1,145 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Hatchling custom build hook for Iris.
+
+Regenerates protobuf files from .proto sources and rebuilds the Vue dashboard
+when source files are newer than their generated outputs. This runs automatically
+during ``uv sync`` / ``pip install -e .`` / wheel builds, eliminating the need
+to check generated files into git or manually run build steps.
+"""
+
+import logging
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+
+logger = logging.getLogger(__name__)
+
+# Glob patterns for source and generated files, relative to the iris package root.
+_PROTO_SOURCE_GLOBS = ["src/iris/rpc/*.proto"]
+_PROTO_OUTPUT_GLOBS = ["src/iris/rpc/*_pb2.py", "src/iris/rpc/*_pb2.pyi", "src/iris/rpc/*_connect.py"]
+
+_DASHBOARD_SOURCE_GLOBS = ["dashboard/src/**/*", "dashboard/package.json", "dashboard/rsbuild.config.ts"]
+_DASHBOARD_OUTPUT_DIR = "dashboard/dist"
+
+
+def _newest_mtime(root: Path, globs: list[str]) -> float:
+    """Return the newest mtime across all files matching the given globs."""
+    newest = 0.0
+    for pattern in globs:
+        for path in root.glob(pattern):
+            if path.is_file():
+                newest = max(newest, path.stat().st_mtime)
+    return newest
+
+
+def _oldest_mtime(root: Path, globs: list[str]) -> float:
+    """Return the oldest mtime across all files matching the given globs.
+
+    Returns 0.0 if no files match (meaning outputs don't exist yet).
+    """
+    oldest = float("inf")
+    found = False
+    for pattern in globs:
+        for path in root.glob(pattern):
+            if path.is_file():
+                found = True
+                oldest = min(oldest, path.stat().st_mtime)
+    return oldest if found else 0.0
+
+
+def _outputs_exist(root: Path, output_globs: list[str]) -> bool:
+    """Return True if at least one output file exists."""
+    return _oldest_mtime(root, output_globs) > 0.0
+
+
+def _needs_rebuild(root: Path, source_globs: list[str], output_globs: list[str]) -> bool:
+    """Return True if any source file is newer than the oldest output file."""
+    source_newest = _newest_mtime(root, source_globs)
+    output_oldest = _oldest_mtime(root, output_globs)
+    return source_newest > output_oldest
+
+
+class CustomBuildHook(BuildHookInterface):
+    PLUGIN_NAME = "iris-build"
+
+    def initialize(self, version: str, build_data: dict) -> None:
+        root = Path(self.root)
+        self._maybe_generate_protos(root)
+        self._maybe_build_dashboard(root)
+
+    def _maybe_generate_protos(self, root: Path) -> None:
+        outputs_present = _outputs_exist(root, _PROTO_OUTPUT_GLOBS)
+
+        if outputs_present and not _needs_rebuild(root, _PROTO_SOURCE_GLOBS, _PROTO_OUTPUT_GLOBS):
+            logger.info("Protobuf outputs are up-to-date, skipping generation")
+            return
+
+        generate_script = root / "scripts" / "generate_protos.py"
+        if not generate_script.exists():
+            if not outputs_present:
+                raise RuntimeError(
+                    "Protobuf outputs are missing and scripts/generate_protos.py not found. "
+                    "Cannot build iris without generated protobuf files."
+                )
+            logger.warning("scripts/generate_protos.py not found, using existing protobuf outputs")
+            return
+
+        if shutil.which("npx") is None:
+            if not outputs_present:
+                raise RuntimeError(
+                    "Protobuf outputs are missing and npx is not installed. "
+                    "Install Node.js (which provides npx) to generate protobuf files: "
+                    "https://nodejs.org/ or run `make install_node`"
+                )
+            logger.warning("npx not found, using existing (possibly stale) protobuf outputs")
+            return
+
+        logger.info("Regenerating protobuf files from .proto sources...")
+        result = subprocess.run(
+            [sys.executable, str(generate_script)],
+            cwd=root,
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            raise RuntimeError(f"Protobuf generation failed:\n{result.stdout}\n{result.stderr}")
+        logger.info("Protobuf generation complete")
+
+    def _maybe_build_dashboard(self, root: Path) -> None:
+        dashboard_dir = root / "dashboard"
+        if not (dashboard_dir / "package.json").exists():
+            logger.info("Dashboard source not found, skipping build")
+            return
+
+        dist_dir = root / _DASHBOARD_OUTPUT_DIR
+        dist_present = dist_dir.exists() and any(dist_dir.iterdir())
+
+        if shutil.which("npm") is None:
+            if not dist_present:
+                logger.warning(
+                    "npm not found and dashboard/dist is missing. "
+                    "Dashboard will not be available. Install Node.js to build it."
+                )
+            return
+
+        source_newest = _newest_mtime(root, _DASHBOARD_SOURCE_GLOBS)
+        if dist_present and source_newest > 0:
+            output_oldest = _oldest_mtime(root, [f"{_DASHBOARD_OUTPUT_DIR}/**/*"])
+            if output_oldest > 0 and source_newest <= output_oldest:
+                logger.info("Dashboard assets are up-to-date, skipping build")
+                return
+
+        logger.info("Building dashboard assets...")
+        result = subprocess.run(["npm", "ci"], cwd=dashboard_dir, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(f"npm ci failed:\n{result.stdout}\n{result.stderr}")
+
+        result = subprocess.run(["npm", "run", "build"], cwd=dashboard_dir, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(f"Dashboard build failed:\n{result.stdout}\n{result.stderr}")
+        logger.info("Dashboard build complete")