adding integration with fsspec FileSystem by artemru · Pull Request #1126 · facebookresearch/fairseq2

artemru · 2025-04-18T13:12:56Z

What does this PR do? Please describe:

This PR allows to save and reload the checkpoints from an remote location using fsspec integration!

Wrapping the standard fsspec interface to fs2 FileSystem API

Using Wrapped FS instead of Native one
Adding some basic path resolvers (they should be used in dataset loading and model registry)
Adding some basic configurations for s3, hf, gcp
adding cat method (for possible faster files reading in concrete fsspec implementations)
chaning is_local to is_local_path
introducing checkpoint-dir support in CLI
fixing moving files logic for hf and normal checkpoints (s3 cannot do move operation as on nfs)

Tested with e2e example "

python -m recipes.lm.sft "/tmp/tmp/fs2_s3test_1" --checkpoint-dir s3://bucket/folder/ \
  --config-file recipes/lm/sft/configs/llama3_2_1b_instruct_gsm8k.yaml \
  --config regime.checkpoint_every_n_steps=200

with restarting as well!

Does your PR introduce any breaking changes? If yes, please list them:
List of all backwards-incompatible changes.

Check list:

Was the content of this PR discussed and approved via a GitHub issue? (no need for typos or documentation improvements)
Did you read the contributor guideline?
Did you make sure that your PR does only one thing instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests?
Did you verify new and existing tests pass locally with your changes?
Did you update the CHANGELOG? (no need for typos, documentation, or minor internal changes)

src/fairseq2/setup/_root.py

artemru · 2026-01-26T11:50:41Z

src/fairseq2/file_system.py

    def open(self, path: Path, mode: FileMode = FileMode.READ) -> BinaryIO: ...

+    @abstractmethod
+    def cat(self, path: Path) -> bytes:


it's useful for faster files downloading

cbalioglu

I believe one major thing missing right now is some benchmark numbers as well as verification that the new file system implementation works end-to-end (i.e. a real world training run, asset metadata loading, Hugging Face model exports, and any other place where we used regular file system calls work without any regressions). Have you been able to verify those?

cbalioglu · 2026-01-26T12:40:17Z

src/fairseq2/assets/download_manager.py

