Remove uses of FUSE (#2308)

## Description

Remove all uses of gcsfuse etc, replace with either reading direct from
gcs or in one case downloading from gcs to a tmp dir.

Author: dlwh

AGENTS.md

@@ -36,6 +36,7 @@
 - Put all imports at the top of the file. Avoid local imports unless technically necessary (for example, to break circular dependencies or guard optional dependencies).
 - Prefer top-level functions when code does not mutate shared state; use classes to encapsulate data when that improves clarity.
 - Prefer top-level Python tests and fixtures.
+- Disprefer internal mutation of function arguments, especially config dataclasses; prefer returning a modified copy (e.g., via `dataclasses.replace`) so call sites remain predictable and side effects are explicit.
 - Use early returns (`if not x: return None`) when they reduce nesting.
 - Do not introduce ad-hoc compatibility hacks like `hasattr(m, "old_attr")`; update the code consistently instead.
 - Do not use `from future import ...` statements.

experiments/exp1342_gemstones_scaling_law.py

@@ -21,16 +21,15 @@
 Usage:
 1. Import the model step you want to use
 2. Run executor_main([model_step]) to download
-3. Use get_model_local_path(model_step) to get the local path
+3. Use the model step's output path for downstream jobs
 
 Example:
 ```
 from gemstones import gemstone_768x45
-from marin.execution.executor import executor_main
-from experiments.models import get_model_local_path
+from marin.execution.executor import executor_main, output_path_of
 
 executor_main([gemstone_768x45])
-local_path = get_model_local_path(gemstone_768x45)
+model_path = output_path_of(gemstone_768x45)
 ```
 """
 

experiments/models.py

@@ -17,20 +17,16 @@
 Usage:
 1. If you have a model you want to download from huggingface, add the repo name and config in MODEL_NAME_TO_CONFIG.
 2. Run download_model_step(MODEL_NAME_TO_CONFIG[model_name]) to download the model.
-3. Use get_model_local_path(model_name) to get the local path of the model.
 
 Example:
 ```
 model_name = "meta-llama/Llama-3.1-8B-Instruct"
 model_config = MODEL_NAME_TO_CONFIG[model_name]
 download_step = download_model_step(model_config)
 executor_main([download_step])
-
-local_path = get_model_local_path(model_name)
 ```
 """
 
-import os
 from dataclasses import dataclass
 
 from marin.download.huggingface.download_hf import DownloadConfig, download_hf
@@ -44,19 +40,14 @@ class ModelConfig:
     hf_revision: str
 
 
-# We utilize GCSFuse because our disk space is limited on TPUs.
-# This means that for certain large models (e.g. Llama 70B), we will not be able
-# to fit the models on local disk. We use GCSFuse to mount the GCS bucket to the local filesystem
-# to be able to download and use these large models.
-LOCAL_PREFIX = "/opt"
-GCS_FUSE_MOUNT_PATH = "gcsfuse_mount/models"
+MODEL_OUTPUT_SUBDIR = "models"
 
 
 def download_model_step(model_config: ModelConfig) -> ExecutorStep:
     model_name = get_directory_friendly_name(model_config.hf_repo_id)
     model_revision = get_directory_friendly_name(model_config.hf_revision)
     download_step = ExecutorStep(
-        name=f"{GCS_FUSE_MOUNT_PATH}/{model_name}--{model_revision}",
+        name=f"{MODEL_OUTPUT_SUBDIR}/{model_name}--{model_revision}",
         fn=download_hf,
         config=DownloadConfig(
             hf_dataset_id=model_config.hf_repo_id,
@@ -67,17 +58,12 @@ def download_model_step(model_config: ModelConfig) -> ExecutorStep:
         ),
         # must override because it because if we don't then it will end in a hash
         # if it ends in a hash, then we cannot determine the local path
-        override_output_path=f"{GCS_FUSE_MOUNT_PATH}/{model_name}--{model_revision}",
+        override_output_path=f"{MODEL_OUTPUT_SUBDIR}/{model_name}--{model_revision}",
     )
 
     return download_step
 
 
-def get_model_local_path(step: ExecutorStep) -> str:
-    model_repo_name = step.name[len(GCS_FUSE_MOUNT_PATH) + 1 :]
-    return os.path.join(LOCAL_PREFIX, GCS_FUSE_MOUNT_PATH, model_repo_name)
-
-
 smollm2_1_7b_instruct = download_model_step(
     ModelConfig(
         hf_repo_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",

infra/marin-big-run.yaml

@@ -37,6 +37,13 @@ docker:
       - --shm-size=100gb
       - -v
       - "/tmp:/tmp"
+      - -e MARIN_PREFIX=gs://marin-us-central2
+      - -e BUCKET=marin-us-central2
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
       # this lets the worker run docker commands and have them run as sibling containers
       - -v "/var/run/docker.sock:/var/run/docker.sock"
 
@@ -48,6 +55,13 @@ docker:
       - -v "/tmp:/tmp"
       - --ulimit nofile=1048576:1048576
       - -e RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS=64
+      - -e MARIN_PREFIX=gs://marin-us-central2
+      - -e BUCKET=marin-us-central2
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
 
       - -e RAY_AUTH_MODE=token
       - -e RAY_AUTH_TOKEN_PATH=/home/ray/.ray/auth_token
@@ -79,15 +93,9 @@ setup_commands:
   - gcloud secrets versions access latest --secret=RAY_AUTH_TOKEN > $HOME/.ray/auth_token
   - chmod 600 $HOME/.ray/auth_token
 
-  - echo 'export MARIN_PREFIX="gs://marin-us-central2"' >> $HOME/.bashrc
-  - echo 'export BUCKET="marin-us-central2"' >> $HOME/.bashrc
+  - mkdir -p /tmp/marin-cache
   # cf https://github.com/ray-project/ray/blob/0bc6ec86ffd0fc0d4e43fb339ffe0ac03ee5531b/python/ray/autoscaler/_private/constants.py#L66
   # this is set to 30s by default, which is much too short for our use case
-  - echo 'export AUTOSCALER_HEARTBEAT_TIMEOUT_S=600' >> $HOME/.bashrc
-  - echo 'export TPU_MIN_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_STDERR_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_LOG_DIR=disabled' >> $HOME/.bashrc
-  - gcsfuse --implicit-dirs --client-protocol grpc --cache-dir /dev/shm --file-cache-max-size-mb 160000 --only-dir gcsfuse_mount $BUCKET /opt/gcsfuse_mount || true
 
 worker_setup_commands:
   # delete any old ray session data

infra/marin-cluster-template.yaml

@@ -33,6 +33,13 @@ docker:
       - --shm-size=100gb
       - -v
       - "/tmp:/tmp"
+      - -e MARIN_PREFIX=gs://{{BUCKET}}
+      - -e BUCKET={{BUCKET}}
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
       # this lets the worker run docker commands and have them run as sibling containers
       - -v "/var/run/docker.sock:/var/run/docker.sock"
 {% if RAY_AUTH_MODE == "token" %}
@@ -44,6 +51,13 @@ docker:
       - -v "/tmp:/tmp"
       - --ulimit nofile=1048576:1048576
       - -e RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS=64
+      - -e MARIN_PREFIX=gs://{{BUCKET}}
+      - -e BUCKET={{BUCKET}}
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
 {% if RAY_AUTH_MODE == "token" %}
       - -e RAY_AUTH_MODE=token
       - -e RAY_AUTH_TOKEN_PATH=/home/ray/.ray/auth_token
@@ -75,15 +89,9 @@ setup_commands:
   - gcloud secrets versions access latest --secret={{ RAY_AUTH_SECRET }} > $HOME/.ray/auth_token
   - chmod 600 $HOME/.ray/auth_token
 {% endif %}
-  - echo 'export MARIN_PREFIX="gs://{{BUCKET}}"' >> $HOME/.bashrc
-  - echo 'export BUCKET="{{BUCKET}}"' >> $HOME/.bashrc
+  - mkdir -p /tmp/marin-cache
   # cf https://github.com/ray-project/ray/blob/0bc6ec86ffd0fc0d4e43fb339ffe0ac03ee5531b/python/ray/autoscaler/_private/constants.py#L66
   # this is set to 30s by default, which is much too short for our use case
-  - echo 'export AUTOSCALER_HEARTBEAT_TIMEOUT_S=600' >> $HOME/.bashrc
-  - echo 'export TPU_MIN_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_STDERR_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_LOG_DIR=disabled' >> $HOME/.bashrc
-  - gcsfuse --implicit-dirs --client-protocol grpc --cache-dir /dev/shm --file-cache-max-size-mb 160000 --only-dir gcsfuse_mount $BUCKET /opt/gcsfuse_mount || true
 
 worker_setup_commands:
   # delete any old ray session data

infra/marin-eu-west4-a.yaml

@@ -37,6 +37,13 @@ docker:
       - --shm-size=100gb
       - -v
       - "/tmp:/tmp"
+      - -e MARIN_PREFIX=gs://marin-eu-west4
+      - -e BUCKET=marin-eu-west4
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
       # this lets the worker run docker commands and have them run as sibling containers
       - -v "/var/run/docker.sock:/var/run/docker.sock"
 
@@ -48,6 +55,13 @@ docker:
       - -v "/tmp:/tmp"
       - --ulimit nofile=1048576:1048576
       - -e RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS=64
+      - -e MARIN_PREFIX=gs://marin-eu-west4
+      - -e BUCKET=marin-eu-west4
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
 
       - -e RAY_AUTH_MODE=token
       - -e RAY_AUTH_TOKEN_PATH=/home/ray/.ray/auth_token
@@ -79,15 +93,9 @@ setup_commands:
   - gcloud secrets versions access latest --secret=RAY_AUTH_TOKEN > $HOME/.ray/auth_token
   - chmod 600 $HOME/.ray/auth_token
 
-  - echo 'export MARIN_PREFIX="gs://marin-eu-west4"' >> $HOME/.bashrc
-  - echo 'export BUCKET="marin-eu-west4"' >> $HOME/.bashrc
+  - mkdir -p /tmp/marin-cache
   # cf https://github.com/ray-project/ray/blob/0bc6ec86ffd0fc0d4e43fb339ffe0ac03ee5531b/python/ray/autoscaler/_private/constants.py#L66
   # this is set to 30s by default, which is much too short for our use case
-  - echo 'export AUTOSCALER_HEARTBEAT_TIMEOUT_S=600' >> $HOME/.bashrc
-  - echo 'export TPU_MIN_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_STDERR_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_LOG_DIR=disabled' >> $HOME/.bashrc
-  - gcsfuse --implicit-dirs --client-protocol grpc --cache-dir /dev/shm --file-cache-max-size-mb 160000 --only-dir gcsfuse_mount $BUCKET /opt/gcsfuse_mount || true
 
 worker_setup_commands:
   # delete any old ray session data

infra/marin-eu-west4-vllm.yaml

@@ -34,6 +34,13 @@ docker:
       - --shm-size=200gb
       - -v
       - "/tmp:/tmp"
+      - -e MARIN_PREFIX=gs://marin-eu-west4
+      - -e BUCKET=marin-eu-west4
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
       # this lets the worker run docker commands and have them run as sibling containers
       - -v "/var/run/docker.sock:/var/run/docker.sock"
 
@@ -45,6 +52,13 @@ docker:
       - -v "/tmp:/tmp"
       - --ulimit nofile=1048576:1048576
       - -e RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS=64
+      - -e MARIN_PREFIX=gs://marin-eu-west4
+      - -e BUCKET=marin-eu-west4
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
 
       - -e RAY_AUTH_MODE=token
       - -e RAY_AUTH_TOKEN_PATH=/home/ray/.ray/auth_token
@@ -76,15 +90,9 @@ setup_commands:
   - gcloud secrets versions access latest --secret=RAY_AUTH_TOKEN > $HOME/.ray/auth_token
   - chmod 600 $HOME/.ray/auth_token
 
-  - echo 'export MARIN_PREFIX="gs://marin-eu-west4"' >> $HOME/.bashrc
-  - echo 'export BUCKET="marin-eu-west4"' >> $HOME/.bashrc
+  - mkdir -p /tmp/marin-cache
   # cf https://github.com/ray-project/ray/blob/0bc6ec86ffd0fc0d4e43fb339ffe0ac03ee5531b/python/ray/autoscaler/_private/constants.py#L66
   # this is set to 30s by default, which is much too short for our use case
-  - echo 'export AUTOSCALER_HEARTBEAT_TIMEOUT_S=600' >> $HOME/.bashrc
-  - echo 'export TPU_MIN_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_STDERR_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_LOG_DIR=disabled' >> $HOME/.bashrc
-  - gcsfuse --implicit-dirs --client-protocol grpc --cache-dir /dev/shm --file-cache-max-size-mb 160000 --only-dir gcsfuse_mount $BUCKET /opt/gcsfuse_mount || true
 
 worker_setup_commands:
   # delete any old ray session data

infra/marin-eu-west4.yaml

@@ -37,6 +37,13 @@ docker:
       - --shm-size=100gb
       - -v
       - "/tmp:/tmp"
+      - -e MARIN_PREFIX=gs://marin-eu-west4
+      - -e BUCKET=marin-eu-west4
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
       # this lets the worker run docker commands and have them run as sibling containers
       - -v "/var/run/docker.sock:/var/run/docker.sock"
 
@@ -48,6 +55,13 @@ docker:
       - -v "/tmp:/tmp"
       - --ulimit nofile=1048576:1048576
       - -e RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS=64
+      - -e MARIN_PREFIX=gs://marin-eu-west4
+      - -e BUCKET=marin-eu-west4
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
 
       - -e RAY_AUTH_MODE=token
       - -e RAY_AUTH_TOKEN_PATH=/home/ray/.ray/auth_token
@@ -79,15 +93,9 @@ setup_commands:
   - gcloud secrets versions access latest --secret=RAY_AUTH_TOKEN > $HOME/.ray/auth_token
   - chmod 600 $HOME/.ray/auth_token
 
-  - echo 'export MARIN_PREFIX="gs://marin-eu-west4"' >> $HOME/.bashrc
-  - echo 'export BUCKET="marin-eu-west4"' >> $HOME/.bashrc
+  - mkdir -p /tmp/marin-cache
   # cf https://github.com/ray-project/ray/blob/0bc6ec86ffd0fc0d4e43fb339ffe0ac03ee5531b/python/ray/autoscaler/_private/constants.py#L66
   # this is set to 30s by default, which is much too short for our use case
-  - echo 'export AUTOSCALER_HEARTBEAT_TIMEOUT_S=600' >> $HOME/.bashrc
-  - echo 'export TPU_MIN_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_STDERR_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_LOG_DIR=disabled' >> $HOME/.bashrc
-  - gcsfuse --implicit-dirs --client-protocol grpc --cache-dir /dev/shm --file-cache-max-size-mb 160000 --only-dir gcsfuse_mount $BUCKET /opt/gcsfuse_mount || true
 
 worker_setup_commands:
   # delete any old ray session data

infra/marin-us-central1-vllm.yaml

@@ -34,6 +34,13 @@ docker:
       - --shm-size=200gb
       - -v
       - "/tmp:/tmp"
+      - -e MARIN_PREFIX=gs://marin-us-central1
+      - -e BUCKET=marin-us-central1
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
       # this lets the worker run docker commands and have them run as sibling containers
       - -v "/var/run/docker.sock:/var/run/docker.sock"
 
@@ -45,6 +52,13 @@ docker:
       - -v "/tmp:/tmp"
       - --ulimit nofile=1048576:1048576
       - -e RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS=64
+      - -e MARIN_PREFIX=gs://marin-us-central1
+      - -e BUCKET=marin-us-central1
+      - -e MARIN_LOCAL_CACHE_DIR=/tmp/marin-cache
+      - -e AUTOSCALER_HEARTBEAT_TIMEOUT_S=600
+      - -e TPU_MIN_LOG_LEVEL=3
+      - -e TPU_STDERR_LOG_LEVEL=3
+      - -e TPU_LOG_DIR=disabled
 
       - -e RAY_AUTH_MODE=token
       - -e RAY_AUTH_TOKEN_PATH=/home/ray/.ray/auth_token
@@ -76,15 +90,9 @@ setup_commands:
   - gcloud secrets versions access latest --secret=RAY_AUTH_TOKEN > $HOME/.ray/auth_token
   - chmod 600 $HOME/.ray/auth_token
 
-  - echo 'export MARIN_PREFIX="gs://marin-us-central1"' >> $HOME/.bashrc
-  - echo 'export BUCKET="marin-us-central1"' >> $HOME/.bashrc
+  - mkdir -p /tmp/marin-cache
   # cf https://github.com/ray-project/ray/blob/0bc6ec86ffd0fc0d4e43fb339ffe0ac03ee5531b/python/ray/autoscaler/_private/constants.py#L66
   # this is set to 30s by default, which is much too short for our use case
-  - echo 'export AUTOSCALER_HEARTBEAT_TIMEOUT_S=600' >> $HOME/.bashrc
-  - echo 'export TPU_MIN_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_STDERR_LOG_LEVEL=3' >> $HOME/.bashrc
-  - echo 'export TPU_LOG_DIR=disabled' >> $HOME/.bashrc
-  - gcsfuse --implicit-dirs --client-protocol grpc --cache-dir /dev/shm --file-cache-max-size-mb 160000 --only-dir gcsfuse_mount $BUCKET /opt/gcsfuse_mount || true
 
 worker_setup_commands:
   # delete any old ray session data