Wire temporary checkpoint roots into Marin launches

Author: dlwh

experiments/grug/base/launch.py

@@ -23,6 +23,7 @@
 from levanter.trainer import TrainerConfig
 from marin.execution.executor import ExecutorStep, executor_main, this_output_path, versioned
 from marin.processing.tokenize import add_validation_sets_to_mixture
+from marin.training.training import temporary_checkpoint_base_path
 
 from experiments.defaults import default_validation_sets
 from experiments.grug.base.model import GrugModelConfig
@@ -99,6 +100,7 @@ def run_grug_base_trial(config: GrugBaseLaunchConfig) -> None:
         allow_nondivisible_batch_size=False,
         checkpointer=CheckpointerConfig(
             base_path=os.path.join(config.output_path, "checkpoints"),
+            temporary_base_path=temporary_checkpoint_base_path(config.output_path),
             append_run_id_to_base_path=False,
             save_interval=timedelta(minutes=10),
             keep=[{"every": 1000}],

experiments/grug/modular_opt/launch.py

@@ -26,6 +26,7 @@
 from levanter.utils.jax_utils import leaf_key_paths
 from marin.execution.executor import ExecutorStep, executor_main, this_output_path, versioned
 from marin.processing.tokenize import add_validation_sets_to_mixture
+from marin.training.training import temporary_checkpoint_base_path
 
 from experiments.defaults import default_validation_sets
 from experiments.grug.modular_opt.model import GrugModelConfig
@@ -205,6 +206,7 @@ def run_grug_modular_opt_trial(config: GrugModularOptLaunchConfig) -> None:
         allow_nondivisible_batch_size=False,
         checkpointer=CheckpointerConfig(
             base_path=os.path.join(config.output_path, "checkpoints"),
+            temporary_base_path=temporary_checkpoint_base_path(config.output_path),
             append_run_id_to_base_path=False,
             save_interval=timedelta(minutes=10),
             keep=[{"every": 1000}],

experiments/grug/moe/launch.py

@@ -24,6 +24,7 @@
 from levanter.utils.mesh import MeshConfig
 from marin.execution.executor import ExecutorStep, executor_main, this_output_path, versioned
 from marin.processing.tokenize import add_validation_sets_to_mixture
+from marin.training.training import temporary_checkpoint_base_path
 
 from experiments.defaults import default_validation_sets
 from experiments.grug.moe.heuristic import build_from_heuristic
@@ -92,6 +93,7 @@ def run_grug_moe_trial(config: GrugMoeLaunchConfig) -> None:
         allow_nondivisible_batch_size=False,
         checkpointer=CheckpointerConfig(
             base_path=os.path.join(config.output_path, "checkpoints"),
+            temporary_base_path=temporary_checkpoint_base_path(config.output_path),
             append_run_id_to_base_path=False,
             save_interval=timedelta(minutes=10),
             keep=[{"every": 1000}],

lib/iris/tests/test_marin_fs.py

@@ -21,6 +21,7 @@
     marin_prefix,
     marin_region,
     marin_temp_bucket,
+    marin_temp_bucket_for_prefix,
     open_url,
     region_from_metadata,
     region_from_prefix,
@@ -136,6 +137,18 @@ def test_marin_temp_bucket_from_env_prefix():
         assert marin_temp_bucket(ttl_days=3, prefix="zephyr") == "gs://marin-tmp-us-east1/ttl=3d/zephyr"
 
 
+def test_marin_temp_bucket_for_prefix_uses_source_region():
+    with (
+        patch("rigging.filesystem.urllib.request.urlopen", side_effect=OSError("not on GCP")),
+        patch.dict(os.environ, {"MARIN_PREFIX": "gs://marin-us-central1/scratch"}),
+    ):
+        assert marin_temp_bucket_for_prefix(
+            ttl_days=14,
+            source_prefix="gs://marin-us-east5/experiments/grug/run",
+            prefix="checkpoints-temp/marin-us-east5/experiments/grug/run/checkpoints",
+        ) == ("gs://marin-tmp-us-east5/ttl=14d/" "checkpoints-temp/marin-us-east5/experiments/grug/run/checkpoints")
+
+
 def test_marin_temp_bucket_falls_back_to_marin_prefix_when_no_region():
     # Unknown region in MARIN_PREFIX → no entry in REGION_TO_TMP_BUCKET → falls back to marin_prefix/tmp
     with (

lib/marin/src/marin/training/training.py

@@ -5,9 +5,10 @@
 import importlib
 import logging
 import os
+import urllib.parse
 from copy import deepcopy
-from dataclasses import dataclass, replace
 from collections.abc import Callable
+from dataclasses import dataclass, replace
 from typing import TypeVar
 
 import draccus
@@ -23,7 +24,7 @@
 )
 from mergedeep import mergedeep
 
-from rigging.filesystem import check_gcs_paths_same_region, marin_temp_bucket
+from rigging.filesystem import check_gcs_paths_same_region, marin_temp_bucket, marin_temp_bucket_for_prefix
 from marin.training.run_environment import add_run_env_variables
 
 logger = logging.getLogger(__name__)
@@ -84,12 +85,34 @@ class TrainDpoOnPodConfig:
 
 DEFAULT_CHECKPOINTS_PATH = "checkpoints"
 DEFAULT_HF_CHECKPOINTS_PATH = "hf"
+TEMPORARY_CHECKPOINT_TTL_DAYS = 14
+TEMPORARY_CHECKPOINTS_PATH = "checkpoints-temp"
 
 
 def _cli_helpers_module():
     return importlib.import_module("levanter.infra.cli_helpers")
 
 
+def _output_path_temp_component(output_path: str) -> str:
+    parsed = urllib.parse.urlparse(output_path)
+    if parsed.scheme and parsed.netloc:
+        return f"{parsed.netloc}{parsed.path}".strip("/")
+    if parsed.scheme:
+        return f"{parsed.scheme}{parsed.path}".strip("/")
+    return output_path.strip("/")
+
+
+def temporary_checkpoint_base_path(output_path: str) -> str:
+    """Return the region-local temporary checkpoint base for an executor output path."""
+    output_component = _output_path_temp_component(output_path)
+    temp_prefix = os.path.join(TEMPORARY_CHECKPOINTS_PATH, output_component, DEFAULT_CHECKPOINTS_PATH)
+    return marin_temp_bucket_for_prefix(
+        ttl_days=TEMPORARY_CHECKPOINT_TTL_DAYS,
+        source_prefix=output_path,
+        prefix=temp_prefix,
+    )
+
+
 def _update_config_to_use_out_path(pod_config: TrainOnPodConfigT) -> TrainOnPodConfigT:
     """
     Update the config to use the out_path as the base output directory for training.
@@ -109,7 +132,7 @@ def _update_config_to_use_out_path(pod_config: TrainOnPodConfigT) -> TrainOnPodC
         checkpointer=replace(
             pod_config.train_config.trainer.checkpointer,
             base_path=os.path.join(pod_config.output_path, DEFAULT_CHECKPOINTS_PATH),
-            temporary_base_path=marin_temp_bucket(ttl_days=14, prefix="checkpoints-temp"),
+            temporary_base_path=temporary_checkpoint_base_path(pod_config.output_path),
         ),
     )
 

lib/rigging/src/rigging/filesystem.py

@@ -157,6 +157,12 @@ def marin_region() -> str | None:
     return region_from_metadata() or region_from_prefix(os.environ.get("MARIN_PREFIX", ""))
 
 
+def _append_path_prefix(path: str, prefix: str) -> str:
+    if prefix:
+        return f"{path}/{prefix.strip('/')}"
+    return path
+
+
 def marin_temp_bucket(ttl_days: int, prefix: str = "") -> str:
     """Return a path on region-local temp storage. Never returns ``None``.
 
@@ -186,16 +192,26 @@ def marin_temp_bucket(ttl_days: int, prefix: str = "") -> str:
             bucket = REGION_TO_TMP_BUCKET.get(region)
             if bucket:
                 path = f"gs://{bucket}/ttl={ttl_days}d"
-                if prefix:
-                    path = f"{path}/{prefix.strip('/')}"
-                return path
+                return _append_path_prefix(path, prefix)
 
     if "://" not in mp:
         mp = f"file://{mp}"
     path = f"{mp}/tmp"
-    if prefix:
-        path = f"{path}/{prefix.strip('/')}"
-    return path
+    return _append_path_prefix(path, prefix)
+
+
+def marin_temp_bucket_for_prefix(ttl_days: int, source_prefix: str, prefix: str = "") -> str:
+    """Return temp storage in the same region as ``source_prefix`` when possible.
+
+    This is useful when configuring a remote job from a launcher that may be in
+    a different region than the job output path.
+    """
+    region = region_from_prefix(source_prefix)
+    if region:
+        bucket = REGION_TO_TMP_BUCKET.get(region)
+        if bucket:
+            return _append_path_prefix(f"gs://{bucket}/ttl={ttl_days}d", prefix)
+    return marin_temp_bucket(ttl_days, prefix=prefix)
 
 
 # ---------------------------------------------------------------------------

tests/test_grug_launch_checkpoint_paths.py

@@ -0,0 +1,49 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Any
+from unittest.mock import patch
+
+from fray.cluster import ResourceConfig
+from levanter.optim import AdamConfig
+from levanter.tracker import NoopConfig
+
+from experiments.grug.base.launch import GRUG_130M_MODEL, GrugBaseLaunchConfig, run_grug_base_trial
+
+_DUMMY_DATA: Any = object()
+
+
+def test_grug_base_launch_sets_temporary_checkpoint_base_path():
+    with (
+        patch("rigging.filesystem.urllib.request.urlopen", side_effect=OSError("not on GCP")),
+        patch.dict(os.environ, {"MARIN_PREFIX": "gs://marin-us-central1/scratch"}),
+        patch("experiments.grug.base.launch.run_grug") as run_grug,
+    ):
+        run_grug_base_trial(
+            GrugBaseLaunchConfig(
+                model=GRUG_130M_MODEL,
+                data=_DUMMY_DATA,
+                output_path="gs://marin-us-east5/experiments/grug/base-trial",
+                run_id="grug-temp-path-test",
+                resources=ResourceConfig.with_cpu(),
+                steps=1,
+                batch_size=1,
+                seed=0,
+                mp="params=float32,compute=bfloat16,output=bfloat16",
+                tracker=NoopConfig(),
+                optimizer=AdamConfig(),
+                eval=None,
+            )
+        )
+
+    run_config = run_grug.call_args.args[0]
+    checkpointer = run_config.trainer.trainer.checkpointer
+    assert checkpointer.base_path == "gs://marin-us-east5/experiments/grug/base-trial/checkpoints"
+    assert checkpointer.temporary_base_path == (
+        "gs://marin-tmp-us-east5/ttl=14d/" "checkpoints-temp/marin-us-east5/experiments/grug/base-trial/checkpoints"
+    )
+    assert run_config.trainer.trainer.checkpoint_search_paths("grug-temp-path-test") == [
+        "gs://marin-us-east5/experiments/grug/base-trial/checkpoints",
+        "gs://marin-tmp-us-east5/ttl=14d/" "checkpoints-temp/marin-us-east5/experiments/grug/base-trial/checkpoints",
+    ]

tests/test_training.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import dataclasses
+import os
 from pathlib import Path
 from unittest.mock import patch
 
@@ -14,6 +15,8 @@
 from marin.training.training import (
     TrainLmOnPodConfig,
     _doublecheck_paths,
+    _update_config_to_use_out_path,
+    temporary_checkpoint_base_path,
 )
 
 
@@ -64,6 +67,38 @@ def test_lm_config_with_train_urls_allowed_out_of_region(trainer_config):
         _doublecheck_paths(config)
 
 
+def test_temporary_checkpoint_base_path_follows_output_path_region():
+    with (
+        patch("rigging.filesystem.urllib.request.urlopen", side_effect=OSError("not on GCP")),
+        patch.dict(os.environ, {"MARIN_PREFIX": "gs://marin-us-central1/scratch"}),
+    ):
+        assert temporary_checkpoint_base_path("gs://marin-us-east5/experiments/grug/base-trial") == (
+            "gs://marin-tmp-us-east5/ttl=14d/" "checkpoints-temp/marin-us-east5/experiments/grug/base-trial/checkpoints"
+        )
+
+
+def test_update_config_to_use_out_path_sets_run_specific_temp_checkpoints(trainer_config):
+    with (
+        patch("rigging.filesystem.urllib.request.urlopen", side_effect=OSError("not on GCP")),
+        patch.dict(os.environ, {"MARIN_PREFIX": "gs://marin-us-central1/scratch"}),
+    ):
+        config = TrainLmOnPodConfig(
+            train_config=train_lm.TrainLmConfig(
+                trainer=trainer_config,
+            ),
+            resources=ResourceConfig.with_tpu("v4-8"),
+            output_path="gs://marin-us-east5/experiments/grug/base-trial",
+        )
+
+        updated = _update_config_to_use_out_path(config)
+
+        checkpointer = updated.train_config.trainer.checkpointer
+        assert checkpointer.base_path == "gs://marin-us-east5/experiments/grug/base-trial/checkpoints"
+        assert checkpointer.temporary_base_path == (
+            "gs://marin-tmp-us-east5/ttl=14d/" "checkpoints-temp/marin-us-east5/experiments/grug/base-trial/checkpoints"
+        )
+
+
 def test_recursive_path_checking(trainer_config):
     """Paths are checked recursively in nested structures."""
     with (