increase timeout, add more suffixes

kscalelabs · Sep 4, 2024 · d2877cb · d2877cb
1 parent a0b0fad
commit d2877cb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/kscale/store/client.py b/kscale/store/client.py
@@ -26,6 +26,7 @@ def __init__(self, base_url: str = get_api_root()) -> None:
         self.client = httpx.AsyncClient(
             base_url=self.base_url,
             headers={"Authorization": f"Bearer {get_api_key()}"},
+            timeout=httpx.Timeout(30.0),
         )
 
     async def _request(

diff --git a/kscale/store/urdf.py b/kscale/store/urdf.py
@@ -21,6 +21,17 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+ALLOWED_SUFFIXES = {
+    ".urdf",
+    ".mjcf",
+    ".stl",
+    ".obj",
+    ".dae",
+    ".png",
+    ".jpg",
+    ".jpeg",
+}
+
 
 def get_cache_dir() -> Path:
     return Path(Settings.load().store.cache_dir).expanduser().resolve()
@@ -75,7 +86,7 @@ def create_tarball(folder_path: Path, output_filename: str, cache_dir: Path) ->
     tarball_path = cache_dir / output_filename
     with tarfile.open(tarball_path, "w:gz") as tar:
         for file_path in folder_path.rglob("*"):
-            if file_path.is_file() and file_path.suffix.lower() in (".urdf", ".mjcf", ".stl", ".obj", ".dae"):
+            if file_path.is_file() and file_path.suffix.lower() in ALLOWED_SUFFIXES:
                 tar.add(file_path, arcname=file_path.relative_to(folder_path))
                 logger.info("Added %s to tarball", file_path)
             else: