Add support for fsspec paths for CSVLoggers

src/lightning/fabric/CHANGELOG.md

@@ -18,6 +18,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for automatically calling `set_epoch` on the `dataloader.batch_sampler.sampler` ([#16841](https://github.com/Lightning-AI/lightning/pull/16841))
 
 
+- Added support for writing logs remote file systems on `CSVLoggers`. ([#16880](https://github.com/Lightning-AI/lightning/pull/16880))
+
 ### Changed
 
 - Fabric now chooses `accelerator="auto", strategy="auto", devices="auto"` as defaults ([#16842](https://github.com/Lightning-AI/lightning/pull/16842))

src/lightning/fabric/loggers/csv_logs.py

@@ -18,6 +18,7 @@
 from argparse import Namespace
 from typing import Any, Dict, List, Optional, Union
 
+import fsspec
 from torch import Tensor
 
 from lightning.fabric.loggers.logger import Logger, rank_zero_experiment
@@ -62,7 +63,8 @@ def __init__(
         flush_logs_every_n_steps: int = 100,
     ):
         super().__init__()
-        self._root_dir = os.fspath(root_dir)
+        self.fs, _ = fsspec.core.url_to_fs(str(root_dir))
+        self._root_dir = str(root_dir)
         self._name = name or ""
         self._version = version
         self._prefix = prefix
@@ -150,13 +152,14 @@ def finalize(self, status: str) -> None:
     def _get_next_version(self) -> int:
         root_dir = self.root_dir
 
-        if not os.path.isdir(root_dir):
+        if not self.fs.isdir(root_dir):
             log.warning("Missing logger folder: %s", root_dir)
             return 0
 
         existing_versions = []
-        for d in os.listdir(root_dir):
-            if os.path.isdir(os.path.join(root_dir, d)) and d.startswith("version_"):
+        for d in self.fs.listdir(root_dir, detail=False):
+            d = d[len(root_dir) + 1 :]  # removes parent directories
+            if self.fs.isdir(os.path.join(root_dir, d)) and d.startswith("version_"):
                 existing_versions.append(int(d.split("_")[1]))
 
         if len(existing_versions) == 0:
@@ -178,13 +181,14 @@ class _ExperimentWriter:
     def __init__(self, log_dir: str) -> None:
         self.metrics: List[Dict[str, float]] = []
 
+        self.fs, _ = fsspec.core.url_to_fs(log_dir)
         self.log_dir = log_dir
-        if os.path.exists(self.log_dir) and os.listdir(self.log_dir):
+        if self.fs.exists(self.log_dir) and self.fs.listdir(self.log_dir):
             rank_zero_warn(
                 f"Experiment logs directory {self.log_dir} exists and is not empty."
                 " Previous log files in this directory will be deleted when the new ones are saved!"
             )
-        os.makedirs(self.log_dir, exist_ok=True)
+        self.fs.makedirs(self.log_dir, exist_ok=True)
 
         self.metrics_file_path = os.path.join(self.log_dir, self.NAME_METRICS_FILE)
 
@@ -213,7 +217,7 @@ def save(self) -> None:
             last_m.update(m)
         metrics_keys = list(last_m.keys())
 
-        with open(self.metrics_file_path, "w", newline="") as f:
+        with self.fs.open(self.metrics_file_path, "w", newline="") as f:
             writer = csv.DictWriter(f, fieldnames=metrics_keys)
             writer.writeheader()
             writer.writerows(self.metrics)

src/lightning/pytorch/CHANGELOG.md

@@ -53,6 +53,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added a `Trainer(barebones=True)` argument where all features that may impact raw speed are disabled ([#16854](https://github.com/Lightning-AI/lightning/pull/16854))
 
+
+- Added support for writing logs remote file systems on `CSVLoggers`. ([#16880](https://github.com/Lightning-AI/lightning/pull/16880))
+
 ### Changed
 
 

src/lightning/pytorch/loggers/csv_logs.py

@@ -154,6 +154,6 @@ def experiment(self) -> _FabricExperimentWriter:
         if self._experiment is not None:
             return self._experiment
 
-        os.makedirs(self.root_dir, exist_ok=True)
+        self.fs.makedirs(self.root_dir, exist_ok=True)
         self._experiment = ExperimentWriter(log_dir=self.log_dir)
         return self._experiment

tests/tests_pytorch/loggers/test_csv.py

@@ -14,6 +14,7 @@
 import os
 from unittest.mock import MagicMock
 
+import fsspec
 import pytest
 import torch
 
@@ -119,6 +120,20 @@ def test_fit_csv_logger(tmpdir):
     assert os.path.isfile(metrics_file)
 
 
+@RunIf(sklearn=True)
+def test_fit_csv_logger_remotefs():
+    dm = ClassifDataModule()
+    model = ClassificationModel()
+    logger = CSVLogger(save_dir="memory://test_fit_csv_logger_remotefs")
+    trainer = Trainer(
+        default_root_dir="memory://test_fit_csv_logger_remotefs", max_steps=10, logger=logger, log_every_n_steps=1
+    )
+    trainer.fit(model, datamodule=dm)
+    metrics_file = os.path.join(logger.log_dir, ExperimentWriter.NAME_METRICS_FILE)
+    fs, _ = fsspec.core.url_to_fs("memory://test_fit_csv_logger_remotefs")
+    assert fs.isfile(metrics_file)
+
+
 def test_flush_n_steps(tmpdir):
     logger = CSVLogger(tmpdir, flush_logs_every_n_steps=2)
     metrics = {"float": 0.3, "int": 1, "FloatTensor": torch.tensor(0.1), "IntTensor": torch.tensor(1)}