ENG-1524: Change retry mechansim default in LightningClient

src/lightning_app/CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+## [unreleased] - 2022-MM-DD
+
+### Fixed
+
+- Fixed an issue with the `lightning` CLI taking a long time to error out when the cloud is not reachable ([#15412](https://github.com/Lightning-AI/lightning/pull/15412))
+
 
 ## [1.8.0] - 2022-11-01
 

src/lightning_app/runners/backends/cloud.py

@@ -11,7 +11,7 @@
 class CloudBackend(Backend):
     def __init__(self, entrypoint_file, queue_id: Optional[str] = None, status_update_interval: int = None):
         super().__init__(entrypoint_file, queues=QueuingSystem.MULTIPROCESS, queue_id=queue_id)
-        self.client = LightningClient()
+        self.client = LightningClient(retry=True)
 
     def create_work(self, app: "lightning_app.LightningApp", work: "lightning_app.LightningWork") -> None:
         raise NotImplementedError

src/lightning_app/utilities/cloud.py

@@ -18,7 +18,7 @@ def _get_project(client: LightningClient, project_id: str = LIGHTNING_CLOUD_PROJ
                 break
         else:
             raise ValueError(
-                "Environment variable LIGHTNING_CLOUD_PROJECT_ID is set " "but could not find an associated project."
+                "Environment variable LIGHTNING_CLOUD_PROJECT_ID is set but could not find an associated project."
             )
         return membership
 

src/lightning_app/utilities/network.py

@@ -104,28 +104,26 @@ def wrapped(*args: Any, **kwargs: Any) -> Any:
     return wrapped
 
 
-class _MethodsRetryWrapperMeta(type):
-    """This wrapper metaclass iterates through all methods of the type and all bases of it to wrap them into the
-    :func:`_retry_wrapper`. It applies to all bound callables except the ``__init__`` method.
-    """
-
-    def __new__(mcs, name, bases, dct):
-        new_class = super().__new__(mcs, name, bases, dct)
-        for base in new_class.__mro__[1:-1]:
-            for key, value in base.__dict__.items():
-                if callable(value) and value.__name__ != "__init__":
-                    setattr(new_class, key, _retry_wrapper(value))
-        return new_class
-
-
-class LightningClient(GridRestClient, metaclass=_MethodsRetryWrapperMeta):
+class LightningClient(GridRestClient):
     """The LightningClient is a wrapper around the GridRestClient.
 
     It wraps all methods to monitor connection exceptions and employs a retry strategy.
+
+    Args:
+        retry: Whether API calls should follow a retry mechanism with exponential backoff.
     """
 
-    def __init__(self) -> None:
+    def __new__(cls, *args: Any, **kwargs: Any) -> "LightningClient":
+        if kwargs.get("retry", False):
+            for base_class in GridRestClient.__mro__:
+                for name, attribute in base_class.__dict__.items():
+                    if callable(attribute) and attribute.__name__ != "__init__":
+                        setattr(cls, name, _retry_wrapper(attribute))
+        return super().__new__(cls)
+
+    def __init__(self, retry: bool = False) -> None:
         super().__init__(api_client=create_swagger_client())
+        self._retry = retry
 
 
 class CustomRetryAdapter(HTTPAdapter):

tests/tests_app/cli/test_cloud_cli.py

@@ -37,9 +37,6 @@ class FakeResponse:
 
 
 class FakeLightningClient:
-    def __init__(self, response, api_client=None):
-        self._response = response
-
     def lightningapp_instance_service_list_lightningapp_instances(self, *args, **kwargs):
         return V1ListLightningappInstancesResponse(lightningapps=[])
 
@@ -102,7 +99,7 @@ class ExceptionResponse:
 
 class FakeLightningClientCreate(FakeLightningClient):
     def __init__(self, *args, create_response, **kwargs):
-        super().__init__(*args, **kwargs)
+        super().__init__()
         self.create_response = create_response
 
     def lightningapp_v2_service_list_lightningapps_v2(self, *args, **kwargs):
@@ -129,7 +126,7 @@ def test_start_app(create_response, monkeypatch):
     monkeypatch.setattr(
         cloud_backend,
         "LightningClient",
-        partial(FakeLightningClientCreate, response=FakeResponse(), create_response=create_response),
+        partial(FakeLightningClientCreate, create_response=create_response),
     )
     monkeypatch.setattr(cloud, "LocalSourceCodeDir", MagicMock())
     monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", MagicMock())
@@ -171,8 +168,8 @@ def run():
 
 
 class FakeLightningClientException(FakeLightningClient):
-    def __init__(self, *args, message, api_client=None, **kwargs):
-        super().__init__(*args, api_client=api_client, **kwargs)
+    def __init__(self, *args, message, **kwargs):
+        super().__init__()
         self.message = message
 
     def lightningapp_v2_service_list_lightningapps_v2(self, *args, **kwargs):
@@ -215,7 +212,7 @@ def test_start_app_exception(message, monkeypatch, caplog):
 
     runner = CliRunner()
 
-    fake_grid_rest_client = partial(FakeLightningClientException, response=FakeResponse(), message=message)
+    fake_grid_rest_client = partial(FakeLightningClientException, message=message)
     with caplog.at_level(logging.ERROR):
         with mock.patch("lightning_app.runners.backends.cloud.LightningClient", fake_grid_rest_client):
             result = runner.invoke(run_app, [_FILE_PATH, "--cloud", "--open-ui=False"], catch_exceptions=False)

tests/tests_app/utilities/test_network.py

@@ -1,5 +1,17 @@
-from lightning_app.utilities.network import find_free_network_port
+from unittest.mock import patch
+
+from lightning_app.utilities.network import find_free_network_port, LightningClient
 
 
 def test_port():
     assert find_free_network_port()
+
+
+def test_lightning_client_retry_enabled():
+    with patch("lightning_app.utilities.network._retry_wrapper") as wrapper:
+        LightningClient()  # default: retry=False
+        wrapper.assert_not_called()
+
+    with patch("lightning_app.utilities.network._retry_wrapper") as wrapper:
+        LightningClient(retry=True)
+        wrapper.assert_called()