test: integration Trainer's ddp_spawn with Prod (#62)

* test: integration Trainer's `ddp_spawn` with Prod
* Apply suggestions from code review

---------

Co-authored-by: Ethan Harris <[email protected]>

Author: Borda

_requirements/extra.txt

@@ -1 +1,2 @@
 lightning >= 2.0.0
+numpy <2.0.0 ; platform_system == "Darwin" # compatibility fix for Torch

tests/integrations/test_cloud.py

@@ -1,6 +1,7 @@
 import os
 from contextlib import redirect_stdout
 from io import StringIO
+from typing import Optional
 
 import pytest
 from lightning_sdk import Teamspace
@@ -15,13 +16,13 @@
 
 
 def _prepare_variables(test_name: str) -> tuple[Teamspace, str, str]:
-    model_name = f"litmodels_test_integrations_{test_name}+{os.urandom(8).hex()}"
+    model_name = f"ci-test_integrations_{test_name}+{os.urandom(8).hex()}"
     teamspace = _resolve_teamspace(org=LIT_ORG, teamspace=LIT_TEAMSPACE, user=None)
     org_team = f"{teamspace.owner.name}/{teamspace.name}"
     return teamspace, org_team, model_name
 
 
-def _cleanup_model(teamspace: Teamspace, model_name: str) -> None:
+def _cleanup_model(teamspace: Teamspace, model_name: str, expected_num_versions: Optional[int] = None) -> None:
     """Cleanup model from the teamspace."""
     client = GridRestClient()
     # cleaning created models as each test run shall have unique model name
@@ -30,7 +31,10 @@ def _cleanup_model(teamspace: Teamspace, model_name: str) -> None:
         project_name=teamspace.name,
         model_name=model_name,
     )
-    client.models_store_delete_model(project_id=teamspace.id, model_id=model.id)
+    if expected_num_versions is not None:
+        versions = client.models_store_list_model_versions(project_id=model.project_id, model_id=model.id)
+        assert expected_num_versions == len(versions.versions)
+    client.models_store_delete_model(project_id=model.project_id, model_id=model.id)
 
 
 @pytest.mark.cloud()
@@ -62,7 +66,7 @@ def test_upload_download_model(tmp_path):
         assert os.path.isfile(os.path.join(tmp_path, file))
 
     # CLEANING
-    _cleanup_model(teamspace, model_name)
+    _cleanup_model(teamspace, model_name, expected_num_versions=1)
 
 
 @pytest.mark.parametrize(
@@ -93,7 +97,7 @@ def test_lightning_default_checkpointing(importing, tmp_path):
     trainer.fit(BoringModel())
 
     # CLEANING
-    _cleanup_model(teamspace, model_name)
+    _cleanup_model(teamspace, model_name, expected_num_versions=2)
 
 
 @pytest.mark.parametrize("trainer_method", ["fit", "validate", "test", "predict"])
@@ -109,7 +113,7 @@ def test_lightning_default_checkpointing(importing, tmp_path):
 )
 @pytest.mark.cloud()
 # todo: mock env variables as it would run in studio
-def test_lightning_resume(trainer_method, registry, importing, tmp_path):
+def test_lightning_plain_resume(trainer_method, registry, importing, tmp_path):
     if importing == "lightning":
         from lightning import Trainer
         from lightning.pytorch.demos.boring_classes import BoringModel
@@ -124,6 +128,7 @@ def test_lightning_resume(trainer_method, registry, importing, tmp_path):
     # model name with random hash
     teamspace, org_team, model_name = _prepare_variables(f"resume_{trainer_method}")
     upload_model(model=checkpoint_path, name=f"{org_team}/{model_name}")
+    expected_num_versions = 1
 
     trainer_kwargs = {"model_registry": f"{org_team}/{model_name}"} if "<model>" not in registry else {}
     trainer = Trainer(
@@ -138,6 +143,8 @@ def test_lightning_resume(trainer_method, registry, importing, tmp_path):
     registry = registry.replace("<model>", f"{org_team}/{model_name}")
     if trainer_method == "fit":
         trainer.fit(BoringModel(), ckpt_path=registry)
+        if trainer_kwargs:
+            expected_num_versions += 1
     elif trainer_method == "validate":
         trainer.validate(BoringModel(), ckpt_path=registry)
     elif trainer_method == "test":
@@ -148,4 +155,41 @@ def test_lightning_resume(trainer_method, registry, importing, tmp_path):
         raise ValueError(f"Unknown trainer method: {trainer_method}")
 
     # CLEANING
-    _cleanup_model(teamspace, model_name)
+    _cleanup_model(teamspace, model_name, expected_num_versions=expected_num_versions)
+
+
+@pytest.mark.parametrize(
+    "importing",
+    [
+        pytest.param("lightning", marks=_SKIP_IF_LIGHTNING_BELLOW_2_5_1),
+        pytest.param("pytorch_lightning", marks=_SKIP_IF_PYTORCHLIGHTNING_BELLOW_2_5_1),
+    ],
+)
+@pytest.mark.cloud()
+def test_lightning_checkpoint_ddp(importing, tmp_path):
+    if importing == "lightning":
+        from lightning import Trainer
+        from lightning.pytorch.demos.boring_classes import BoringModel
+    elif importing == "pytorch_lightning":
+        from pytorch_lightning import Trainer
+        from pytorch_lightning.demos.boring_classes import BoringModel
+
+    # model name with random hash
+    teamspace, org_team, model_name = _prepare_variables("checkpoint_resume")
+    trainer_args = {
+        "default_root_dir": tmp_path,
+        "accelerator": "cpu",
+        "strategy": "ddp_spawn",
+        "devices": 4,
+        "model_registry": f"{org_team}/{model_name}",
+    }
+
+    trainer = Trainer(max_epochs=2, **trainer_args)
+    trainer.fit(BoringModel())
+
+    # FIXME: seems like barrier is not respected in the test, but in real life it correctly waits for all GPUs
+    # trainer = Trainer(max_epochs=5, **trainer_args)
+    # trainer.fit(BoringModel(), ckpt_path="registry")
+
+    # CLEANING
+    _cleanup_model(teamspace, model_name, expected_num_versions=2)