bump timeout

kscalelabs · Jan 16, 2025 · 4a6e040 · 4a6e040
1 parent 0cb0fe8
commit 4a6e040
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/kscale/requirements.txt b/kscale/requirements.txt
@@ -3,7 +3,6 @@
 # Configuration
 omegaconf
 email_validator
-colorlogging
 
 # HTTP requests
 aiohttp
@@ -17,6 +16,7 @@ yarl
 # CLI
 aiofiles
 click
+colorlogging
 tabulate
 
 # Async

diff --git a/kscale/web/clients/robot_class.py b/kscale/web/clients/robot_class.py
@@ -12,6 +12,9 @@
 
 logger = logging.getLogger(__name__)
 
+UPLOAD_TIMEOUT = 300.0
+DOWNLOAD_TIMEOUT = 60.0
+
 
 class RobotClassClient(BaseClient):
     async def get_robot_classes(self) -> list[RobotClass]:
@@ -77,7 +80,9 @@ async def upload_robot_class_urdf(self, class_name: str, urdf_file: str | Path)
             auth=True,
         )
         response = RobotUploadURDFResponse.model_validate(data)
-        async with httpx.AsyncClient() as client:
+        async with httpx.AsyncClient(
+            timeout=httpx.Timeout(UPLOAD_TIMEOUT),
+        ) as client:
             async with client.stream(
                 "PUT",
                 response.url,
@@ -97,7 +102,9 @@ async def download_robot_class_urdf(self, class_name: str, *, cache: bool = True
         cache_path.parent.mkdir(parents=True, exist_ok=True)
 
         logger.info("Downloading URDF file from %s", response.url)
-        async with httpx.AsyncClient() as client:
+        async with httpx.AsyncClient(
+            timeout=httpx.Timeout(DOWNLOAD_TIMEOUT),
+        ) as client:
             with open(cache_path, "wb") as file:
                 hash_value = hashlib.md5()
                 async with client.stream("GET", response.url) as r: