Repository: skypilot-org/skypilot
Pull request: Fix Lambda rate limit by backoff and retry when exceeding
<!-- Describe the changes in this PR -->

This PR fixes #2261 by:
* Add retry with backoff to all Lambda API requests when the rate limit exceeds;
* Retry launching when exceeding the rate limit.

TODO: debug and verify it works properly

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[Lambda] Handle rate limit error crashing `sky launch` 
User reported:
> For lambda, I'm getting api rate limit errors. 
> However, the problem is that I'm guessing there are some secondary api calls that causes sky launch to fail because of the rate limi, so the job never gets sent, but the node itself gets provisioned. 
> And because sky launch fails, it doesn't autodown, so it's just left running

Error:
```
...
File "/tmp/skypilot_ray_up_niigobvr.py", line 76, in <module>
sdk.create_or_update_cluster('/root/.sky/generated/training-228-53355c15.yml', **{'no_restart': True})
File "/usr/local/lib/python3.8/site-packages/ray/autoscaler/sdk/sdk.py", line 38, in create_or_update_cluster
return commands.create_or_update_cluster(
File "/usr/local/lib/python3.8/site-packages/ray/autoscaler/_private/commands.py", line 282, in create_or_update_cluster
get_or_create_head_node(
File "/usr/local/lib/python3.8/site-packages/ray/autoscaler/_private/commands.py", line 726, in get_or_create_head_node
provider.create_node(head_node_config, head_node_tags, 1)
File "/usr/local/lib/python3.8/site-packages/sky/skylet/providers/lambda_cloud/node_provider.py", line 287, in create_node
vms = self._list_instances_in_cluster()
File "/usr/local/lib/python3.8/site-packages/sky/skylet/providers/lambda_cloud/node_provider.py", line 127, in _list_instances_in_cluster
vms = self.lambda_client.list_instances()
File "/usr/local/lib/python3.8/site-packages/sky/skylet/providers/lambda_cloud/lambda_utils.py", line 160, in list_instances
raise_lambda_error(response)
File "/usr/local/lib/python3.8/site-packages/sky/skylet/providers/lambda_cloud/lambda_utils.py", line 78, in raise_lambda_error
raise LambdaCloudError('Your API requests are being rate limited.')
sky.skylet.providers.lambda_cloud.lambda_utils.LambdaCloudError: Your API requests are being rate limited.
```

Hints from issue discussion:
Can reproduce this with:
```
for i in $(seq 5); do sky launch --cpus 1+ --cloud lambda --down -y -c lambda-$i &; done
```

Commit message: merge to master

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index d3dd34db2b4..5b28cd43660 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -1787,6 +1787,12 @@ def need_ray_up(
                         'error.')
                     return True
 
+            if isinstance(to_provision_cloud, clouds.Lambda):
+                if 'Your API requests are being rate limited.' in stderr:
+                    logger.info(
+                        'Retrying due to Lambda API rate limit exceeded.')
+                    return True
+
             if 'rsync: command not found' in stderr:
                 logger.info('Skipping retry due to `rsync` not found in '
                             'the specified image.')
diff --git a/sky/skylet/providers/lambda_cloud/lambda_utils.py b/sky/skylet/providers/lambda_cloud/lambda_utils.py
index 3ccc80a8ff2..7abff31aa94 100644
--- a/sky/skylet/providers/lambda_cloud/lambda_utils.py
+++ b/sky/skylet/providers/lambda_cloud/lambda_utils.py
@@ -2,10 +2,17 @@
 import json
 import os
 import requests
+import time
 from typing import Any, Dict, List, Optional, Tuple
 
+from sky.utils import common_utils
+
 CREDENTIALS_PATH = '~/.lambda_cloud/lambda_keys'
 API_ENDPOINT = 'https://cloud.lambdalabs.com/api/v1'
+# TODO(tian): Determine best backoff factors.
+INITIAL_BACKOFF = 3
+MAX_BACKOFF_FACTOR = 10
+MAX_ATTEMPTS = 6
 
 
 class LambdaCloudError(Exception):
@@ -85,6 +92,34 @@ def raise_lambda_error(response: requests.Response) -> None:
     raise LambdaCloudError(f'{code}: {message}')
 
 
