try Auth before training starts (#41)

* try Auth before training starts
* ConnectionError
* linter

Author: Borda

src/litmodels/integrations/lightning_checkpoint.py

@@ -1,5 +1,7 @@
 from typing import Any
 
+from lightning_sdk.lightning_cloud.login import Auth
+
 from litmodels import upload_model
 from litmodels.integrations.imports import _LIGHTNING_AVAILABLE, _PYTORCHLIGHTNING_AVAILABLE
 
@@ -28,6 +30,13 @@ def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.model_name = model_name
 
+        try:
+            # authenticate before anything else starts
+            auth = Auth()
+            auth.authenticate()
+        except Exception:
+            raise ConnectionError("Unable to authenticate with Lightning Cloud. Check your credentials.")
+
     def _save_checkpoint(self, trainer: Trainer, filepath: str) -> None:
         super()._save_checkpoint(trainer, filepath)
         # todo: uploading on background so training does nt stops

tests/integrations/test_lightning.py

@@ -13,7 +13,8 @@
 
 
 @mock.patch("litmodels.io.cloud.sdk_upload_model")
-def test_lightning_checkpoint_callback(mock_upload_model, tmp_path):
+@mock.patch("litmodels.integrations.lightning_checkpoint.Auth")
+def test_lightning_checkpoint_callback(mock_auth, mock_upload_model, tmp_path):
     mock_upload_model.return_value.name = "org-name/teamspace/model-name"
 
     trainer = Trainer(