Skip to content

Commit

Permalink
fix the order of yield and async infer
Browse files Browse the repository at this point in the history
  • Loading branch information
picciama committed Dec 21, 2023
1 parent dd676d2 commit 3300f37
Showing 1 changed file with 7 additions and 5 deletions.
12 changes: 7 additions & 5 deletions oktoberfest/predict/koina.py
Original file line number Diff line number Diff line change
Expand Up @@ -408,7 +408,13 @@ def __async_predict_batch(
batch_outputs = self.__get_batch_outputs(self.model_outputs.keys())
batch_inputs = self.__get_batch_inputs(data)

for _ in range(retries):
for i in range(retries):
if i > 0: # need to yield first, before doing sth, but only after first time
yield
if isinstance(infer_results.get(request_id), InferResult):
break

Check warning on line 415 in oktoberfest/predict/koina.py

View check run for this annotation

Codecov / codecov/patch

oktoberfest/predict/koina.py#L415

Added line #L415 was not covered by tests
del infer_results[request_id] # avoid race condition in case inference is slower than tqdm loop

self.client.async_infer(
model_name=self.model_name,
request_id=str(request_id),
Expand All @@ -417,9 +423,6 @@ def __async_predict_batch(
outputs=batch_outputs,
client_timeout=timeout,
)
yield
if isinstance(infer_results.get(request_id), InferResult):
break

def predict(
self,
Expand Down Expand Up @@ -503,7 +506,6 @@ def __predict_async(
pbar.n += 1
else: # unexpected result / exception -> try again
try:
del infer_results[j] # avoid race condition in case inference is slower than loop
next(tasks[j])
new_unfinished_tasks.append(j)
except StopIteration:
Expand Down

0 comments on commit 3300f37

Please sign in to comment.