chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_lightning_integration.py

@@ -0,0 +1,177 @@
+import os
+
+import pytest
+
+import ray
+from ray.train import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
+from ray.train.lightning import (
+    RayDDPStrategy,
+    RayFSDPStrategy,
+    RayLightningEnvironment,
+    RayTrainReportCallback,
+)
+from ray.train.lightning._lightning_utils import import_lightning
+from ray.train.tests.lightning_test_utils import DummyDataModule, LinearModule
+from ray.train.torch import TorchTrainer
+from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
+from ray.train.v2.api.report_config import CheckpointUploadMode
+from ray.train.v2.api.validation_config import ValidationConfig, ValidationTaskConfig
+
+pl = import_lightning()
+
+
+@pytest.fixture(autouse=True)
+def reduce_health_check_interval(monkeypatch):
+    monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "0.2")
+    yield
+
+
+@pytest.mark.parametrize("strategy_name", ["ddp", "fsdp"])
+@pytest.mark.parametrize("accelerator", ["cpu"])
+# @pytest.mark.parametrize("accelerator", ["cpu", "gpu"])  # TODO: Enable GPU test
+@pytest.mark.parametrize("datasource", ["dataloader", "datamodule"])
+def test_trainer_with_native_dataloader(
+    ray_start_4_cpus, strategy_name, accelerator, datasource
+):
+    """Test basic ddp and fsdp training with dataloader and datamodule."""
+
+    if accelerator == "cpu" and strategy_name == "fsdp":
+        return
+
+    num_workers = 2
+    num_epochs = 1
+    batch_size = 8
+    dataset_size = 256
+
+    strategy_map = {"ddp": RayDDPStrategy(), "fsdp": RayFSDPStrategy()}
+
+    def train_loop():
+        model = LinearModule(input_dim=32, output_dim=4, strategy=strategy_name)
+
+        strategy = strategy_map[strategy_name]
+
+        trainer = pl.Trainer(
+            max_epochs=num_epochs,
+            devices="auto",
+            accelerator=accelerator,
+            strategy=strategy,
+            plugins=[RayLightningEnvironment()],
+            callbacks=[RayTrainReportCallback()],
+        )
+
+        datamodule = DummyDataModule(batch_size, dataset_size)
+
+        if datasource == "dataloader":
+            trainer.fit(
+                model,
+                train_dataloaders=datamodule.train_dataloader(),
+                val_dataloaders=datamodule.val_dataloader(),
+            )
+        if datasource == "datamodule":
+            trainer.fit(model, datamodule=datamodule)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=(accelerator == "gpu")),
+    )
+
+    results = trainer.fit()
+    assert results.metrics["epoch"] == num_epochs - 1
+    assert (
+        results.metrics["step"] == num_epochs * dataset_size / num_workers / batch_size
+    )
+    assert "loss" in results.metrics
+    assert "val_loss" in results.metrics
+
+
+def test_async_checkpointing_and_validation(ray_start_4_cpus, tmp_path):
+    """Test lightning training with async checkpointing and validation."""
+
+    num_workers = 2
+    num_epochs = 2
+    batch_size = 8
+    dataset_size = 256
+
+    @ray.remote
+    class TmpdirPrefixActor:
+        def __init__(self):
+            self.tmpdir_prefixes = []
+
+        def set_tmpdir_prefix(self, tmpdir_prefix):
+            self.tmpdir_prefixes.append(tmpdir_prefix)
+
+        def get_tmpdir_prefixes(self):
+            return self.tmpdir_prefixes
+
+    tmpdir_prefix_actor = TmpdirPrefixActor.remote()
+
+    def validation_fn(checkpoint):
+        assert checkpoint.path is not None
+        checkpoint_file = checkpoint.path + "/checkpoint.ckpt"
+        assert os.path.exists(
+            checkpoint_file
+        ), f"Checkpoint file not found: {checkpoint_file}"
+        return {"val_score": 1}
+
+    def train_loop():
+        checkpoint = ray.train.get_checkpoint()
+        model = LinearModule(input_dim=32, output_dim=4, strategy="ddp", fail_epoch=1)
+        callback = RayTrainReportCallback(
+            checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
+            validation=ValidationTaskConfig(fn_kwargs={}),
+        )
+        # Only track tmpdirs from the post-resume attempt.
+        # TODO: fix bug where async checkpoint upload does not clean up tmpdir if worker fails.
+        if checkpoint is not None:
+            ray.get(
+                tmpdir_prefix_actor.set_tmpdir_prefix.remote(callback.tmpdir_prefix)
+            )
+        trainer = pl.Trainer(
+            max_epochs=num_epochs,
+            devices="auto",
+            accelerator="cpu",
+            strategy=RayDDPStrategy(),
+            plugins=[RayLightningEnvironment()],
+            callbacks=[callback],
+        )
+
+        datamodule = DummyDataModule(batch_size, dataset_size)
+        if checkpoint is not None:
+            with checkpoint.as_directory() as ckpt_dir:
+                ckpt_path = os.path.join(
+                    ckpt_dir, RayTrainReportCallback.CHECKPOINT_NAME
+                )
+                trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path)
+        else:
+            trainer.fit(model, datamodule=datamodule)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        scaling_config=ScalingConfig(num_workers=num_workers),
+        validation_config=ValidationConfig(fn=validation_fn),
+        run_config=RunConfig(
+            storage_path=str(tmp_path),
+            checkpoint_config=CheckpointConfig(
+                num_to_keep=1, checkpoint_score_attribute="val_score"
+            ),
+            failure_config=FailureConfig(max_failures=1),
+        ),
+    )
+
+    results = trainer.fit()
+    assert results.error is None
+    assert "loss" in results.metrics
+    assert results.best_checkpoints is not None
+    assert len(results.best_checkpoints) == 1
+    assert results.best_checkpoints[0][1]["val_score"] == 1
+    recorded_prefixes = ray.get(tmpdir_prefix_actor.get_tmpdir_prefixes.remote())
+    assert len(recorded_prefixes) == num_workers
+    # Seems pyarrow.fs.FileSystem's delete_dir can leave an empty dir behind.
+    for path in recorded_prefixes:
+        assert not os.path.exists(path) or not any(os.scandir(path))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))