[executor] Add mirrored() for region-agnostic data references (#3990)

Add MirroredValue/mirrored() config wrapper that marks input paths for
cross-region mirroring with per-path transfer budgets. The executor
resolves mirrored paths to the local marin prefix and copies data from
whichever regional bucket has it before step execution. Also fixes
MirrorFileSystem.ls to union local and remote prefixes so glob()
discovers files across regions, and makes the transfer budget
configurable via MARIN_MIRROR_BUDGET_GB env var.

Converts hardcoded gs:// input paths in exp934_hq_vs_pt.py
(stackexchange,
wikipedia, ar5iv, medu QA) to use mirrored() so the initial ferry can
run from any region.

Fixes #3989

Author: rjpower

docs/explanations/executor.md

@@ -35,6 +35,42 @@ Coordination between multiple pipelines is handled via lease files. This
 prevents duplicate execution if, for example, 2 Executor pipelines share common
 ancestor steps.
 
+## Mirrored inputs
+
+Some datasets live in a specific regional bucket (e.g.
+`gs://marin-us-central2/documents/stackexchange/...`) but experiments may run
+from any region.  The `mirrored()` wrapper marks an input path for
+**cross-region mirroring** so that the executor copies the data to the local
+marin prefix before the step runs.
+
+```python
+from marin.execution.executor import mirrored, versioned
+
+step = ExecutorStep(
+    name="train",
+    fn=my_training_fn,
+    config=TrainConfig(
+        dataset=mirrored(versioned("documents/stackexchange/v1"), budget_gb=50),
+    ),
+)
+```
+
+At config instantiation time, `mirrored()` rewrites the path to use the
+`mirror://` protocol.  When the step's function opens the path via `fsspec`,
+the `MirrorFileSystem` transparently copies data from whichever regional bucket
+has it into the local marin prefix, respecting the per-path transfer budget.
+
+**Key details:**
+
+- `budget_gb` (default 10) caps how much data (in GB) a single step may copy
+  cross-region.  The budget is enforced via the `mirror_budget` context manager
+  from `iris.marin_fs`.
+- Paths that already exist in the local prefix are not re-copied.
+- `mirrored()` can wrap plain strings or `VersionedValue` / `InputName`
+  references.
+- To adjust the global mirror budget default, set the `MARIN_MIRROR_BUDGET_GB`
+  environment variable before the process starts.
+
 ## Ray
 
 Recall that a step's function can either be a normal Python function or in most realistic cases,

experiments/exp934_hq_vs_pt.py

@@ -8,7 +8,7 @@
 datasets used by various training experiments.
 """
 
-from marin.execution.executor import ExecutorStep, this_output_path, versioned
+from marin.execution.executor import ExecutorStep, mirrored, this_output_path, versioned
 from marin.schemas.web.convert import HtmlToMarkdownConfig, ResiliparseConfig
 from marin.schemas.web.selectors import ARXIV_BLACKLISTED_SELECTORS, WIKI_BLACKLISTED_SELECTORS
 from marin.transform.ar5iv.transform_ar5iv import Ar5ivExtractionConfig, process_ar5iv_dump
@@ -29,7 +29,7 @@
     name="documents/stackexchange-resiliparse-custom-fork",
     fn=process_stackexchange_dump,
     config=StackExchangeExtractionConfig(
-        input_path=versioned("gs://marin-us-central2/documents/stackexchange/v2024-04-02/md-complete"),
+        input_path=mirrored(versioned("documents/stackexchange/v2024-04-02/md-complete"), budget_gb=50),
         output_path=this_output_path(),
         extract_method="resiliparse",
         extract_config=ResiliparseConfig(
@@ -48,7 +48,7 @@
         name="documents/wikipedia-resiliparse-custom-fork",
         fn=process_wiki_dump,
         config=WikiExtractionConfig(
-            input_path="gs://marin-us-central2/raw/wikipedia-a7dad0/20241201",
+            input_path=mirrored("raw/wikipedia-a7dad0/20241201", budget_gb=1),
             revision=versioned("20241201"),
             output_path=this_output_path(),
             extract_method="resiliparse",
@@ -72,7 +72,7 @@
     name="documents/ar5iv/ar5iv-04-2024-no-problem",
     fn=process_ar5iv_dump,
     config=Ar5ivExtractionConfig(
-        input_path="gs://marin-us-central2/raw/ar5iv/ar5iv-04-2024-no-problem-49c4e3/202404",
+        input_path=mirrored("raw/ar5iv/ar5iv-04-2024-no-problem-49c4e3/202404", budget_gb=1),
         revision="042024",
         output_path=this_output_path("resiliparse-custom-fork"),
         extract_method=versioned("resiliparse"),
@@ -88,7 +88,7 @@
 # MMLU Science QA tokenization
 medu_mmlu_science_qa_tokenized = default_tokenize(
     name="medu-mmlu-science-qa",
-    dataset="gs://marin-us-east1/documents/medu-mmlu-science-llama8b-qa-whole-1a419d",
+    dataset=mirrored("documents/medu-mmlu-science-llama8b-qa-whole-1a419d", budget_gb=30),
     tokenizer=llama3_tokenizer,
 ).with_output_path("tokenized/medu-mmlu-science-qa-c64fda")
 

lib/iris/src/iris/marin_fs.py

@@ -23,6 +23,8 @@
   the ``MARIN_I_WILL_PAY_FOR_ALL_FEES`` env var to override the guard.
 """
 
+import contextlib
+import contextvars
 import dataclasses
 import functools
 import logging
@@ -33,7 +35,7 @@
 import time
 import urllib.error
 import urllib.request
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Generator, Sequence
 from pathlib import PurePath
 from typing import Any
 
@@ -375,7 +377,18 @@ def _normalize_path_like(path: str | os.PathLike) -> str:
 # ---------------------------------------------------------------------------
 
 MARIN_CROSS_REGION_OVERRIDE_ENV: str = "MARIN_I_WILL_PAY_FOR_ALL_FEES"
-CROSS_REGION_TRANSFER_LIMIT_BYTES: int = 10 * 1024 * 1024 * 1024  # 10 GB
+MARIN_MIRROR_BUDGET_ENV: str = "MARIN_MIRROR_BUDGET_GB"
+_DEFAULT_TRANSFER_LIMIT_GB: int = 10
+
+
+def _transfer_limit_bytes() -> int:
+    raw = os.environ.get(MARIN_MIRROR_BUDGET_ENV, "")
+    if raw:
+        return int(float(raw) * 1024 * 1024 * 1024)
+    return _DEFAULT_TRANSFER_LIMIT_GB * 1024 * 1024 * 1024
+
+
+CROSS_REGION_TRANSFER_LIMIT_BYTES: int = _transfer_limit_bytes()
 
 # GCS multi-region bucket locations are returned as "us", "eu", or "asia"
 # rather than a specific region like "us-central1".  European regions use the
@@ -450,6 +463,34 @@ def reset(self, limit_bytes: int | None = None) -> None:
 
 _global_transfer_budget = TransferBudget()
 
+_mirror_budget_ctx: contextvars.ContextVar[TransferBudget | None] = contextvars.ContextVar(
+    "_mirror_budget_ctx", default=None
+)
+
+
+def set_mirror_budget(budget_gb: float) -> contextvars.Token:
+    """Set the MirrorFileSystem transfer budget for the current context.
+
+    Returns a token that can be used to reset the budget.
+    """
+    budget = TransferBudget(limit_bytes=int(budget_gb * 1024 * 1024 * 1024))
+    return _mirror_budget_ctx.set(budget)
+
+
+def reset_mirror_budget(token: contextvars.Token) -> None:
+    """Reset the MirrorFileSystem transfer budget to its previous value."""
+    _mirror_budget_ctx.reset(token)
+
+
+@contextlib.contextmanager
+def mirror_budget(budget_gb: float) -> Generator[None, None, None]:
+    """Context manager to scope a MirrorFileSystem transfer budget."""
+    token = set_mirror_budget(budget_gb)
+    try:
+        yield
+    finally:
+        reset_mirror_budget(token)
+
 
 @functools.lru_cache(maxsize=1)
 def _cached_marin_region() -> str | None:
@@ -684,6 +725,15 @@ def __init__(
         self._budget = budget if budget is not None else _global_transfer_budget
         self._worker_id = default_worker_id()
 
+    # -- budget resolution ----------------------------------------------------
+
+    def _active_budget(self) -> TransferBudget:
+        """Return the contextvar budget if set, otherwise the instance budget."""
+        ctx_budget = _mirror_budget_ctx.get()
+        if ctx_budget is not None:
+            return ctx_budget
+        return self._budget
+
     # -- underlying fs helpers ------------------------------------------------
 
     def _get_fs_and_path(self, url: str) -> tuple[Any, str]:
@@ -755,7 +805,7 @@ def _copy_to_local(self, source_prefix: str, path: str) -> None:
 
             size = self._fs_size(remote_url)
             if size is not None:
-                self._budget.record(size, remote_url)
+                self._active_budget().record(size, remote_url)
 
             logger.info("Mirror: copying %s → %s", remote_url, local_url)
             self._fs_copy(remote_url, local_url)
@@ -785,27 +835,38 @@ def _info(self, path: str, **kwargs: Any) -> dict[str, Any]:
         info["name"] = path
         return info
 
+    @staticmethod
+    def _stripped_prefix(bucket_prefix: str) -> str:
+        """Return the bucket prefix without scheme, with trailing slash."""
+        return bucket_prefix.rstrip("/").replace("gs://", "").replace("file://", "") + "/"
+
     def ls(self, path: str, detail: bool = True, **kwargs: Any) -> list[Any]:
         path = self._strip_protocol(path)
-        local_url = self._local_url(path)
-        fs, fspath = self._get_fs_and_path(local_url)
-        try:
-            results = fs.ls(fspath, detail=detail, **kwargs)
-        except FileNotFoundError:
-            results = []
-
-        prefix = self._local_prefix.rstrip("/") + "/"
-        # For GCS, fsspec strips the scheme, so the prefix in results won't have gs://
-        stripped_prefix = prefix.replace("gs://", "").replace("file://", "")
-
+        # Union listings from local + all remote prefixes so that glob()
+        # discovers files that only exist in other regions.  Local entries
+        # take precedence when a relative path appears in multiple buckets.
+        seen: dict[str, dict[str, Any]] = {}
+
+        for prefix in [self._local_prefix, *self._remote_prefixes]:
+            url = f"{prefix}/{path}"
+            fs, fspath = self._get_fs_and_path(url)
+            try:
+                entries = fs.ls(fspath, detail=True, **kwargs)
+            except FileNotFoundError:
+                continue
+
+            stripped = self._stripped_prefix(prefix)
+            for entry in entries:
+                rel_name = entry["name"]
+                if rel_name.startswith(stripped):
+                    rel_name = rel_name[len(stripped) :]
+                if rel_name not in seen:
+                    seen[rel_name] = {**entry, "name": rel_name}
+
+        results = list(seen.values())
         if detail:
-            for entry in results:
-                name = entry["name"]
-                if name.startswith(stripped_prefix):
-                    entry["name"] = name[len(stripped_prefix) :]
             return results
-        else:
-            return [r[len(stripped_prefix) :] if r.startswith(stripped_prefix) else r for r in results]
+        return [e["name"] for e in results]
 
     def exists(self, path: str, **kwargs: Any) -> bool:
         path = self._strip_protocol(path)

lib/iris/tests/test_mirror_fs.py

@@ -104,3 +104,88 @@ def test_second_read_uses_local_cache(mirror_fs, mirror_env):
 def test_read_finds_file_in_second_remote(mirror_fs, mirror_env):
     _write_file(mirror_env["remote2"], "data/remote2.txt", b"from-remote2")
     assert mirror_fs.cat_file("data/remote2.txt") == b"from-remote2"
+
+
+# ---------------------------------------------------------------------------
+# ls / glob tests
+# ---------------------------------------------------------------------------
+
+
+def test_ls_returns_local_entries(mirror_fs, mirror_env):
+    _write_file(mirror_env["local_dir"], "data/a.jsonl", b"a")
+    _write_file(mirror_env["local_dir"], "data/b.jsonl", b"b")
+
+    entries = mirror_fs.ls("data", detail=False)
+    assert sorted(entries) == ["data/a.jsonl", "data/b.jsonl"]
+
+
+def test_ls_discovers_remote_only_entries(mirror_fs, mirror_env):
+    _write_file(mirror_env["remote1"], "data/remote.jsonl", b"r")
+
+    entries = mirror_fs.ls("data", detail=False)
+    assert "data/remote.jsonl" in entries
+
+
+def test_ls_unions_local_and_remote(mirror_fs, mirror_env):
+    _write_file(mirror_env["local_dir"], "data/local.jsonl", b"l")
+    _write_file(mirror_env["remote1"], "data/remote.jsonl", b"r")
+
+    entries = mirror_fs.ls("data", detail=False)
+    assert sorted(entries) == ["data/local.jsonl", "data/remote.jsonl"]
+
+
+def test_ls_local_takes_precedence(mirror_fs, mirror_env):
+    _write_file(mirror_env["local_dir"], "data/file.jsonl", b"local-version")
+    _write_file(mirror_env["remote1"], "data/file.jsonl", b"remote-version")
+
+    entries = mirror_fs.ls("data", detail=True)
+    assert len(entries) == 1
+    assert entries[0]["name"] == "data/file.jsonl"
+
+
+def test_ls_unions_multiple_remotes(mirror_fs, mirror_env):
+    _write_file(mirror_env["remote1"], "data/from_r1.jsonl", b"r1")
+    _write_file(mirror_env["remote2"], "data/from_r2.jsonl", b"r2")
+
+    entries = mirror_fs.ls("data", detail=False)
+    assert sorted(entries) == ["data/from_r1.jsonl", "data/from_r2.jsonl"]
+
+
+def test_ls_empty_when_path_missing_everywhere(mirror_fs):
+    entries = mirror_fs.ls("nonexistent/path", detail=False)
+    assert entries == []
+
+
+def test_glob_discovers_remote_files(mirror_fs, mirror_env):
+    _write_file(mirror_env["remote1"], "docs/a.jsonl.gz", b"a")
+    _write_file(mirror_env["remote1"], "docs/b.jsonl.gz", b"b")
+    _write_file(mirror_env["remote1"], "docs/skip.txt", b"skip")
+
+    matched = mirror_fs.glob("docs/*.jsonl.gz")
+    assert sorted(matched) == ["docs/a.jsonl.gz", "docs/b.jsonl.gz"]
+
+
+# ---------------------------------------------------------------------------
+# mirror_budget context manager tests
+# ---------------------------------------------------------------------------
+
+
+def test_mirror_budget_context_manager(mirror_fs, mirror_env):
+    """Transfer budget set via context manager is used for copies."""
+    from iris.marin_fs import mirror_budget
+
+    _write_file(mirror_env["remote1"], "data/big.bin", b"x" * 1000)
+
+    with mirror_budget(budget_gb=0.001):  # ~1MB
+        assert mirror_fs.cat_file("data/big.bin") == b"x" * 1000
+
+
+def test_mirror_budget_context_manager_blocks_over_budget(mirror_fs, mirror_env):
+    from iris.marin_fs import mirror_budget
+
+    mirror_fs._budget.reset(limit_bytes=10 * 1024 * 1024 * 1024)  # high instance budget
+    _write_file(mirror_env["remote1"], "data/big.bin", b"x" * 2000)
+
+    with mirror_budget(budget_gb=0.000001):  # ~1KB — too small
+        with pytest.raises(TransferBudgetExceeded):
+            mirror_fs.cat_file("data/big.bin")