Add timeout to requests calls (#1)

jax-ml · Oct 1, 2024 · 4f8c127 · 4f8c127
1 parent cd6eeea
commit 4f8c127
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/jax/_src/clusters/cloud_tpu_cluster.py b/jax/_src/clusters/cloud_tpu_cluster.py
@@ -44,7 +44,7 @@ def get_metadata(key):
   while retry_count < 6:
     api_resp = requests.get(
         f'{gce_metadata_endpoint}/computeMetadata/v1/instance/attributes/{key}',
-        headers={'Metadata-Flavor': 'Google'})
+        headers={'Metadata-Flavor': 'Google'}, timeout=60)
     if api_resp.status_code == 200:
       break
     retry_count += 1

diff --git a/jax/experimental/jax2tf/examples/serving/model_server_request.py b/jax/experimental/jax2tf/examples/serving/model_server_request.py
@@ -92,7 +92,7 @@ def serving_call_mnist(images):
     # You can see the name of the input ("inputs") in the SavedModel dump.
     data = f'{{"inputs": {images_json}}}'
     predict_url = f"http://{_PREDICTION_SERVICE_ADDR.value}/v1/models/{_MODEL_SPEC_NAME.value}:predict"
-    response = requests.post(predict_url, data=data)
+    response = requests.post(predict_url, data=data, timeout=60)
     if response.status_code != 200:
       msg = (f"Received error response {response.status_code} from model "
              f"server: {response.text}")

diff --git a/jax/experimental/jax2tf/examples/tf_js/quickdraw/input_pipeline.py b/jax/experimental/jax2tf/examples/tf_js/quickdraw/input_pipeline.py
@@ -39,7 +39,7 @@ def download_dataset(dir_path, nb_classes):
       continue
     with open(cls_file_path, "wb") as save_file:
       try:
-        response = requests.get(url + cls_filename.replace('_', ' '))
+        response = requests.get(url + cls_filename.replace('_', ' '), timeout=60)
         save_file.write(response.content)
         print(f'Successfully fetched {cls_filename}')
       except: