ENG-1524: Change retry mechansim default in LightningClient (#15412)

* Change retry mechansim default in LightningClient
* add changelog
* fix weird patching mechanism
* remove unused import
* hack the system
* try again, maybe something has changed in the package

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Jirka <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>
(cherry picked from commit 7decc50)

Author: 3 people

src/lightning_app/CHANGELOG.md

@@ -4,16 +4,17 @@ 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] - 202Y-MM-DD
 
 
 ### Added
 
 - Added the `start` method to the work ([#15523](https://github.com/Lightning-AI/lightning/pull/15523))
 
+
 - Added a `MultiNode` Component to run with distributed computation with any frameworks ([#15524](https://github.com/Lightning-AI/lightning/pull/15524))
 
+
 - Expose `RunWorkExecutor` to the work and provides default ones for the `MultiNode` Component ([#15561](https://github.com/Lightning-AI/lightning/pull/15561))
 
 - Added a `start_with_flow` flag to the `LightningWork` which can be disabled to prevent the work from starting at the same time as the flow ([#15591](https://github.com/Lightning-AI/lightning/pull/15591))
@@ -34,7 +35,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 -
 
--
 
 -
 
@@ -43,7 +43,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 -
 
--
 
 -
 
@@ -53,11 +52,15 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed writing app name and id in connect.txt file for the command CLI ([#15443](https://github.com/Lightning-AI/lightning/pull/15443))
 
+
 - Fixed missing root flow among the flows of the app ([#15531](https://github.com/Lightning-AI/lightning/pull/15531))
 
+
 - Fixed bug with Multi Node Component and add some examples ([#15557](https://github.com/Lightning-AI/lightning/pull/15557))
 
 
+- 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())
@@ -182,8 +179,8 @@ def getheaders(self):
 
 
 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):
@@ -216,7 +213,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()

tests/tests_app/utilities/test_secrets.py

@@ -5,6 +5,7 @@
 import pytest
 from lightning_cloud.openapi import V1ListMembershipsResponse, V1ListSecretsResponse, V1Membership, V1Secret
 
+import lightning_app
 from lightning_app.utilities.secrets import _names_to_ids
 
 
@@ -29,18 +30,21 @@
     ],
 )
 @mock.patch("lightning_cloud.login.Auth.authenticate", MagicMock())
-@mock.patch("lightning_app.utilities.network.LightningClient.secret_service_list_secrets")
-@mock.patch("lightning_app.utilities.network.LightningClient.projects_service_list_memberships")
 def test_names_to_ids(
-    list_memberships: MagicMock,
-    list_secrets: MagicMock,
     secret_names: List[str],
     secrets: List[V1Secret],
     expected: Dict[str, str],
     expected_exception: bool,
+    monkeypatch,
 ):
-    list_memberships.return_value = V1ListMembershipsResponse(memberships=[V1Membership(project_id="default-project")])
-    list_secrets.return_value = V1ListSecretsResponse(secrets=secrets)
+    class FakeLightningClient:
+        def projects_service_list_memberships(self, *_, **__):
+            return V1ListMembershipsResponse(memberships=[V1Membership(project_id="default-project")])
+
+        def secret_service_list_secrets(self, *_, **__):
+            return V1ListSecretsResponse(secrets=secrets)
+
+    monkeypatch.setattr(lightning_app.utilities.secrets, "LightningClient", FakeLightningClient)
 
     if expected_exception:
         with pytest.raises(ValueError):