With Portkey, you can perform batch inference operations with Vertex AI models. This is the most efficient way to:

  • Process large volumes of data with Vertex AI models
  • Test your data with different foundation models
  • Perform A/B testing with different foundation models

Upload a file for batch inference

from portkey_ai import Portkey

# Initialize the Portkey client
portkey = Portkey(
    api_key="PORTKEY_API_KEY", # Replace with your Portkey API key
    virtual_key="VERTEX_VIRTUAL_KEY", # Add your Vertex virtual key
    vertex_storage_bucket_name="your_bucket_name", # Specify the GCS bucket name
    provider_file_name="your_file_name.jsonl", # Specify the file name in GCS
    provider_model="gemini-1.5-flash-001" # Specify the model to use
)

# Upload a file for batch inference
file = portkey.files.create(
    file=open("dataset.jsonl", "rb"),
    purpose="batch"
)

print(file)

Create a batch job

from portkey_ai import Portkey

# Initialize the Portkey client
portkey = Portkey(
    api_key="PORTKEY_API_KEY", # Replace with your Portkey API key
    virtual_key="VERTEX_VIRTUAL_KEY" # Add your Vertex virtual key
)

# Create a batch inference job
batch_job = portkey.batches.create(
    input_file_id="<file_id>", # File ID from the upload step
    endpoint="/v1/chat/completions", # API endpoint to use
    completion_window="24h", # Time window for completion
    model="gemini-1.5-flash-001"
)

print(batch_job)

List batch jobs

from portkey_ai import Portkey

# Initialize the Portkey client
portkey = Portkey(
    api_key="PORTKEY_API_KEY", # Replace with your Portkey API key
    virtual_key="VERTEX_VIRTUAL_KEY" # Add your Vertex virtual key
)

# List all batch jobs
jobs = portkey.batches.list(
    limit=10  # Optional: Number of jobs to retrieve (default: 20)
)

print(jobs)

Get a batch job

from portkey_ai import Portkey

# Initialize the Portkey client
portkey = Portkey(
    api_key="PORTKEY_API_KEY", # Replace with your Portkey API key
    virtual_key="VERTEX_VIRTUAL_KEY" # Add your Vertex virtual key
)

# Retrieve a specific batch job
job = portkey.batches.retrieve(
    "job_id"  # The ID of the batch job to retrieve
)

print(job)

Get batch job output

curl -X GET --header 'x-portkey-api-key: <portkey_api_key>' \
--header 'x-portkey-virtual-key: <vertex_virtual_key>' \
'https://api.portkey.ai/v1/batches/<job_id>/output'

Was this page helpful?