fray: rename fray.v2.* to fray.* (#4453) (#5140)

## Summary

Stage 3i (final code-side stage) of the Ray removal (parent #4453). With
v1 gone (#5137) and the Ray backend gone (#5138), `fray.v2` is the only
surviving version — drop the `v2` suffix and move the modules up to the
`fray` root.

## What this PR does

- `git mv
lib/fray/src/fray/v2/{actor,client,device_flops,iris_backend,local_backend,types}.py`
→ `lib/fray/src/fray/` (6 files, history preserved)
- Collapse the `fray.v2/__init__.py` content into `fray/__init__.py`;
delete the v2 shim
- Rewrite `fray.v2.X` → `fray.X` across 76 files (code, tests, docs,
`AGENTS.md`, `.pyrefly-baseline.json`)
- Rename `lib/fray/tests/test_v2_*.py` → `lib/fray/tests/test_*.py` (4
files)
- No runtime behavior changes — this is a pure rename

## What stays unchanged

- `fray.cluster` public API — `from fray.cluster import ResourceConfig`
still works for the ~60 external call sites; the shim now re-exports
from `fray.types` instead of `fray.v2.types`
- Top-level `fray` exports (already the canonical public interface
before this PR)
- `lib/fray/pyproject.toml` `[tool.hatch.build.targets.wheel]` (already
`packages = ["src/fray"]`)

## Post-rename doc + CI cleanup

Follow-up commit on this PR scrubs stale Ray references that the three
prior stage-3 PRs left behind and simplifies the levanter CI matrix. No
runtime behavior changes.

Comment/docstring updates (6):
1. `lib/levanter/src/levanter/infra/docker.py` — "(if set by ray runtime
env vars)" → "(set by the orchestrator, e.g. Iris/Fray)"
2. `lib/levanter/src/levanter/main/train_lm.py` — "(as happens w/ ray)"
→ "(as happens under Iris/Fray)"
3. `lib/levanter/src/levanter/main/eval_lm.py` — drop ray-specific
framing from the manual-`finish()` comment
4. `lib/levanter/src/levanter/data/sharded_datasource.py` —
`num_cpus`/`num_gpus` docstring no longer says "passed to ray" (params
remain — still wired through `_BatchMapTransform`); nearby comment about
HF dataset multiprocessing rephrased to drop Ray mention
5.
`lib/marin/src/marin/processing/classification/deduplication/connected_components.py`
— `zephyr/ray/pyarrow` → `zephyr/pyarrow`
6. `lib/marin/src/marin/training/training.py` — `run_levanter_train_lm`
docstring rewritten to describe Fray submission (was describing
`auto-ray-start`/`auto-worker-start` flags that no longer exist)

Function rename (1):
7. `lib/marin/src/marin/tokenize/slice_cache.py` — `_slice_cache_in_ray`
→ `_slice_cache_entrypoint` (only caller is the `ExecutorStep.fn=`
reference in the same module)

CI simplification (1):
8. `.github/workflows/levanter-tests.yaml` — drop the
`levanter-ray-tests` job and the `-m "not ray"` filter on the main + TPU
levanter jobs; remove the `ray` pytest marker registration in
`lib/levanter/pyproject.toml` and the two `@pytest.mark.ray` decorators.
Only one live ray-marked test remained
(`test_chat_dataset_build_and_pack`, ~9s, not slow — folded into the
main suite). The other (`test_hf_audio_ray_pipeline`) is
`@pytest.mark.skip` anyway. Also scrubs leftover `not ray` filters from
`lib/levanter/AGENTS.md` and `lib/levanter/scripts/launch_tpu_tests.sh`.

## Verification

- [x] `./infra/pre-commit.py --all-files --fix` → OK
- [x] `uv run --with pyrefly pyrefly check` → 150 errors / 152
suppressed (identical to `origin/main`, no regressions; the two
`fray/types.py` entries are already baselined at the new path)
- [x] `uv run --directory lib/fray --group fray_test pytest tests -x
--timeout=60` → 57 passed (matches post-3g count)
- [x] `uv lock --check` → clean
- [x] `test_chat_dataset_build_and_pack` passes standalone in ~9s
(confirming the CI fold-in is cheap)
- [x] No remaining `fray.v2` references outside `.agents/projects/`
(historical design docs; left untouched, matching stage 3f's handling)

## Scope

- 11 file renames + 76 content edits = 87 files total. 155 insertions /
155 deletions on content plus 73 deletions from the collapsed v2 shim.
- Cleanup commit: 13 files, +19/-51.

## Next steps

Per parent #4453:
- §2 (RAY_* secrets) + §3 (`marin_cluster*` digest cleanup) are parked
on the `marin-big-run` Ray cluster retirement
- Close #4453 when those two land

---------

Co-authored-by: Romain Yon <1596570+yonromai@users.noreply.github.com>

Author: yonromai

.github/workflows/levanter-tests.yaml

@@ -58,36 +58,7 @@ jobs:
       - name: Test with pytest
         run: |
           # Test with specific JAX version, excluding TPU tests
-          PYTHONPATH=tests:src:. uv run --package marin-levanter --frozen --with "jax[cpu]==0.8.0" pytest tests -m "not entry and not slow and not ray and not tpu" --durations=20
-
-  levanter-ray-tests:
-    needs: changes
-    if: needs.changes.outputs.should_run == 'true'
-    runs-on: ubuntu-latest
-    defaults:
-      run:
-        working-directory: lib/levanter
-
-    steps:
-      - uses: actions/checkout@v5
-      - name: Install uv and Python
-        uses: astral-sh/setup-uv@v6
-        with:
-          version: "0.7.20"
-          python-version: "3.11"
-          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
-        run: uv sync --package marin-levanter --dev --group test --frozen
-      - name: Test with pytest
-        run: |
-          PYTHONPATH=tests:src:. uv run --package marin-levanter --frozen pytest tests -m "ray" --durations=20
+          PYTHONPATH=tests:src:. uv run --package marin-levanter --frozen --with "jax[cpu]==0.8.0" pytest tests -m "not entry and not slow and not tpu" --durations=20
 
   levanter-entry-tests:
     needs: changes
@@ -241,7 +212,7 @@ jobs:
               timeout --kill-after=5 --signal=TERM 890 \
               uv run --package marin-levanter --frozen --group test --with 'jax[tpu]==$JAX_VERSION' \
               pytest -n 0 lib/levanter/tests \
-              -m 'not entry and not ray and not slow and not torch' \
+              -m 'not entry 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 \

.pyrefly-baseline.json

@@ -5,7 +5,7 @@
       "column": 9,
       "stop_line": 508,
       "stop_column": 22,
-      "path": "lib/fray/src/fray/v2/types.py",
+      "path": "lib/fray/src/fray/types.py",
       "code": -2,
       "name": "invalid-annotation",
       "description": "`Self` cannot be used in a static method",
@@ -17,7 +17,7 @@
       "column": 9,
       "stop_line": 516,
       "stop_column": 20,
-      "path": "lib/fray/src/fray/v2/types.py",
+      "path": "lib/fray/src/fray/types.py",
       "code": -2,
       "name": "invalid-annotation",
       "description": "`Self` cannot be used in a static method",

docs/debug-log-zephyr-coordinator-thread-shutdown.md

@@ -19,7 +19,7 @@ The real `ZephyrContext.execute()` teardown path leaves the coordinator actor's
 ## Changes to make
 
 - Inspect `_run_coordinator_job()` in `lib/zephyr/src/zephyr/execution.py`
-- Inspect `LocalClient.host_actor()` in `lib/fray/src/fray/v2/local_backend.py`
+- Inspect `LocalClient.host_actor()` in `lib/fray/src/fray/local_backend.py`
 - Add a regression test in `lib/zephyr/tests/test_execution.py` that exercises `execute()` and waits for `zephyr-coordinator-loop` to disappear
 
 ## Future Work

docs/references/resource-config.md

@@ -21,20 +21,20 @@ gpu_auto = ResourceConfig.with_gpu()  # auto-detect GPU type
 cpu_config = ResourceConfig.with_cpu()
 ```
 
-::: fray.v2.types.ResourceConfig
+::: fray.types.ResourceConfig
 
 ## Device Configurations
 
 These are the underlying device types wrapped by `ResourceConfig`:
 
 ### CPU
 
-::: fray.v2.types.CpuConfig
+::: fray.types.CpuConfig
 
 ### GPU
 
-::: fray.v2.types.GpuConfig
+::: fray.types.GpuConfig
 
 ### TPU
 
-::: fray.v2.types.TpuConfig
+::: fray.types.TpuConfig

experiments/dedup/poc_nemotron.py

@@ -8,7 +8,7 @@
 import os
 from typing import TypeVar
 
-from fray.v2 import ResourceConfig
+from fray import ResourceConfig
 from rigging.filesystem import marin_temp_bucket, region_from_metadata, check_path_in_region
 
 from marin.datakit.normalize import NormalizedData, normalize_step

experiments/defaults.py

@@ -14,7 +14,7 @@
 from typing import Any
 
 import jmp
-from fray.v2 import ResourceConfig
+from fray import ResourceConfig
 from marin.execution.remote import remote
 from haliax.partitioning import ResourceAxis
 from haliax.quantization import QuantizationConfig

experiments/exp_model_perplexity_gap_fineweb2_multilingual.py

@@ -1,7 +1,7 @@
 # Copyright The Marin Authors
 # SPDX-License-Identifier: Apache-2.0
 
-from fray.v2.types import ResourceConfig
+from fray.types import ResourceConfig
 
 from experiments.defaults import default_raw_validation_sets
 from experiments.evals.fineweb2_multilingual import fineweb2_multilingual_raw_validation_sets

experiments/exp_model_perplexity_gap_long_tail_runnable.py

@@ -6,7 +6,7 @@
 See https://github.com/marin-community/marin/issues/5005.
 """
 
-from fray.v2.types import ResourceConfig
+from fray.types import ResourceConfig
 
 from experiments.evals.long_tail_ppl_runnable import runnable_long_tail_raw_validation_sets
 from marin.evaluation.perplexity_gap import GapFinderModelConfig, default_model_perplexity_gap

experiments/exp_model_perplexity_gap_marin_vs_llama.py

@@ -1,7 +1,7 @@
 # Copyright The Marin Authors
 # SPDX-License-Identifier: Apache-2.0
 
-from fray.v2.types import ResourceConfig
+from fray.types import ResourceConfig
 
 from experiments.defaults import default_raw_validation_sets
 from marin.evaluation.perplexity_gap import (

experiments/grug/dispatch.py

@@ -9,9 +9,9 @@
 from typing import TypeVar
 
 from fray.cluster import ResourceConfig
-from fray.v2.client import current_client
-from fray.v2.types import Entrypoint, JobRequest, create_environment
-from fray.v2.types import GpuConfig, TpuConfig
+from fray.client import current_client
+from fray.types import Entrypoint, JobRequest, create_environment
+from fray.types import GpuConfig, TpuConfig
 
 logger = logging.getLogger(__name__)