Merge pull request #1068 from julep-ai/x/continue-as-new-bug

fix(agents-api): Fix blob-store interceptor implementation

Author: Ahmad-mtos

agents-api/agents_api/activities/sync_items_remote.py

@@ -1,24 +1,23 @@
-import asyncio
 from typing import Any
 
 from beartype import beartype
 from temporalio import activity
 
-from ..common.protocol.remote import RemoteObject
+from ..worker.codec import RemoteObject
 
 
 @beartype
 async def save_inputs_remote_fn(inputs: list[Any]) -> list[Any | RemoteObject]:
     from ..common.interceptors import offload_if_large
 
-    return await asyncio.gather(*[offload_if_large(input) for input in inputs])
+    return [offload_if_large(input) for input in inputs]
 
 
 @beartype
 async def load_inputs_remote_fn(inputs: list[Any | RemoteObject]) -> list[Any]:
     from ..common.interceptors import load_if_remote
 
-    return await asyncio.gather(*[load_if_remote(input) for input in inputs])
+    return [load_if_remote(input) for input in inputs]
 
 
 save_inputs_remote = activity.defn(name="save_inputs_remote")(save_inputs_remote_fn)

agents-api/agents_api/activities/task_steps/evaluate_step.py

@@ -24,4 +24,4 @@ async def evaluate_step(
 
     except BaseException as e:
         activity.logger.error(f"Error in evaluate_step: {e}")
-        return StepOutcome(error=str(e) or repr(e))
+        return StepOutcome(error=str(e) or repr(e), output=None)

agents-api/agents_api/app.py

@@ -2,15 +2,18 @@
 from contextlib import asynccontextmanager
 from typing import Protocol
 
-from aiobotocore.client import AioBaseClient
-from aiobotocore.session import get_session
-from asyncpg.pool import Pool
-from fastapi import APIRouter, FastAPI
-from prometheus_fastapi_instrumentator import Instrumentator
-from scalar_fastapi import get_scalar_api_reference
+from temporalio import workflow
 
-from .clients.pg import create_db_pool
-from .env import api_prefix, hostname, pool_max_size, protocol, public_port
+with workflow.unsafe.imports_passed_through():
+    from aiobotocore.client import AioBaseClient
+    from aiobotocore.session import get_session
+    from asyncpg.pool import Pool
+    from fastapi import APIRouter, FastAPI
+    from prometheus_fastapi_instrumentator import Instrumentator
+    from scalar_fastapi import get_scalar_api_reference
+
+    from .clients.pg import create_db_pool
+    from .env import api_prefix, hostname, pool_max_size, protocol, public_port
 
 
 class State(Protocol):
@@ -28,7 +31,7 @@ async def lifespan(container: FastAPI | ObjectWithState):
     # INIT POSTGRES #
     pg_dsn = os.environ.get("PG_DSN")
 
-    pool = await create_db_pool(pg_dsn, max_size=pool_max_size)
+    pool = await create_db_pool(pg_dsn, max_size=pool_max_size, min_size=min(pool_max_size, 10))
 
     if hasattr(container, "state") and not getattr(container.state, "postgres_pool", None):
         container.state.postgres_pool = pool

agents-api/agents_api/clients/async_s3.py

@@ -3,25 +3,22 @@
 
 with workflow.unsafe.imports_passed_through():
     import botocore
+    from aiobotocore.client import AioBaseClient
     from async_lru import alru_cache
     from xxhash import xxh3_64_hexdigest as xxhash_key
 
-    from ..env import (
-        blob_store_bucket,
-        blob_store_cutoff_kb,
-    )
+    from ..env import blob_store_bucket
 
 
 @alru_cache(maxsize=1)
-async def setup():
+async def setup() -> AioBaseClient:
     from ..app import app
 
-    if not app.state.s3_client:
+    client: AioBaseClient | None = getattr(app.state, "s3_client", None)
+    if client is None:
         msg = "S3 client not initialized"
         raise RuntimeError(msg)
 
-    client = app.state.s3_client
-
     try:
         await client.head_bucket(Bucket=blob_store_bucket)
     except botocore.exceptions.ClientError as e:
@@ -68,7 +65,6 @@ async def add_object(key: str, body: bytes, replace: bool = False) -> None:
     await client.put_object(Bucket=blob_store_bucket, Key=key, Body=body)
 
 
-@alru_cache(maxsize=256 * 1024 // max(1, blob_store_cutoff_kb))  # 256mb in cache
 @beartype
 async def get_object(key: str) -> bytes:
     client = await setup()

agents-api/agents_api/clients/sync_s3.py

@@ -0,0 +1,95 @@
+import os
+from functools import lru_cache
+
+from beartype import beartype
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    import botocore
+    from xxhash import xxh3_64_hexdigest as xxhash_key
+
+    from ..env import blob_store_bucket
+
+
+@lru_cache(maxsize=1)
+def setup():
+    # INIT S3 #
+    s3_access_key = os.environ.get("S3_ACCESS_KEY")
+    s3_secret_key = os.environ.get("S3_SECRET_KEY")
+    s3_endpoint = os.environ.get("S3_ENDPOINT")
+
+    session = botocore.session.Session()
+    client = session.create_client(
+        "s3",
+        endpoint_url=s3_endpoint,
+        aws_access_key_id=s3_access_key,
+        aws_secret_access_key=s3_secret_key,
+        config=botocore.config.Config(signature_version="s3v4", retries={"max_attempts": 3}),
+    )
+
+    try:
+        client.head_bucket(Bucket=blob_store_bucket)
+    except botocore.exceptions.ClientError as e:
+        if e.response["Error"]["Code"] == "404":
+            client.create_bucket(Bucket=blob_store_bucket)
+        else:
+            raise e
+
+    return client
+
+
+@lru_cache(maxsize=1024)
+def list_buckets() -> list[str]:
+    client = setup()
+
+    data = client.list_buckets()
+    return [bucket["Name"] for bucket in data["Buckets"]]
+
+
+@lru_cache(maxsize=10_000)
+def exists(key: str) -> bool:
+    client = setup()
+
+    try:
+        client.head_object(Bucket=blob_store_bucket, Key=key)
+        return True
+    except botocore.exceptions.ClientError as e:
+        if e.response["Error"]["Code"] == "404":
+            return False
+        raise e
+
+
+@beartype
+def add_object(key: str, body: bytes, replace: bool = False) -> None:
+    client = setup()
+
+    if replace:
+        client.put_object(Bucket=blob_store_bucket, Key=key, Body=body)
+        return
+
+    if exists(key):
+        return
+
+    client.put_object(Bucket=blob_store_bucket, Key=key, Body=body)
+
+
+@beartype
+def get_object(key: str) -> bytes:
+    client = setup()
+
+    response = client.get_object(Bucket=blob_store_bucket, Key=key)
+    return response["Body"].read()
+
+
+@beartype
+def delete_object(key: str) -> None:
+    client = setup()
+    client.delete_object(Bucket=blob_store_bucket, Key=key)
+
+
+@beartype
+def add_object_with_hash(body: bytes, replace: bool = False) -> str:
+    key = xxhash_key(body)
+    add_object(key, body, replace=replace)
+
+    return key

agents-api/agents_api/clients/temporal.py

@@ -114,7 +114,7 @@ async def run_task_execution_workflow(
     execution_id_key = SearchAttributeKey.for_keyword("CustomStringField")
 
     old_args = execution_input.arguments
-    execution_input.arguments = await offload_if_large(old_args)
+    execution_input.arguments = offload_if_large(old_args)
 
     current_input: dict = current_input or execution_input.arguments