googleapis / googleapis/python-aiplatform

Make `Endpoint.predict` method async

未关闭
#1,998 2 条评论 5 个 reaction 已指派 0 人 在 GitHub 查看
api: vertex-ai type: feature request
主要语言
Python
星标
905
派生
465
平均合并
1 天 13 小时
30 天内合并 PR
44

描述

### Problem
I want to request predictions on my image classifier endpoint. Since there is a limit of 1.5 MB per request, if I want to get predictions for several images I have to do the following:

```python
import base64
import time
from io import BytesIO

from PIL import Image
from pympler.asizeof import asizeof

def sync_predict(endpoint, instances):
return endpoint.predict(instances=instances).predictions

def make_sync_predictions(image_list, endpoint, max_mbs=1.5):
predictions = []
batch, batch_mbs = [], 0.0
for i, image_path in enumerate(image_list):
image = Image.open(image_path).convert("RGB")
buffer = BytesIO()
image.save(buffer, format="JPEG")
enc_image = base64.b64encode(buffer.getvalue())
b64_image = str(enc_image.decode("utf-8"))
image_mbs = asizeof(b64_image) / (1000**2)

if batch_mbs + image_mbs >= max_mbs:
predictions += sync_predict(endpoint=endpoint, instances=batch)
batch, batch_mbs = [{"data": {"b64": b64_image}}], image_mbs
else:
batch.append({"data": {"b64": b64_image}})
batch_mbs += image_mbs

if i == len(image_list) - 1:
predictions += sync_predict(endpoint=endpoint, instances=batch)

image_names = [image.name for image in image_list]
return zip(image_names, predictions)

start = time.perf_counter()
result = make_sync_predictions(image_list=test_images, endpoint=endpoint)
print(f"Predictions: {[item for item in result]}")
print(f"Total elapsed time: {time.perf_counter() - start} s")
```

But obviously, this way I cannot benefit from having multiple replicas, for example an deployed model with `min_replica_count=2`. So I change to this:

```python
import asyncio
import base64
import time
from io import BytesIO

from PIL import Image
from pympler.asizeof import asizeof

async def async_predict(endpoint, batch, queue):
image_names, instances = batch["image_names"], batch["instances"]
response = await endpoint.predict(instances=instances) # this doesn't work because it isn't async
for image_name, prediction in zip(image_names, response.predictions):
await queue.put((image_name, prediction))

async def make_async_predictions(image_list, endpoint, max_mbs=1.5):
batches = []
batch_names, batch_data, batch_mbs = [], [], 0.0
for i, image_path in enumerate(image_list):
image = Image.open(image_path).convert("RGB")
buffer = BytesIO()
image.save(buffer, format="JPEG")
enc_image = base64.b64encode(buffer.getvalue())
b64_image = str(enc_image.decode("utf-8"))
image_mbs = asizeof(b64_image) / (1000**2)

if batch_mbs + image_mbs >= max_mbs:
batches.append({"image_names": batch_names, "instances": batch_data})
batch_names, batch_data, batch_mbs = [image_path.name], [{"data": {"b64": b64_image}}], image_mbs
else:
batch_names.append(image_path.name)
batch_data.append({"data": {"b64": b64_image}})
batch_mbs += image_mbs

if i == len(image_list) - 1:
batches.append({"image_names": batch_names, "instances": batch_data})

predictions_queue = asyncio.Queue()
predictors = [asyncio.create_task(async_predict(endpoint, batch, predictions_queue)) for batch in batches]
await asyncio.gather(*predictors)
result = []
while not predictions_queue.empty():
item = await predictions_queue.get()
result.append(item)
predictions_queue.task_done()

return result

start = time.perf_counter()
result = await make_async_predictions(image_list=test_images, endpoint=endpoint)
print(f"Predictions: {[item for item in result]}")
print(f"Total elapsed time: {time.perf_counter() - start} s")
```

### Workaround
I can solve this by changing the `endpoint.predict` line to:

```python
response = await asyncio.get_event_loop().run_in_executor(None, endpoint.predict, instances)
```

But I think there should be an `async_predict` method, or maybe there should be a parameter `sync: bool` that would make the call blocking or non-blocking depending on the parameter.

贡献指南

打开贡献指南

评估

这个 Issue 还没有评估数据。

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。