Use iris bundle for Ray working_dir to include pb2 files

rjpower · rjpower · commit 8e888345d9bf · 2026-03-26T16:14:17.000-07:00
Extract create_workspace_zip() from BundleCreator and use it in fray
v1/v2 backends and ray_run.py so that gitignored protobuf artifacts
(pb2 files) are included in Ray job bundles.
diff --git a/lib/fray/src/fray/v1/cluster/ray/cluster.py b/lib/fray/src/fray/v1/cluster/ray/cluster.py
@@ -17,6 +17,8 @@
 from ray.job_submission import JobStatus as RayJobStatus
 from ray.job_submission import JobSubmissionClient
 
+from iris.cluster.client.bundle import create_workspace_zip
+
 from fray.v1.cluster.base import (
     Cluster,
     EnvironmentConfig,
@@ -276,8 +278,12 @@ def _get_runtime_env(self, request: JobRequest) -> dict | None:
                 pip_packages=list(environment.pip_packages),
                 env_vars=env_vars,
             )
-            runtime_env["working_dir"] = environment.workspace
-            runtime_env["excludes"] = [".git", "tests/", "docs/", "**/*.pack"]
+            runtime_env["working_dir"] = create_workspace_zip(
+                environment.workspace,
+                exclude_dirs={"tests", "docs"},
+                exclude_extensions={".pack"},
+                max_size_bytes=None,
+            )
             runtime_env["config"] = {"setup_timeout_seconds": 1800}
         else:
             # No runtime package installation: rely on the existing environment.
diff --git a/lib/fray/src/fray/v2/ray_backend/backend.py b/lib/fray/src/fray/v2/ray_backend/backend.py
@@ -17,6 +17,8 @@
 from ray.job_submission import JobStatus as RayJobStatus
 from ray.job_submission import JobSubmissionClient
 
+from iris.cluster.client.bundle import create_workspace_zip
+
 from fray.v2.actor import (
     ActorContext,
     ActorFuture,
@@ -209,8 +211,12 @@ def build_runtime_env(request: JobRequest) -> dict:
             pip_packages=list(environment.pip_packages),
             env_vars=env_vars,
         )
-        runtime_env["working_dir"] = environment.workspace
-        runtime_env["excludes"] = [".git", "tests/", "docs/", "**/*.pack"]
+        runtime_env["working_dir"] = create_workspace_zip(
+            environment.workspace,
+            exclude_dirs={"tests", "docs"},
+            exclude_extensions={".pack"},
+            max_size_bytes=None,
+        )
         runtime_env["config"] = {"setup_timeout_seconds": 1800}
     else:
         python_path = build_python_path(submodules_dir=os.path.join(environment.workspace, "submodules"))
diff --git a/lib/iris/src/iris/cluster/client/bundle.py b/lib/iris/src/iris/cluster/client/bundle.py
@@ -35,21 +35,49 @@
     "tests/snapshot",
 }
 
+# Glob patterns for generated files that are gitignored but required at runtime.
+# These are produced by build hooks (e.g. hatch_build.py protobuf generation)
+# and must be included in task bundles so that `uv sync` inside containers can
+# skip regeneration.
+GENERATED_ARTIFACT_GLOBS = [
+    "src/iris/rpc/*_pb2.py",
+    "src/iris/rpc/*_pb2.pyi",
+    "src/iris/rpc/*_connect.py",
+    "lib/iris/src/iris/rpc/*_pb2.py",
+    "lib/iris/src/iris/rpc/*_pb2.pyi",
+    "lib/iris/src/iris/rpc/*_connect.py",
+]
 
-def _should_exclude(relative: Path) -> bool:
+
+def _should_exclude(
+    relative: Path,
+    extra_dirs: set[str] | None = None,
+    extra_extensions: set[str] | None = None,
+    extra_subpaths: set[str] | None = None,
+) -> bool:
     """Check whether a relative path should be excluded from the bundle."""
-    if relative.suffix in EXCLUDE_EXTENSIONS:
+    all_extensions = EXCLUDE_EXTENSIONS | (extra_extensions or set())
+    all_dirs = EXCLUDE_DIRS | (extra_dirs or set())
+    all_subpaths = EXCLUDE_SUBPATHS | (extra_subpaths or set())
+
+    if relative.suffix in all_extensions:
         return True
     # e.g. foo.egg-info
     if any(part.endswith(".egg-info") for part in relative.parts):
         return True
-    if any(part in EXCLUDE_DIRS for part in relative.parts):
+    if any(part in all_dirs for part in relative.parts):
         return True
     rel_str = str(relative)