+def try_post_with_backoff(url: str, data: str, headers: Dict[str, str]):
+    backoff = common_utils.Backoff(initial_backoff=INITIAL_BACKOFF,
+                                   max_backoff_factor=MAX_BACKOFF_FACTOR)
+    for i in range(MAX_ATTEMPTS):
+        response = requests.post(url, data=data, headers=headers)
+        # If rate limited, wait and try again
+        if response.status_code == 429 and i != MAX_ATTEMPTS - 1:
+            time.sleep(backoff.current_backoff())
+            continue
+        if response.status_code == 200:
+            return response
+        raise_lambda_error(response)
+
+
+def try_get_with_backoff(url: str, headers: Dict[str, str]):
+    backoff = common_utils.Backoff(initial_backoff=INITIAL_BACKOFF,
+                                   max_backoff_factor=MAX_BACKOFF_FACTOR)
+    for i in range(MAX_ATTEMPTS):
+        response = requests.get(url, headers=headers)
+        # If rate limited, wait and try again
+        if response.status_code == 429 and i != MAX_ATTEMPTS - 1:
+            time.sleep(backoff.current_backoff())
+            continue
+        if response.status_code == 200:
+            return response
+        raise_lambda_error(response)
+
+
 class LambdaCloudClient:
     """Wrapper functions for Lambda Cloud API."""
 
@@ -133,10 +168,10 @@ def create_instances(self,
             'quantity': quantity,
             'name': name
         })
-        response = requests.post(f'{API_ENDPOINT}/instance-operations/launch',
-                                 data=data,
-                                 headers=self.headers)
-        raise_lambda_error(response)
+        response = try_post_with_backoff(
+            f'{API_ENDPOINT}/instance-operations/launch',
+            data=data,
+            headers=self.headers)
         return response.json().get('data', []).get('instance_ids', [])
 
     def remove_instances(self, *instance_ids: str) -> Dict[str, Any]:
@@ -146,25 +181,22 @@ def remove_instances(self, *instance_ids: str) -> Dict[str, Any]:
                 instance_ids[0]  # TODO(ewzeng) don't hardcode
             ]
         })
