andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Improve remote call

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-11 14:11:20 +02:00
parent 5800557168
commit 6d57622ce3
No known key found for this signature in database
GPG key ID: 0EA1BC97D0AB076E
2 changed files with 19 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

7
great_ai/remote/call_remote_great_ai.py
View file

@ -17,6 +17,7 @@ def call_remote_great_ai(
base_uri: str, base_uri: str,
data: Mapping[str, Any], data: Mapping[str, Any],
retry_count: int = 4, retry_count: int = 4,
timeout_in_seconds: Optional[int] = 300,
model_class: Optional[Type[T]] = None, model_class: Optional[Type[T]] = None,
) -> Trace[T]: ) -> Trace[T]:
try: try:
@ -28,7 +29,11 @@ def call_remote_great_ai(
pass pass
future = call_remote_great_ai_async( future = call_remote_great_ai_async(
base_uri=base_uri, data=data, retry_count=retry_count, model_class=model_class base_uri=base_uri,
data=data,
retry_count=retry_count,
timeout_in_seconds=timeout_in_seconds,
model_class=model_class,
) )
return asyncio.run(future) return asyncio.run(future)

6
great_ai/remote/call_remote_great_ai_async.py
View file

@ -13,7 +13,7 @@ async def call_remote_great_ai_async(
base_uri: str, base_uri: str,
data: Mapping[str, Any], data: Mapping[str, Any],
retry_count: int = 4, retry_count: int = 4,
timeout_in_seconds: int = 300, timeout_in_seconds: Optional[int] = 300,
model_class: Optional[Type[T]] = None, model_class: Optional[Type[T]] = None,
) -> Trace[T]: ) -> Trace[T]:
@ -24,6 +24,8 @@ async def call_remote_great_ai_async(
base_uri = f"{base_uri}/predict" base_uri = f"{base_uri}/predict"
transport = httpx.AsyncHTTPTransport(retries=retry_count) transport = httpx.AsyncHTTPTransport(retries=retry_count)
try:
async with httpx.AsyncClient( async with httpx.AsyncClient(
transport=transport, timeout=timeout_in_seconds transport=transport, timeout=timeout_in_seconds
) as client: ) as client:
@ -32,6 +34,8 @@ async def call_remote_great_ai_async(
response.raise_for_status() response.raise_for_status()
except Exception: except Exception:
raise RemoteCallError("Unexpected status code") raise RemoteCallError("Unexpected status code")
except Exception as e:
raise RemoteCallError from e
try: try:
trace = response.json() trace = response.json()