-    return any(subpath in rel_str for subpath in EXCLUDE_SUBPATHS)
+    return any(subpath in rel_str for subpath in all_subpaths)
 
 
-def _get_git_non_ignored_files(workspace: Path) -> set[Path] | None:
+def get_git_non_ignored_files(
+    workspace: Path,
+    *,
+    exclude_dirs: set[str] | None = None,
+    exclude_extensions: set[str] | None = None,
+    exclude_subpaths: set[str] | None = None,
+) -> set[Path] | None:
     """Get files that are not ignored by git.
 
     Returns None if git is not available or this isn't a git repo.
@@ -63,41 +91,109 @@ def _get_git_non_ignored_files(workspace: Path) -> set[Path] | None:
             check=True,
         )
         files = [Path(f) for f in result.stdout.splitlines() if f]
-        files = [f for f in files if not _should_exclude(f)]
+        files = [f for f in files if not _should_exclude(f, exclude_dirs, exclude_extensions, exclude_subpaths)]
         return {workspace / f for f in files}
     except (subprocess.CalledProcessError, FileNotFoundError) as e:
         logger.debug("Git not available, using pattern-based exclusion: %s", e)
         return None
 
 
-# Glob patterns for generated files that are gitignored but required at runtime.
-# These are produced by build hooks (e.g. hatch_build.py protobuf generation)
-# and must be included in task bundles so that `uv sync` inside containers can
-# skip regeneration.
-_GENERATED_ARTIFACT_GLOBS = [
-    "src/iris/rpc/*_pb2.py",
-    "src/iris/rpc/*_pb2.pyi",
-    "src/iris/rpc/*_connect.py",
-    "lib/iris/src/iris/rpc/*_pb2.py",
-    "lib/iris/src/iris/rpc/*_pb2.pyi",
-    "lib/iris/src/iris/rpc/*_connect.py",
-]
-
-
-def _include_generated_build_artifacts(workspace: Path, files: set[Path]) -> None:
+def include_generated_build_artifacts(
+    workspace: Path,
+    files: set[Path],
+    *,
+    exclude_dirs: set[str] | None = None,
+    exclude_extensions: set[str] | None = None,
+    exclude_subpaths: set[str] | None = None,
+) -> None:
     """Add generated build artifacts that exist on disk but are gitignored."""
     added = 0
-    for pattern in _GENERATED_ARTIFACT_GLOBS:
+    for pattern in GENERATED_ARTIFACT_GLOBS:
         for path in workspace.glob(pattern):
-            if path.is_file() and path not in files and not _should_exclude(path.relative_to(workspace)):
+            if (
+                path.is_file()
+                and path not in files
+                and not _should_exclude(path.relative_to(workspace), exclude_dirs, exclude_extensions, exclude_subpaths)
+            ):
                 files.add(path)
                 added += 1
     if added:
         logger.debug("Included %d generated build artifact(s) in bundle", added)
 
 
+def create_workspace_zip(
+    workspace: str | Path,
+    *,
+    exclude_dirs: set[str] | None = None,
+    exclude_extensions: set[str] | None = None,
+    exclude_subpaths: set[str] | None = None,
+    max_size_bytes: int | None = MAX_BUNDLE_SIZE_BYTES,
+) -> str:
+    """Create a zip of the workspace suitable for Ray's working_dir or Iris bundles.
+
+    Uses git ls-files to determine which files to include (respecting .gitignore),
+    then adds back generated protobuf artifacts that are gitignored but needed at
+    runtime. When git is unavailable, falls back to pattern-based exclusion.
+
+    Args:
+        workspace: Root directory to bundle.
+        exclude_dirs: Additional directory names to exclude (merged with defaults).
+        exclude_extensions: Additional file extensions to exclude (merged with defaults).
+        exclude_subpaths: Additional subpath strings to exclude (merged with defaults).
+        max_size_bytes: Maximum allowed zip size. Pass None to disable the check.
+
+    Returns:
+        Path to the created zip file (in a temp directory; caller should not delete
+        the parent directory while the zip is in use).
+    """
+    workspace = Path(workspace)
+
+    git_files = get_git_non_ignored_files(
+        workspace,
+        exclude_dirs=exclude_dirs,
+        exclude_extensions=exclude_extensions,
+        exclude_subpaths=exclude_subpaths,
+    )
+    if git_files is not None:
+        include_generated_build_artifacts(
+            workspace,
+            git_files,
+            exclude_dirs=exclude_dirs,
+            exclude_extensions=exclude_extensions,
+            exclude_subpaths=exclude_subpaths,
+        )
+
+    # Use a persistent temp directory (not a context manager) so the caller
+    # can use the zip path after this function returns.
+    td = tempfile.mkdtemp(prefix="workspace_zip_")
+    zip_path = Path(td) / "workspace.zip"
+
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        if git_files is not None:
+            for file in git_files:
+                if file.is_file():
+                    zf.write(file, file.relative_to(workspace))
+        else:
+            for file in workspace.rglob("*"):
+                rel = file.relative_to(workspace)
+                if file.is_file() and not _should_exclude(rel, exclude_dirs, exclude_extensions, exclude_subpaths):
+                    zf.write(file, rel)
+
+    if max_size_bytes is not None:
+        zip_size = zip_path.stat().st_size
+        if zip_size > max_size_bytes:
+            zip_size_mb = zip_size / (1024 * 1024)
+            max_size_mb = max_size_bytes / (1024 * 1024)
+            raise ValueError(
+                f"Bundle size {zip_size_mb:.1f}MB exceeds maximum {max_size_mb:.0f}MB. "
+                "Consider excluding large files or using .gitignore."
+            )
+
+    return str(zip_path)
+
+
 class BundleCreator:
-    """Helper for creating workspace bundles.
+    """Helper for creating workspace bundles for Iris job submission.
 
     Bundles a user's workspace directory (containing pyproject.toml, uv.lock,
     and source code) into a zip file for job execution.
@@ -118,31 +214,5 @@ def create_bundle(self) -> bytes:
         Raises:
             ValueError: If bundle size exceeds MAX_BUNDLE_SIZE_BYTES
         """
