Skip to content

refactor(alloydb): migrate batch processing notebook to genai SDK #14537

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
XrossFox wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

refactor(alloydb): migrate batch processing notebook to genai SDK #14537

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5de36ae
refactor(alloydb): migrate batch processing notebook to genai SDK
XrossFox Aug 25, 2026
9e96091
Addressed a comment.
XrossFox Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 97 additions & 80 deletions alloydb/notebooks/embeddings_batch_processing.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,8 @@
" google-cloud-alloydb-connector[asyncpg]==1.4.0 \\\n",
" sqlalchemy==2.0.36 \\\n",
" pandas==2.2.3 \\\n",
" vertexai==1.70.0 \\\n",
" google-cloud-aiplatform==1.165.1 \\\n",
" google-genai==2.19.0 \\\n",
" asyncio==3.4.3 \\\n",
" greenlet==3.1.1 \\\n",
" --quiet"
Expand Down Expand Up @@ -792,119 +793,133 @@
},
"outputs": [],
"source": [
"from google.api_core.exceptions import ResourceExhausted\n",
"from typing import Union\n",
"from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n",
"from typing import Any, AsyncIterator, List, Optional, Union\n",
"\n",
"from google import genai\n",
"from google.genai import types\n",
"\n",
"async def embed_text(\n",
" batch_data: List[dict[str, Any]],\n",
" model: TextEmbeddingModel,\n",
" cols_to_embed: List[str],\n",
" client: genai.Client,\n",
" model_name: str = \"text-embedding-004\",\n",
" task_type: str = \"SEMANTIC_SIMILARITY\",\n",
" retries: int = 100,\n",
" delay: int = 30,\n",
") -> List[dict[str, Union[List[float], str]]]:\n",
" \"\"\"Embeds text data from a batch of records using a Vertex AI embedding model.\n",
" \"\"\"Embeds text data from a batch of records using the google-genai SDK.\n",
"\n",
" Args:\n",
" batch_data: A data batch containing records with text data to embed.\n",
" model: The Vertex AI `TextEmbeddingModel` to use for generating embeddings.\n",
" cols_to_embed: A list of column names containing the data to be embedded.\n",
" task_type: The task type for the embedding model. Defaults to\n",
" \"SEMANTIC_SIMILARITY\".\n",
" Supported task types: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types\n",
" retries: The maximum number of times to retry embedding generation in case\n",
" of errors. Defaults to 100.\n",
" delay: The delay in seconds between retries. Defaults to 30.\n",
" batch_data: A data batch containing records with text data to embed.\n",
" cols_to_embed: A list of column names containing the data to be embedded.\n",
" model_name: The embedding model ID.\n",
" task_type: The task type for the embedding model.\n",
" retries: The maximum number of retry attempts in case of errors.\n",
" delay: The delay in seconds between retries.\n",
" client: An optional genai.Client instance for connection reuse.\n",
"\n",
" Returns:\n",
" A list of records containing ids and embeddings.\n",
" Example:\n",
" [\n",
" {\n",
" 'id': 'id1',\n",
" 'col1_embedding': [1.0, 1.1, ...],\n",
" 'col2_embedding': [2.0, 2.1, ...],\n",
" ...\n",
" },\n",
" ...\n",
" ]\n",
" where col1 and col2 are columns containing data to be embedded.\n",
" A list of records containing IDs and mapped column embeddings.\n",
" Raises:\n",
Comment thread
XrossFox marked this conversation as resolved.
" Exception: Raises the encountered exception if all retries fail.\n",
" Exception: Raises the encountered exception if all retries fail.\n",
" \"\"\"\n",
" logger = logging.getLogger(\"embed_objects\")\n",
" global total_char_count\n",
"\n",
" # Place all of the embeddings into a single list\n",
" # Extract non-empty text strings to embed\n",
" inputs = []\n",
" for row in batch_data:\n",
" for col in cols_to_embed:\n",
" if col in row and row[col]:\n",
" inputs.append(TextEmbeddingInput(row[col], task_type))\n",
" inputs.append(str(row[col]))\n",
"\n",
" # Retry loop\n",
" for attempt in range(retries):\n",
" try:\n",
" # Get embeddings for the text data\n",
" embeddings = await model.get_embeddings_async(inputs)\n",
" # Asynchronous API call using client_instance.aio\n",
" response = await client.aio.models.embed_content(\n",
" model=model_name,\n",
" contents=inputs,\n",
" config=types.EmbedContentConfig(\n",
" task_type=task_type,\n",
" ),\n",
" )\n",
"\n",
" # Track character metrics\n",
" total_char_count += sum(len(text) for text in inputs)\n",
"\n",
" # Increase total char count\n",
" total_char_count += sum([len(input.text) for input in inputs])\n",
" # Map response embeddings back to dataset record structure\n",
" embeddings_list = (\n",
" response.embeddings\n",
" if response.embeddings\n",
" else [response.embedding]\n",
" )\n",
Comment thread
XrossFox marked this conversation as resolved.
" embedding_iter = iter(embeddings_list)\n",
"\n",
" # group the results together by id\n",
" embedding_iter = iter(embeddings)\n",
" results = []\n",
" for row in batch_data:\n",
" r = {\"id\": row[\"id\"]}\n",
" r = {\"id\": row.get(\"id\")}\n",
" for col in cols_to_embed:\n",
" if col in row and row[col]:\n",
" r[f\"{col}_embedding\"] = str(next(embedding_iter).values)\n",
" embedding_obj = next(embedding_iter)\n",
" r[f\"{col}_embedding\"] = str(embedding_obj.values)\n",
" else:\n",
" r[f\"{col}_embedding\"] = None\n",
" results.append(r)\n",
"\n",
" return results\n",
"\n",
" except Exception as e:\n",
" if attempt < retries - 1: # Retry only if attempts are left\n",
" logger.warning(f\"Error: {e}. Retrying in {delay} seconds...\")\n",
" await asyncio.sleep(delay) # Wait before retrying\n",
" if attempt < retries - 1:\n",
" logger.warning(\n",
" f\"Error: {e}. Retrying in {delay} seconds (attempt {attempt + 1}/{retries})...\"\n",
" )\n",
" await asyncio.sleep(delay)\n",
" else:\n",
" logger.error(f\"Failed to get embeddings for data: {batch_data} after {retries} attempts.\")\n",
" logger.error(\n",
" f\"Failed to get embeddings after {retries} attempts: {e}\"\n",
" )\n",
" raise e\n",
"\n",
" return []\n",
"\n",
"\n",
"async def embed_objects_concurrently(\n",
" cols_to_embed: List[str],\n",
" batch_data: AsyncIterator[List[dict[str, Any]]],\n",
" model: TextEmbeddingModel,\n",
" task_type: str,\n",
" client: genai.Client,\n",
" model_name: str = \"text-embedding-004\",\n",
" task_type: str = \"SEMANTIC_SIMILARITY\",\n",
" max_concurrency: int = 5,\n",
") -> AsyncIterator[List[dict[str, Union[str, List[float]]]]]:\n",
" \"\"\"Embeds text data concurrently from an asynchronous batch data generator.\n",
"\n",
" Args:\n",
" cols_to_embed: A list of column names containing the data to be embedded.\n",
" batch_data: A data batch containing records with text data to embed.\n",
" model: The Vertex AI `TextEmbeddingModel` to use for generating embeddings.\n",
" task_type: The task type for the embedding model.\n",
" Supported task types: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types\n",
" max_concurrency: The maximum number of embedding tasks to run concurrently.\n",
" Defaults to 5.\n",
" cols_to_embed: A list of column names containing the data to be embedded.\n",
" batch_data: An async generator yielding data batches with records.\n",
" model_name: The embedding model ID.\n",
" task_type: The task type for the embedding model.\n",
" max_concurrency: The maximum number of concurrent tasks to execute.\n",
" client: An optional genai.Client instance for connection reuse.\n",
"\n",
" Yields:\n",
" A list of records containing ids and embeddings.\n",
" A list of records containing IDs and mapped column embeddings.\n",
" \"\"\"\n",
" logger = logging.getLogger(\"embed_objects\")\n",
"\n",
" # Keep track of pending tasks\n",
" pending: set[asyncio.Task] = set()\n",
" has_next = True\n",
"\n",
" while pending or has_next:\n",
" while len(pending) < max_concurrency and has_next:\n",
" try:\n",
" data = await batch_data.__anext__()\n",
" coro = embed_text(data, model, cols_to_embed, task_type)\n",
" coro = embed_text(\n",
" batch_data=data,\n",
" cols_to_embed=cols_to_embed,\n",
" model_name=model_name,\n",
" task_type=task_type,\n",
" client=client,\n",
" )\n",
" pending.add(asyncio.ensure_future(coro))\n",
" except StopAsyncIteration:\n",
" has_next = False\n",
Expand All @@ -915,7 +930,9 @@
" )\n",
" for task in done:\n",
" result = task.result()\n",
" logger.info(f\"Embedding task completed: Processed {len(result)} rows.\")\n",
" logger.info(\n",
" f\"Embedding task completed: Processed {len(result)} rows.\"\n",
" )\n",
" yield result"
]
},
Expand Down Expand Up @@ -1045,9 +1062,9 @@
},
"outputs": [],
"source": [
"import vertexai\n",
"import time\n",
"from vertexai.language_models import TextEmbeddingModel\n",
"\n",
"from google import genai\n",
"\n",
"### Define variables ###\n",
"\n",
Expand Down Expand Up @@ -1080,40 +1097,40 @@
"):\n",
" \"\"\"Orchestrates the end-to-end workflow for generating and storing embeddings.\n",
"\n",
" The workflow includes the following major steps:\n",
"\n",
" 1. Data Retrieval: Fetches data from the database that requires embedding.\n",
" 2. Batching: Divides the data into batches for optimized processing.\n",
" 3. Embedding Generation: Generates embeddings concurrently for the batched\n",
" data using the Vertex AI model.\n",
" 4. Database Update: Updates the database concurrently with the generated\n",
" embeddings.\n",
" Workflow Steps:\n",
" 1. Connection Initialization: Initializes DB pool.\n",
" 2. Client Initialization: Instantiates shared GenAI client for workflow lifetime.\n",
" 3. Data Retrieval & Batching: Fetches and chunks source data into async streams.\n",
" 4. Embedding Generation: Uses google-genai client asynchronously and concurrently.\n",
" 5. Database Update: Stores output embeddings concurrently in the database.\n",
"\n",
" Args:\n",
" pool_size: The size of the database connection pool. Defaults to 10.\n",
" embed_data_concurrency: The maximum number of concurrent tasks for generating embeddings.\n",
" Defaults to 20.\n",
" batch_update_concurrency: The maximum number of concurrent tasks for updating the database.\n",
" Defaults to 10.\n",
" pool_size: The size of the database connection pool.\n",
" embed_data_concurrency: Max concurrent tasks for generating embeddings.\n",
" batch_update_concurrency: Max concurrent tasks for database updates.\n",
" \"\"\"\n",
" # Set up connections to the database\n",
" # Set up database connection pool\n",
" pool = await init_connection_pool(connector, database_name, pool_size=pool_size)\n",
"\n",
" # Initialise VertexAI and the model to be used to generate embeddings\n",
" vertexai.init(project=project_id, location=region)\n",
" model = TextEmbeddingModel.from_pretrained(model_name)\n",
" # Initialize single GenAI client instance for the top-level workflow lifetime\n",
" client = genai.Client(vertexai=True, project=project_id, location=region)\n",
"\n",
" start_time = time.monotonic()\n",
"\n",
" # Fetch source data from the database\n",
" # Fetch source data from database\n",
" source_data = get_source_data(pool, cols_to_embed)\n",
"\n",
" # Divide the source data into batches for efficient processing\n",
" # Divide source data into asynchronous batches\n",
" batch_data = batch_source_data(source_data, cols_to_embed)\n",
"\n",
" # Generate embeddings for the batched data concurrently\n",
" # Generate embeddings concurrently using the GenAI SDK\n",
" embeddings_data = embed_objects_concurrently(\n",
" cols_to_embed, batch_data, model, task, max_concurrency=embed_data_concurrency\n",
" cols_to_embed=cols_to_embed,\n",
" batch_data=batch_data,\n",
" model_name=model_name,\n",
" task_type=task,\n",
" max_concurrency=embed_data_concurrency,\n",
" client=client,\n",
" )\n",
"\n",
" # Update the database with the generated embeddings concurrently\n",
Expand All @@ -1124,7 +1141,7 @@
" end_time = time.monotonic()\n",
" elapsed_time = end_time - start_time\n",
"\n",
" # Release database connections and close the connector\n",
" # Release database resources\n",
" await pool.dispose()\n",
" await connector.close()\n",
"\n",
Expand Down