-        response = requests.post(
+        response = try_post_with_backoff(
             f'{API_ENDPOINT}/instance-operations/terminate',
             data=data,
             headers=self.headers)
-        raise_lambda_error(response)
         return response.json().get('data', []).get('terminated_instances', [])
 
     def list_instances(self) -> List[Dict[str, Any]]:
         """List existing instances."""
-        response = requests.get(f'{API_ENDPOINT}/instances',
-                                headers=self.headers)
-        raise_lambda_error(response)
+        response = try_get_with_backoff(f'{API_ENDPOINT}/instances',
+                                        headers=self.headers)
         return response.json().get('data', [])
 
     def list_ssh_keys(self) -> List[Dict[str, str]]:
         """List ssh keys."""
-        response = requests.get(f'{API_ENDPOINT}/ssh-keys',
-                                headers=self.headers)
-        raise_lambda_error(response)
+        response = try_get_with_backoff(f'{API_ENDPOINT}/ssh-keys',
+                                        headers=self.headers)
         return response.json().get('data', [])
 
     def get_unique_ssh_key_name(self, prefix: str,
@@ -200,14 +232,12 @@ def get_unique_ssh_key_name(self, prefix: str,
     def register_ssh_key(self, name: str, pub_key: str) -> None:
         """Register ssh key with Lambda."""
         data = json.dumps({'name': name, 'public_key': pub_key})
-        response = requests.post(f'{API_ENDPOINT}/ssh-keys',
-                                 data=data,
-                                 headers=self.headers)
-        raise_lambda_error(response)
+        try_post_with_backoff(f'{API_ENDPOINT}/ssh-keys',
+                              data=data,
+                              headers=self.headers)
 
     def list_catalog(self) -> Dict[str, Any]:
         """List offered instances and their availability."""
-        response = requests.get(f'{API_ENDPOINT}/instance-types',
-                                headers=self.headers)
-        raise_lambda_error(response)
+        response = try_get_with_backoff(f'{API_ENDPOINT}/instance-types',
+                                        headers=self.headers)
         return response.json().get('data', [])


Code review comments:
diff_hunk: @@ -2,10 +2,17 @@
 import json
 import os
 import requests
+import time
 from typing import Any, Dict, List, Optional, Tuple
 
+from sky.utils import common_utils
+
 CREDENTIALS_PATH = '~/.lambda_cloud/lambda_keys'
 API_ENDPOINT = 'https://cloud.lambdalabs.com/api/v1'
+# TODO(tian): Determine best backoff factors.
+INITIAL_BACKOFF = 3
path: sky/skylet/providers/lambda_cloud/lambda_utils.py
text: @user1:
```suggestion
INITIAL_BACKOFF_SECONDS = 3
```

diff_hunk: @@ -85,6 +92,34 @@ def raise_lambda_error(response: requests.Response) -> None:
     raise LambdaCloudError(f'{code}: {message}')
 
 
+def try_post_with_backoff(url: str, data: str, headers: Dict[str, str]):
+    backoff = common_utils.Backoff(initial_backoff=INITIAL_BACKOFF,
+                                   max_backoff_factor=MAX_BACKOFF_FACTOR)
+    for i in range(MAX_ATTEMPTS):
+        response = requests.post(url, data=data, headers=headers)
+        # If rate limited, wait and try again
+        if response.status_code == 429 and i != MAX_ATTEMPTS - 1:
+            time.sleep(backoff.current_backoff())
+            continue
+        if response.status_code == 200:
+            return response
+        raise_lambda_error(response)
+
+
+def try_get_with_backoff(url: str, headers: Dict[str, str]):
path: sky/skylet/providers/lambda_cloud/lambda_utils.py
text: @user1:
Can we refactor these two funcs to use the same code path?

diff_hunk: @@ -2,10 +2,17 @@
 import json
 import os
 import requests
+import time
 from typing import Any, Dict, List, Optional, Tuple
 
+from sky.utils import common_utils
+
 CREDENTIALS_PATH = '~/.lambda_cloud/lambda_keys'
 API_ENDPOINT = 'https://cloud.lambdalabs.com/api/v1'
+# TODO(tian): Determine best backoff factors.
+INITIAL_BACKOFF = 3
path: sky/skylet/providers/lambda_cloud/lambda_utils.py
text: @user1:
I tried the repro and for the first ~3 minutes, no launch can go through:
```
I 07-24 20:59:38 cloud_vm_ray_backend.py:1732] Launching on Lambda us-west-1
⠧ LaunchingI 07-24 20:59:38 cloud_vm_ray_backend.py:1732] Launching on Lambda us-west-1
⠇ LaunchingI 07-24 20:59:38 cloud_vm_ray_backend.py:1732] Launching on Lambda us-west-1
⠋ LaunchingI 07-24 20:59:38 cloud_vm_ray_backend.py:1732] Launching on Lambda us-west-1
⠹ LaunchingI 07-24 20:59:38 cloud_vm_ray_backend.py:1732] Launching on Lambda us-west-1
⠹ LaunchingI 07-24 21:00:18 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:00:18 cloud_vm_ray_backend.py:1826] Retrying launching in 6.8 seconds.
⠦ LaunchingI 07-24 21:00:32 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:00:32 cloud_vm_ray_backend.py:1826] Retrying launching in 5.1 seconds.
⠦ LaunchingI 07-24 21:00:39 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:00:39 cloud_vm_ray_backend.py:1826] Retrying launching in 6.8 seconds.
⠋ LaunchingI 07-24 21:01:14 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:01:14 cloud_vm_ray_backend.py:1826] Retrying launching in 4.0 seconds.
⠇ LaunchingI 07-24 21:01:24 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:01:24 cloud_vm_ray_backend.py:1826] Retrying launching in 7.1 seconds.
I 07-24 21:01:24 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:01:24 cloud_vm_ray_backend.py:1826] Retrying launching in 14.1 seconds.
⠧ LaunchingI 07-24 21:01:45 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:01:45 cloud_vm_ray_backend.py:1826] Retrying launching in 7.8 seconds.
⠧ LaunchingI 07-24 21:02:00 cloud_vm_ray_backend.py:1793] Retrying due to Lambda API rate limit exceeded.
I 07-24 21:02:00 cloud_vm_ray_backend.py:1826] Retrying launching in 17.2 seconds.
⠹ Launching
```

Should we try increasing initial backoff to say like 10 seconds to lessen the contention?

@author:
Tested initial = 10s and it performs very well (not a single retry happened when launching 5 clusters simultaneously). I think 10s is a good choice.