googleapis / googleapis/python-aiplatform
Make `Endpoint.predict` method async
- Ngôn ngữ chính
- Python
- Star
- 905
- Fork
- 465
- Merge trung bình
- 1 ngày 13 giờ
- Pull request đã merge (30 ngày)
- 44
Mô tả
### 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.
Hướng dẫn đóng góp
Đánh giá
Issue này chưa được đánh giá.