googleapis / googleapis/python-aiplatform

400 Error: "bigquery" output format does not support key_field in aiplatform_v1.BatchPredictionJob.InstanceConfig

Open
#4,514 2 comments 0 reactions 1 assignee Claimed by @weichungw View on GitHub
api: vertex-ai
Dominant language
Python
Stars
905
Forks
465
Avg merge
1d 13h
Merged PRs (30d)
44

Description

## Summary
I encountered the following error when trying to specify the key_field in `aiplatform_v1.BatchPredictionJob.InstanceConfig` with a BigQuery input:
```
google.api_core.exceptions.InvalidArgument: 400 "bigquery" output format does not support key_field.
```

## Environment details

- OS type and version: `maxOS Sonoma 14.5`
- Python version: `3.11.10`
- `google-cloud-aiplatform` version: `1.69.0`

## Code example

```python
from google.cloud import aiplatform, aiplatform_v1

LOCATION = "asia-northeast1"
MY_PROJECT = "my-project"

def batch_predict_with_bq(
model: aiplatform.Model,
job_display_name: str,
bq_source_uri: str,
bq_output_uri: str,
machine_type: str,
) -> aiplatform_v1.BatchPredictionJob:
# https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.types.BatchPredictionJob.InputConfig
input_config = aiplatform_v1.BatchPredictionJob.InputConfig(
instances_format="bigquery",
bigquery_source=aiplatform_v1.BigQuerySource(
input_uri=bq_source_uri,
),
)

# https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.types.BatchPredictionJob.InstanceConfig
instance_config = aiplatform_v1.BatchPredictionJob.InstanceConfig(
excluded_fields=["user_id"],
key_field="key",
)

# https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.types.BatchPredictionJob.OutputConfig
output_config = aiplatform_v1.BatchPredictionJob.OutputConfig(
predictions_format="bigquery",
bigquery_destination=aiplatform_v1.BigQueryDestination(
output_uri=bq_output_uri,
),
)

# https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.types.BatchDedicatedResources
batch_dedicated_resources = aiplatform_v1.BatchDedicatedResources(
machine_spec=aiplatform_v1.MachineSpec(machine_type=machine_type),
starting_replica_count=1,
max_replica_count=1,
)

# https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.types.BatchPredictionJob

job = aiplatform_v1.BatchPredictionJob(
name="test",
display_name=job_display_name,
model=model.resource_name,
input_config=input_config,
output_config=output_config,
instance_config=instance_config,
dedicated_resources=batch_dedicated_resources,
)

# https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.services.job_service.JobServiceClient#google_cloud_aiplatform_v1_services_job_service_JobServiceClient_create_batch_prediction_job
client = aiplatform_v1.JobServiceClient(
client_options={"api_endpoint": f"{LOCATION}-aiplatform.googleapis.com"}
)

request = aiplatform_v1.CreateBatchPredictionJobRequest(
parent=f"projects/{MY_PROJECT}/locations/{LOCATION}",
batch_prediction_job=job,
)

response = client.create_batch_prediction_job(request=request)

return response

```

## Stack trace
```
The above exception was the direct cause of the following exception:

Traceback (most recent call last):

...
_ = batch_predict_with_bq(
^^^^^^^^^^^^^^^^^^^^^^
File "/Users/user/workspace/prog/python/vertexai/batch/main.py", line 305, in batch_predict_with_bq
response = client.create_batch_prediction_job(request=request)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/user/workspace/prog/python/vertexai/batch/.venv/lib/python3.11/site-packages/google/cloud/aiplatform_v1/services/job_service/client.py", line 3739, in create_batch_prediction_job
response = rpc(
^^^^
File "/Users/user/workspace/prog/python/vertexai/batch/.venv/lib/python3.11/site-packages/google/api_core/gapic_v1/method.py", line 131, in __call__
return wrapped_func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/user/workspace/prog/python/vertexai/batch/.venv/lib/python3.11/site-packages/google/api_core/grpc_helpers.py", line 78, in error_remapped_callable
raise exceptions.from_grpc_error(exc) from exc
google.api_core.exceptions.InvalidArgument: 400 "bigquery" output format does not support key_field.

```

## Expected Behavior
According to the [documentation](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.types.BatchPredictionJob.InstanceConfig), it seems that specifying `key_field` with a `bigquery` input should be allowed, but the error indicates otherwise.

## Actual Behavior
The job fails with a 400 error stating that the BigQuery output format does not support key_field, which contradicts the information in the documentation.

## Additional Information
If `key_field` is not supported for the `bigquery` format, it would be helpful to update the documentation to reflect this limitation. Otherwise, any guidance on resolving this issue would be greatly appreciated.

Thanks!

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.