fix(agents-api): Remove s3 sync client + minor fixes

Signed-off-by: Diwank Singh Tomer <[email protected]>
julep-ai · Nov 29, 2024 · 14e4c3c · 14e4c3c
1 parent f6bf839
commit 14e4c3c
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 115 deletions.
diff --git a/agents-api/agents_api/clients/s3.py b/agents-api/agents_api/clients/s3.py
diff --git a/agents-api/agents_api/common/storage_handler.py b/agents-api/agents_api/common/storage_handler.py
@@ -67,9 +67,11 @@ async def load_args(
             new_args = await asyncio.gather(
                 *[load_from_blob_store_if_remote(arg) for arg in args]
             )
-            new_kwargs = {
-                k: await load_from_blob_store_if_remote(v) for k, v in kwargs.items()
-            }
+            kwargs_keys, kwargs_values = zip(*kwargs.items())
+            new_kwargs = await asyncio.gather(
+                *[load_from_blob_store_if_remote(v) for v in kwargs_values]
+            )
+            new_kwargs = dict(zip(kwargs_keys, new_kwargs))
 
             if deep:
                 args = new_args
@@ -85,12 +87,12 @@ async def load_args(
                             )
                         )
                     elif isinstance(arg, dict):
-                        new_args.append(
-                            {
-                                key: await load_from_blob_store_if_remote(value)
-                                for key, value in arg.items()
-                            }
+                        keys, values = zip(*arg.items())
+                        values = await asyncio.gather(
+                            *[load_from_blob_store_if_remote(value) for value in values]
                         )
+                        new_args.append(dict(zip(keys, values)))
+
                     elif isinstance(arg, BaseRemoteModel):
                         new_args.append(await arg.unload_all())
 
@@ -110,7 +112,7 @@ async def load_args(
                                     field,
                                     await asyncio.gather(
                                         *[
-                                            await load_from_blob_store_if_remote(item)
+                                            load_from_blob_store_if_remote(item)
                                             for item in getattr(arg, field)
                                         ]
                                     ),
@@ -131,18 +133,19 @@ async def load_args(
 
                 for k, v in kwargs.items():
                     if isinstance(v, list):
-                        new_kwargs[k] = [
-                            await load_from_blob_store_if_remote(item) for item in v
-                        ]
+                        new_kwargs[k] = await asyncio.gather(
+                            *[load_from_blob_store_if_remote(item) for item in v]
+                        )
 
                     elif isinstance(v, dict):
-                        new_kwargs[k] = {
-                            key: await load_from_blob_store_if_remote(value)
-                            for key, value in v.items()
-                        }
+                        keys, values = zip(*v.items())
+                        values = await asyncio.gather(
+                            *[load_from_blob_store_if_remote(value) for value in values]
+                        )
+                        new_kwargs[k] = dict(zip(keys, values))
 
                     elif isinstance(v, BaseRemoteModel):
-                        new_kwargs[k] = v.unload_all()
+                        new_kwargs[k] = await v.unload_all()
 
                     elif isinstance(v, BaseModel):
                         for field in v.model_fields.keys():
@@ -160,7 +163,7 @@ async def load_args(
                                     field,
                                     await asyncio.gather(
                                         *[
-                                            await load_from_blob_store_if_remote(item)
+                                            load_from_blob_store_if_remote(item)
                                             for item in getattr(v, field)
                                         ]
                                     ),

diff --git a/agents-api/agents_api/env.py b/agents-api/agents_api/env.py
@@ -33,12 +33,14 @@
 
 # Blob Store
 # ----------
-use_blob_store_for_temporal: bool = (
-    env.bool("USE_BLOB_STORE_FOR_TEMPORAL", default=False) if not testing else False
-)
+use_blob_store_for_temporal: bool = True
+# use_blob_store_for_temporal: bool = (
+#     env.bool("USE_BLOB_STORE_FOR_TEMPORAL", default=False) if not testing else False
+# )
 
 blob_store_bucket: str = env.str("BLOB_STORE_BUCKET", default="agents-api")
-blob_store_cutoff_kb: int = env.int("BLOB_STORE_CUTOFF_KB", default=64)
+blob_store_cutoff_kb: int = 0
+# blob_store_cutoff_kb: int = env.int("BLOB_STORE_CUTOFF_KB", default=64)
 s3_endpoint: str = env.str("S3_ENDPOINT", default="http://seaweedfs:8333")
 s3_access_key: str | None = env.str("S3_ACCESS_KEY", default=None)
 s3_secret_key: str | None = env.str("S3_SECRET_KEY", default=None)

diff --git a/agents-api/gunicorn_conf.py b/agents-api/gunicorn_conf.py
@@ -1,7 +1,10 @@
 import multiprocessing
+import os
+
+debug = os.getenv("AGENTS_API_DEBUG", "false").lower() == "true"
 
 # Gunicorn config variables
-workers = multiprocessing.cpu_count() - 1
+workers = multiprocessing.cpu_count() - 1 if not debug else 1
 worker_class = "uvicorn.workers.UvicornWorker"
 bind = "0.0.0.0:8080"
 keepalive = 120