googleapis / python-aiplatform

A Python SDK for Vertex AI, a fully managed, end-to-end platform for data science and machine learning.
Apache License 2.0
643 stars 348 forks source link

BatchPredictionJob.create does not support providing a value for instance_config #3031

Open vstrimaitis opened 12 months ago

vstrimaitis commented 12 months ago

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.

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

jackward-countdown commented 2 months ago

included_fields/excluded_fields are also part of instance_config, making these not possible to pass through Model.batch_predict. It would be great to have either instance_config or the individual variables able to be passed through.

tetsu-i commented 1 month ago

While this issue may already have been resolved, I wanted to share my experience using the client library to set up a batch prediction job in Python. I encountered issues where the job ran but execution was not successful. I hope this information can be helpful for anyone facing a similar situation.

Here’s the approach I took using the client library:

from google.cloud import aiplatform_v1

LOCATION = "YOUR_LOCATION"
PROJECT_ID = "YOUR_PROJECT_ID"

def batch_predict_with_gcs(
    model: aiplatform.Model,
    job_display_name: str,
    gcs_source: str,
    gcs_destination_prefix: 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="jsonl",
        gcs_source=aiplatform_v1.GcsSource(
            uris=[gcs_source],
        ),
    )

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

    # https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.types.BatchPredictionJob.OutputConfig
    output_config = aiplatform_v1.BatchPredictionJob.OutputConfig(
        predictions_format="jsonl",
        gcs_destination=aiplatform_v1.GcsDestination(
            output_uri_prefix=gcs_destination_prefix,
        ),
    )

    # 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/{PROJECT_ID}/locations/{LOCATION}",
        batch_prediction_job=job,
    )

    response = client.create_batch_prediction_job(request=request)
    return response

I used the following JSONL data stored in GCS as input:

{"user_id": 1, "date": "2024-10-04", "feature1": 0.8, "feature2": 0.9}
{"user_id": 2, "date": "2024-10-04", "feature1": 0.5, "feature2": 0.7}

Unfortunately, I was unable to find a working example that successfully uses JSONL as input with this method. There may be an issue with my approach, particularly when handling the instance configuration or input format, but I hope this helps others in similar situations.

Additionally, I was able to successfully execute the batch prediction job using BigQuery as the source by modifying other parts of the job for BigQuery input, without specifying the instance_type.