Generative AI Google Brain

gemini-batch-prediction-monitor-batch-prediction-gemini-api

3 min read

Monitor batch prediction with Gemini API

Source notebook

Repo path: gemini/batch-prediction/monitor_batch_prediction_gemini_api.ipynb · Open on GitHub · intermediate

Orchestrates and monitors Gemini batch predictions with Vertex AI Pipelines and BigQuery output.

Summary

The notebook shows how to compensate for Gemini API batch prediction jobs lacking built-in completion notifications by wrapping the job in Vertex AI Pipelines. It prepares Cloud Storage and BigQuery destinations, defines custom Kubeflow Pipelines components to submit and poll a Gemini batch prediction job, then samples the resulting BigQuery table into a Markdown artifact. The workflow compiles a pipeline YAML, runs an aiplatform.PipelineJob, sends email notification through VertexNotificationEmailOp, and includes cleanup steps.

Key code patterns

Create BigQuery output URI

from google.cloud import bigquery
from datetime import datetime
 
def create_bq_table(dataset_id, project_id=PROJECT_ID, location=LOCATION):
    bq_client = bigquery.Client(project=project_id, location=location)
    dataset = bigquery.Dataset(f"{project_id}.{dataset_id}")
    dataset.location = location
    bq_client.create_dataset(dataset, exists_ok=True, timeout=30)
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    return f"bq://{project_id}.{dataset_id}.prediction_result_{timestamp}"

Creates or reuses the BigQuery dataset and gives each batch prediction output a timestamped table URI.

Submit and poll batch job

vertexai.init(project=project, location=location)
job = BatchPredictionJob.submit(
    source_model=model_id,
    input_dataset=input_bq_table,
    output_uri_prefix=output_bq_table,
)
while not job.has_ended:
    time.sleep(60)
    job.refresh()
if not job.has_succeeded:
    sys.exit(1)
return NamedTuple("outputs", dataset_uri=str)(job.output_location)

Runs Gemini batch prediction from BigQuery input and turns polling status into a pipeline component result or failure.

Visualize result sample

output_bq_table = output_bq_table.replace("bq://", "")
query = f"""
    SELECT *
    FROM `{output_bq_table}`
    LIMIT {sample_size}
"""
df = client.query(query).to_dataframe()
processed_records = [extract_text(r) for r in df.to_dict("records")]
processed_df = pd.DataFrame(processed_records)
with open(output_markdown_table.path, "w") as f:
    f.write(markdown_table)

Queries a limited BigQuery sample and writes request and response text as a Markdown pipeline artifact.

Pipeline exit notification

@dsl.pipeline(name="genai-batch-prediction-pipeline")
def pipeline(...):
    notify_email_task = VertexNotificationEmailOp(recipients=recipients)
    create_input_dataset_task = TabularDatasetCreateOp(...)
    with dsl.ExitHandler(notify_email_task, name="Notification handler"):
        run_batch_prediction_task = GenAIModelBatchPredictOp(...).after(create_input_dataset_task)
        VisualizeBatchPredictionTable(...).after(run_batch_prediction_task)

Uses an ExitHandler so the Vertex AI Pipeline can send completion notification around the batch workflow.

Compile and run pipeline

compiler.Compiler().compile(
    pipeline_func=pipeline,
    package_path="pipeline.yaml",
)
job = aiplatform.PipelineJob(
    display_name="census-demo-pipeline",
    parameter_values=parameter_values,
    template_path="pipeline.yaml",
    pipeline_root=PIPELINE_ROOT,
)
job.run()

Separates pipeline definition from execution by compiling YAML and launching it as a Vertex AI PipelineJob.

Models & APIs used

  • Models: gemini-2.0-flash
  • APIs / services: Vertex AI, Vertex AI Pipelines, BigQuery, Cloud Storage, Artifact Registry
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-bigquery, kfp, google-cloud-pipeline-components, vertexai, pandas

When to use this

Use this pattern when Gemini batch predictions need managed orchestration, polling, BigQuery outputs, and email notification through Vertex AI Pipelines.

Gotchas & caveats

  • The Google Cloud project must enable aiplatform.googleapis.com and artifactregistry.googleapis.com.
  • The notebook creates a Cloud Storage bucket with gsutil and uses it as the Vertex AI staging bucket and pipeline root.
  • The compute service account needs roles/aiplatform.user, roles/storage.objectAdmin, and roles/bigquery.dataEditor.
  • Colab authentication is only run when google.colab is present.
  • The monitoring loop polls every 60 seconds with job.refresh(), so completion is not instant.
  • Visualization assumes request.contents[0].parts[0].text and response.candidates[0].content.parts[0].text exist; missing fields become empty strings.
  • Cleanup flags can delete the pipeline job, Cloud Storage bucket, and BigQuery dataset.

Best practices

  • Initialize Vertex AI with project, location, and staging_bucket before running the pipeline.
  • Use timestamped BigQuery output table names to avoid collisions across batch runs.
  • Wrap the batch job in a Vertex AI Pipeline and use VertexNotificationEmailOp because Gemini batch prediction lacks built-in completion notifications.
  • Poll BatchPredictionJob until has_ended and fail the component when has_succeeded is false.
  • Sample prediction results from BigQuery instead of loading the whole output table.
  • Escape pipe characters and newlines before writing Markdown table output.
  • Include cleanup controls for the pipeline job, bucket, and BigQuery dataset.

Graph View

Backlinks