+    Asset download manager for local files and fsspec-compatible URI schemes.
+
+    For local files (file:// scheme), it returns the local path directly.
+    For fsspec schemes (s3://, gs://, etc.), it returns the URI as a Path,


Not really sure whether we should treat schemes beyond file:// as "local" in a download manager. It defeats the purpose of download manager if we do not download and cache a remote file.

what is the exact purpose of that class btw ?
or it's a pure naming issue LocalAssetDownloadManager -> GlobalAssetDownloadManager ?

cbalioglu · 2026-01-26T12:42:24Z

src/fairseq2/checkpoint/manager.py

-        if not self._file_system.is_local:
-            return
-
+    def _sync_distributed_writes(self) -> None:


This whole function can become noop in case the file system is not local (I assume fsspec guarantees some form of strong consistency, if not, this function will likely become tricky). We do not need a redundant barrier call.

just wanted to confirm that

yes, i think s3 should get some transactional consistency indeed

https://aws.amazon.com/blogs/aws/amazon-s3-update-strong-read-after-write-consistency/

src/fairseq2/file_system.py

src/fairseq2/checkpoint/manager.py

caciolai

Some minor comments

artemru · 2026-01-26T18:27:59Z

I believe one major thing missing right now is some benchmark numbers as well as verification that the new file system implementation works end-to-end (i.e. a real world training run, asset metadata loading, Hugging Face model exports, and any other place where we used regular file system calls work without any regressions). Have you been able to verify those?

i verified and continue doing so on real run where the chkpt happening on s3.
is there any e2e workflow that I can reuse for tests that checks that all saved primitives are correct ?

artemru · 2026-01-26T18:31:14Z

i need to move checkpoints/model.yaml to the remove location as well...
Ok, done, model def + hg models :

now we've locally

0       /tmp/tmp/fs2_s3test_12/ws_1.d2b3ae4f/cache
4.0K    /tmp/tmp/fs2_s3test_12/ws_1.d2b3ae4f/config.yaml
192K    /tmp/tmp/fs2_s3test_12/ws_1.d2b3ae4f/logs
212K    /tmp/tmp/fs2_s3test_12/ws_1.d2b3ae4f/metrics
312K    /tmp/tmp/fs2_s3test_12/ws_1.d2b3ae4f/tb

and in s3 folder:
model.yaml step_1000 step_1200 step_1400 step_1600 step_1800 step_2000 step_2200 step_2400 step_2600 step_800
with inner structure:

8.8M    fs2_s3test_12/ws_1.d2b3ae4f/step_1200/data_reader
4.7G    fs2_s3test_12/ws_1.d2b3ae4f/step_1200/hg
4.0K    fs2_s3test_12/ws_1.d2b3ae4f/step_1200/hg.run
4.0K    fs2_s3test_12/ws_1.d2b3ae4f/step_1200/hg.stderr
0       fs2_s3test_12/ws_1.d2b3ae4f/step_1200/hg.stdout
5.6G    fs2_s3test_12/ws_1.d2b3ae4f/step_1200/model
9.3G    fs2_s3test_12/ws_1.d2b3ae4f/step_1200/optimizer
12K     fs2_s3test_12/ws_1.d2b3ae4f/step_1200/trainer

artemru · 2026-01-26T18:58:01Z

the issue is that s3fs removes folder content but does not remove the folders themselves ...

artemru · 2026-01-26T19:04:53Z

src/fairseq2/checkpoint/__init__.py

+
+
+@final
+class CheckpointDir:


what is the right place for this dubby holder class ? maybe we can even remove it

artemru · 2026-01-26T19:35:39Z

@cbalioglu : what tests should run to make sure that it's ok ?

artemru · 2026-01-26T23:23:33Z

ok. catching up hg models folder for copy

artemru · 2026-01-26T23:29:01Z

src/fairseq2/checkpoint/hg.py

+                            cmd, stdout=out_fp, stderr=err_fp
+                        )
+                else:
+                    result = self._process_runner.run_text(cmd, capture_output=True)


generally speaking, we cannot pass self._file_system into subprocess, i dont know that is the best option here ?

i cannot pass env={} here become varenv can be useful to reinit s3 filesystem from aws related varenvs

artemru · 2026-01-26T23:34:33Z

src/fairseq2/file_system.py

+    def remove_directory(self, path: Path) -> None:
+        # For S3/remote filesystems, rmdir only removes empty directories
+        # Use rm with recursive=True to delete directory and all contents
+        self._fsspec.rm(self.get_short_uri(path), recursive=True)


there's an issue with phantoms s3 directories that cannot be removed...

artemru · 2026-01-26T23:38:24Z

src/fairseq2/recipe/internal/metric_recorders.py

-            fp = self._file_system.open_text(run_id_file)
-
-            return fp.read()
+            return self._file_system.cat(run_id_file).decode("utf-8")


i still dont understand why we have filesystem abstraction here if we cannot really use other remote fs for the metrics logging...

artemru · 2026-01-26T23:39:17Z

src/fairseq2/recipe/internal/train_model.py

+        metadata_dumper: ModelMetadataDumper,
+        output_dir: Path,
+        gangs: Gangs,
+        checkpoint_dir: Path | None = None,


i dont know where else we need to propagate checkpoint_dir ...

artemru · 2026-01-26T23:54:47Z

src/fairseq2/checkpoint/hg.py

-                    )
-
-                    result = self._process_runner.run_text(
-                        cmd, stdout=out_fp, stderr=err_fp, env={}


do you know why env={} was necessary here ?

This reverts commit 5358729.

This reverts commit 024d862.

artemru · 2026-02-06T12:48:24Z

some rebase issues - fixing them

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 18, 2025

artemru self-assigned this Apr 18, 2025

artemru force-pushed the adding_fsspec branch from 6b11ad9 to 171ab45 Compare April 23, 2025 10:45

artemru marked this pull request as ready for review April 23, 2025 13:10

artemru requested a review from cbalioglu as a code owner April 23, 2025 13:10

artemru commented Apr 23, 2025

View reviewed changes

src/fairseq2/setup/_root.py Outdated Show resolved Hide resolved

artemru force-pushed the adding_fsspec branch from 3245a68 to 4508cc2 Compare July 7, 2025 16:09

caciolai requested review from MartinGleize, cirquit and zyaoj as code owners January 23, 2026 11:39

artemru commented Jan 26, 2026

View reviewed changes

cbalioglu reviewed Jan 26, 2026

View reviewed changes

caciolai reviewed Jan 26, 2026

View reviewed changes

src/fairseq2/checkpoint/manager.py Outdated Show resolved Hide resolved

caciolai reviewed Jan 26, 2026

View reviewed changes

artemru commented Jan 26, 2026

View reviewed changes

Artyom Kozhevnikov added 5 commits February 6, 2026 12:35

adding integration with fsspec FileSystem

9bb5315

global fs

4ab1b83

adding registery

c17a637

rm RuntimeContext circular import

997716a

using fsspec registery

f83cf85

artemru added 21 commits February 6, 2026 12:40

adding cat

d793cc2

mypy

ca694a8

fixing Path artifacts s3:/

93437af

local files

3a7d001

extra test

9c82e48

black

d75cabc

comment

bdeef30

trying to do e2e run

f9f4eeb

pushing more logging

383e82e

Revert "pushing more logging"

3ef868b

This reverts commit 5358729.

Revert "trying to do e2e run"

f95b651

This reverts commit 024d862.

introduce chekpoint_dir

805a3f4

move model.yaml to remote location

ffe77ba

black

bf84714

taking comments

4b73aa7

some doc

abaac2d

hg move

4df34f4

hg copy to remote

9641f87

rf comm

2126c8c

ff

1096def

env back

1d0203f

artemru force-pushed the adding_fsspec branch from 9a06bc7 to 1d0203f Compare February 6, 2026 12:42

artemru and others added 7 commits February 6, 2026 13:48

Merge branch 'main' into adding_fsspec

db57fde

ff

c53b18c

linters

796cb73

ff

7a66cbb

s3 only protocol

68179dd

sort

da75d83

revert 2f

b3ac484

Conversation

artemru commented Apr 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cbalioglu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

caciolai left a comment

Choose a reason for hiding this comment

Uh oh!

artemru commented Jan 26, 2026

Uh oh!

artemru commented Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

artemru commented Jan 26, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

artemru commented Jan 26, 2026

Uh oh!

artemru commented Jan 26, 2026

Uh oh!

artemru Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

artemru commented Feb 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

artemru commented Apr 18, 2025 •

edited

Loading

artemru commented Jan 26, 2026 •

edited

Loading

artemru Jan 26, 2026 •

edited

Loading