-        git_files = _get_git_non_ignored_files(self._workspace)
-        if git_files is not None:
-            _include_generated_build_artifacts(self._workspace, git_files)
-
-        with tempfile.TemporaryDirectory(prefix="bundle_") as td:
-            bundle_path = Path(td) / "bundle.zip"
-            with zipfile.ZipFile(bundle_path, "w", zipfile.ZIP_DEFLATED) as zf:
-                if git_files is not None:
-                    for file in git_files:
-                        if file.is_file():
-                            zf.write(file, file.relative_to(self._workspace))
-                else:
-                    for file in self._workspace.rglob("*"):
-                        rel = file.relative_to(self._workspace)
-                        if file.is_file() and not _should_exclude(rel):
-                            zf.write(file, rel)
-
-            bundle_bytes = bundle_path.read_bytes()
-            bundle_size_mb = len(bundle_bytes) / (1024 * 1024)
-            max_size_mb = MAX_BUNDLE_SIZE_BYTES / (1024 * 1024)
-
-            if len(bundle_bytes) > MAX_BUNDLE_SIZE_BYTES:
-                raise ValueError(
-                    f"Bundle size {bundle_size_mb:.1f}MB exceeds maximum {max_size_mb:.0f}MB. "
-                    "Consider excluding large files or using .gitignore."
-                )
-
-            return bundle_bytes
+        zip_path = create_workspace_zip(self._workspace, max_size_bytes=MAX_BUNDLE_SIZE_BYTES)
+        return Path(zip_path).read_bytes()
diff --git a/lib/iris/tests/cluster/client/test_bundle.py b/lib/iris/tests/cluster/client/test_bundle.py
@@ -1,15 +1,17 @@
 # Copyright The Marin Authors
 # SPDX-License-Identifier: Apache-2.0
 
-"""Tests for BundleCreator."""
+"""Tests for BundleCreator and create_workspace_zip."""
 
 import io
+import os
 import zipfile
+from pathlib import Path
 from unittest.mock import patch
 
 import pytest
 
-from iris.cluster.client.bundle import MAX_BUNDLE_SIZE_BYTES, BundleCreator
+from iris.cluster.client.bundle import MAX_BUNDLE_SIZE_BYTES, BundleCreator, create_workspace_zip
 
 
 @pytest.fixture
@@ -24,7 +26,7 @@ def workspace(tmp_path):
 
 
 def test_bundle_creator_uses_fallback_when_git_unavailable(workspace):
-    with patch("iris.cluster.client.bundle._get_git_non_ignored_files", return_value=None):
+    with patch("iris.cluster.client.bundle.get_git_non_ignored_files", return_value=None):
         creator = BundleCreator(workspace)
         bundle_bytes = creator.create_bundle()
 
@@ -37,7 +39,7 @@ def test_bundle_creator_uses_fallback_when_git_unavailable(workspace):
 
 def test_bundle_creator_uses_git_files_when_available(workspace):
     git_files = {workspace / "pyproject.toml", workspace / "src" / "main.py"}
