googleapis / googleapis/python-aiplatform

BatchPredictionJob.create does not support providing a value for instance_config

Open
#3,031 3 comments 14 reactions 0 assignees View on GitHub
api: vertex-ai
Dominant language
Python
Stars
905
Forks
465
Avg merge
1d 13h
Merged PRs (30d)
44

Description

**Is your feature request related to a problem? Please describe.**
We have a use case where we'd like to launch a Vertex AI Batch Prediction job using a custom model. In order for the Docker container backing our model to work, the inputs must be sent in as objects (i.e. `{"a": 1, "b": 2}` instead of `[1, 2]`). Currently, the only way to do this is to use the REST API and specify `"instanceConfig": {"instanceType": "object"}` as per [these docs](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions#filter_and_transform_input_data).

**Describe the solution you'd like**
We'd like to be able to set `instanceType` to `object` via the `BatchPredictionJob.create` method.

**Describe alternatives you've considered**
Using the REST API is, to my knowledge, the only alternative. We can use it, but it's not convenient - having this functionality available in the Python SDK would allow us to schedule Batch Prediction jobs in an easier manner. Moreover, it would be possible to update the official Airflow operator for this and we'd be able to leverage that instead of writing custom code.

Another alternative would be to export our data to GCS in JSONL format and use that as the input for our Batch Prediction job. However, that is also an additional step in our pipeline that we could avoid.

**Additional context**
I guess the easiest solution would be to extend the [`BatchPredictionJob.create` method](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/jobs.py#L384) to include a new optional field `instance_type: Optional[Literal["object", "array"]]` and attach the provided value (if any) to the `gapic_batch_prediction_job` object. Another approach would be to enable passing in the whole `instance_config` object. I'd be interested to hear which design sounds the most reasonable to you.

I'd be happy to submit a PR once a design is agreed upon.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.