googleapis / googleapis/python-aiplatform

Make `Endpoint.predict` method async

Abierto
#1,998 2 comentarios 5 reacciones 0 asignados Ver en GitHub
api: vertex-ai type: feature request
Lenguaje dominante
Python
Estrellas
905
Forks
465
Merge medio
1 d 13 h
PR fusionados (30 d)
44

Descripción

### 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.

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.