-    with patch("iris.cluster.client.bundle._get_git_non_ignored_files", return_value=git_files):
+    with patch("iris.cluster.client.bundle.get_git_non_ignored_files", return_value=git_files):
         creator = BundleCreator(workspace)
         bundle_bytes = creator.create_bundle()
 
@@ -58,9 +60,9 @@ def test_bundle_includes_generated_proto_files(workspace):
 
     # Simulate git ls-files returning only tracked files (not the generated ones)
     git_files = {workspace / "pyproject.toml", workspace / "src" / "main.py"}
-    with patch("iris.cluster.client.bundle._get_git_non_ignored_files") as mock_git:
+    with patch("iris.cluster.client.bundle.get_git_non_ignored_files") as mock_git:
         # Call the real function's logic but with controlled git output,
-        # then verify generated files are added via _include_generated_build_artifacts.
+        # then verify generated files are added via include_generated_build_artifacts.
         mock_git.return_value = git_files
         creator = BundleCreator(workspace)
         bundle_bytes = creator.create_bundle()
@@ -74,14 +76,43 @@ def test_bundle_includes_generated_proto_files(workspace):
 
 def test_bundle_creator_rejects_oversized_bundles(workspace):
     """Test that bundles exceeding MAX_BUNDLE_SIZE_BYTES are rejected."""
-    # Create a large file with random data that won't compress well
-    import os
-
     large_file = workspace / "large_file.bin"
     # Use urandom to create incompressible data
     large_file.write_bytes(os.urandom(MAX_BUNDLE_SIZE_BYTES + 1024 * 1024))
 
-    with patch("iris.cluster.client.bundle._get_git_non_ignored_files", return_value=None):
+    with patch("iris.cluster.client.bundle.get_git_non_ignored_files", return_value=None):
         creator = BundleCreator(workspace)
         with pytest.raises(ValueError, match=r"Bundle size .* exceeds maximum"):
             creator.create_bundle()
+
+
+def test_create_workspace_zip_returns_path_with_expected_files(workspace):
+    """create_workspace_zip returns a zip path containing workspace files, excluding __pycache__."""
+    with patch("iris.cluster.client.bundle.get_git_non_ignored_files", return_value=None):
+        zip_path = create_workspace_zip(workspace)
+
+    assert Path(zip_path).exists()
+    assert zip_path.endswith(".zip")
+
+    with zipfile.ZipFile(zip_path) as zf:
+        names = zf.namelist()
+        assert "pyproject.toml" in names
+        assert "src/main.py" in names
+        assert not any("__pycache__" in n for n in names)
+        assert not any(n.endswith(".pyc") for n in names)
+
+
+def test_create_workspace_zip_with_custom_exclude_dirs(workspace):
+    """create_workspace_zip respects custom exclude_dirs in addition to the defaults."""
+    (workspace / "experiments").mkdir()
+    (workspace / "experiments" / "run.py").write_text("# experiment")
+    (workspace / "src" / "core.py").write_text("# core")
+
+    with patch("iris.cluster.client.bundle.get_git_non_ignored_files", return_value=None):
+        zip_path = create_workspace_zip(workspace, exclude_dirs={"experiments"})
+
+    with zipfile.ZipFile(zip_path) as zf:
+        names = zf.namelist()
+        assert "src/core.py" in names
+        assert "pyproject.toml" in names
+        assert not any("experiments" in n for n in names)
diff --git a/lib/marin/src/marin/run/ray_run.py b/lib/marin/src/marin/run/ray_run.py
@@ -29,6 +29,7 @@
 from marin.cluster.config import find_config_by_region
 from fray.v1.cluster.ray import DashboardConfig, ray_dashboard
 from fray.v1.cluster.ray.deps import build_runtime_env_for_packages, accelerator_type_from_extra, AcceleratorType
+from iris.cluster.client.bundle import create_workspace_zip
 from iris.logging import configure_logging
 
 logger = logging.getLogger(__name__)
@@ -156,9 +157,14 @@ async def submit_and_track_job(
     logger.info(f"env_vars: {json.dumps(env_vars, indent=4)}")
 
     runtime_dict = {
-        "working_dir": current_dir,
+        "working_dir": create_workspace_zip(
+            current_dir,
+            exclude_dirs={"docs"},
+            exclude_extensions={".pack"},
+            exclude_subpaths={"lib/levanter/docs"},
+            max_size_bytes=None,
+        ),
         "config": {"setup_timeout_seconds": 1800},
-        "excludes": [".git", "docs/", "**/*.pack", "lib/levanter/docs"],
     }
 
     # add the TPU dependency for cluster jobs.
