Skip to content

Commit

Permalink
fix: disable setup timeout
Browse files Browse the repository at this point in the history
fixes #4. timeout was used to avoid a strange behavior where cleanup would hang, but waiting for processing tasks and then workers (sequentially) seems to resolve this.
  • Loading branch information
ctjlewis committed Feb 12, 2024
1 parent a837869 commit 161d462
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 5 deletions.
8 changes: 3 additions & 5 deletions oaib/Batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -344,11 +344,9 @@ async def run(self, callback=None):
# existing requests first.
if not self.__stopped.is_set():
self.log("FINISHING PROCESSING | 5 second timeout")
await wait(
[*self.__processing, *self.__callbacks, *self.__workers],
return_when=ALL_COMPLETED,
timeout=5
)
await gather(*self.__processing)
await gather(*self.__workers)
await gather(*self.__callbacks)
await self.stop()

self.log("RETURNING OUTPUT")
Expand Down
24 changes: 24 additions & 0 deletions tests/test_large.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
import oaib


async def test_long():
batch = oaib.Auto()

n = 20
m = 20
for i in range(n):
await batch.add(
"chat.completions.create",
model="gpt-4",
max_tokens=4000,
messages=[{"role": "user", "content": "say hello and goodbye " * m}]
)

chats = await batch.run()
assert len(
chats) == n, f"Chat batch should return {n} results, got {len(chats)}"

chat = chats.iloc[0].get("result")
assert chat['choices'], "Should get valid chat completions"

chats

0 comments on commit 161d462

Please sign in to comment.