feat: add smarter batching for incremental sync based on updated parent_ids (#5862)

Author: sahil485

servers/fai-reindexing/src/services/sync.ts

@@ -45,7 +45,7 @@ export async function syncToQueryIndexIncremental(domain: string, parentIds: str
 }
 
 async function pollJobStatus(jobId: string, domain: string): Promise<void> {
-    const log = createDomainLogger(domain);
+    const logger = createDomainLogger(domain);
     const maxAttempts = 1000;
     let attempts = 0;
 
@@ -54,15 +54,17 @@ async function pollJobStatus(jobId: string, domain: string): Promise<void> {
             const statusResponse = await faiClient.index.getJobStatus(jobId);
             const { status, success, error } = statusResponse;
 
-            log.info("Job status", { jobId, status, success, error, attempt: attempts + 1, maxAttempts });
+            logger.info("Job status", { jobId, status, success, error, attempt: attempts + 1, maxAttempts });
 
             if (status === "completed") {
                 if (success === false) {
+                    logger.error("Sync job completed but failed", { jobId, error });
                     const errorMessage = typeof error === "object" ? JSON.stringify(error) : error || "Unknown error";
-                    throw new Error(`Sync job completed but failed: ${errorMessage}`);
+                    throw new Error(`Sync job failed: ${errorMessage}`);
                 }
                 return;
             } else if (status === "failed") {
+                logger.error("Sync job failed", { jobId, error });
                 const errorMessage = typeof error === "object" ? JSON.stringify(error) : error || "Unknown error";
                 throw new Error(`Sync job failed: ${errorMessage}`);
             }
@@ -71,16 +73,19 @@ async function pollJobStatus(jobId: string, domain: string): Promise<void> {
             attempts++;
         } catch (error) {
             const errorMessage = error instanceof Error ? error.message : String(error);
+            if (errorMessage.includes("Sync job failed:")) {
+                throw error;
+            }
             attempts++;
             if (attempts >= maxAttempts) {
-                log.error("Error polling job status - max attempts exceeded", {
+                logger.error("Error polling job status - max attempts exceeded", {
                     error: errorMessage,
                     attempts,
                     maxAttempts
                 });
                 throw error;
             }
-            log.warn("Error polling job status - retrying", { error: errorMessage, attempt: attempts, maxAttempts });
+            logger.warn("Error polling job status - retrying", { error: errorMessage, attempt: attempts, maxAttempts });
             await new Promise((resolve) => setTimeout(resolve, RETRY_CONFIG.SYNC_RETRY_DELAY_MS));
         }
     }

servers/fai/src/fai/utils/turbopuffer/sync.py

@@ -6,6 +6,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 from turbopuffer import (
     NOT_GIVEN,
+    APITimeoutError,
     AsyncTurbopuffer,
 )
 from turbopuffer.types.row import Row
@@ -483,63 +484,95 @@ async def sync_index_to_target_incremental(
     async with AsyncTurbopuffer(
         region=CONFIG.TURBOPUFFER_DEFAULT_REGION,
         api_key=VARIABLES.TURBOPUFFER_API_KEY,
+        max_retries=0,
+        timeout=30.0,
     ) as tpuf_client:
         source_ns = tpuf_client.namespace(source_namespace_id)
+        await source_ns.hint_cache_warm()
         target_ns = tpuf_client.namespace(target_namespace_id)
 
-        try:
-            await target_ns.write(
-                delete_by_filter=("And", [("source", "Eq", source_index_name), ("parent_id", "In", parent_ids)])
-            )
-            LOGGER.info(f"Deleted old records for {len(parent_ids)} parent_ids from target")
-        except Exception as e:
-            LOGGER.warning(f"Batch delete failed, trying individual deletes: {e}")
-            for parent_id in parent_ids:
-                try:
-                    await target_ns.write(
-                        delete_by_filter=("And", [("source", "Eq", source_index_name), ("parent_id", "Eq", parent_id)])
+        delete_parent_id_batch_size = len(parent_ids)
+        while delete_parent_id_batch_size >= 1:
+            try:
+                total_deleted = 0
+                for i in range(0, len(parent_ids), delete_parent_id_batch_size):
+                    batch = parent_ids[i:min(i+delete_parent_id_batch_size, len(parent_ids))]
+                    result = await target_ns.write(
+                        delete_by_filter=("And", [("source", "Eq", source_index_name), ("parent_id", "In", batch)])
                     )
-                except Exception as delete_error:
-                    LOGGER.warning(f"Failed to delete parent_id {parent_id}: {delete_error}")
+                    LOGGER.info(
+                        f"Deleted {result.rows_deleted or 0} records for {len(batch)} parent_ids"
+                    )
+                    total_deleted += result.rows_deleted or 0
+                LOGGER.info(
+                    f"Successfully deleted {total_deleted} records from {target_namespace_id}"
+                )
+                break
+            except APITimeoutError:
+                LOGGER.warning(
+                    f"Batch delete with size {delete_parent_id_batch_size} timed out, "
+                    f"trying batch size {delete_parent_id_batch_size // 2}"
+                )
+                delete_parent_id_batch_size = delete_parent_id_batch_size // 2
+                if delete_parent_id_batch_size < 1:
+                    LOGGER.error("Failed to delete records even with batch size 1")
+                    raise
 
         source_ns_exists = await source_ns.exists()
         if not source_ns_exists:
             LOGGER.warning(f"Source namespace {source_namespace_id} does not exist")
             return
 
         total_synced = 0
-        last_id = None
-
-        while True:
-            filter_conditions = [("parent_id", "In", parent_ids)]
-            if last_id is not None:
-                filter_conditions.append(("id", "Gt", last_id))
-
-            result = await source_ns.query(
-                rank_by=("id", "asc"),
-                top_k=10000,
-                include_attributes=True,
-                filters=("And", filter_conditions),
-            )
-
-            if not result.rows:
-                break
+        sync_parent_id_batch_size = len(parent_ids)
 
-            prefixed_rows = []
-            for row in result.rows:
-                new_row = Row.from_dict(row.model_dump())
-                new_row.id = prefixed_id(source_namespace_id, row.id)
-                new_row.source = source_index_name
-                prefixed_rows.append(new_row)
-
-            await target_ns.write(
-                upsert_rows=prefixed_rows, distance_metric="cosine_distance", schema=get_query_index_tpuf_schema()
-            )
-            total_synced += len(prefixed_rows)
-            LOGGER.info(f"Synced batch: {len(prefixed_rows)} records (total: {total_synced})")
-
-            if len(result.rows) < 10000:
+        while sync_parent_id_batch_size >= 1:
+            try:
+                for parent_idx in range(0, len(parent_ids), sync_parent_id_batch_size):
+                    last_id = None
+                    parent_id_batch = parent_ids[parent_idx:min(parent_idx+sync_parent_id_batch_size, len(parent_ids))]
+                    filter_conditions = [("parent_id", "In", parent_id_batch)]
+                    while True:
+                        if last_id is not None:
+                            filter_conditions.append(("id", "Gt", last_id))
+
+                        result = await source_ns.query(
+                            rank_by=("id", "asc"),
+                            top_k=1000,
+                            include_attributes=True,
+                            filters=("And", filter_conditions),
+                        )
+
+                        if not result.rows:
+                            break
+
+                        prefixed_rows = []
+                        for row in result.rows:
+                            new_row = Row.from_dict(row.model_dump())
+                            new_row.id = prefixed_id(source_namespace_id, row.id)
+                            new_row.source = source_index_name
+                            prefixed_rows.append(new_row)
+
+                        await target_ns.write(
+                            upsert_rows=prefixed_rows,
+                            distance_metric="cosine_distance",
+                            schema=get_query_index_tpuf_schema(),
+                        )
+                        total_synced += len(prefixed_rows)
+                        LOGGER.info(f"Synced batch: {len(prefixed_rows)} records (total: {total_synced})")
+
+                        last_id = result.rows[-1].id
+                        if len(result.rows) < 1000:
+                            break
                 break
-            last_id = result.rows[-1].id
+            except APITimeoutError:
+                LOGGER.warning(
+                    f"Batch sync with size {sync_parent_id_batch_size} timed out, "
+                    f"trying batch size {sync_parent_id_batch_size // 2}"
+                )
+                sync_parent_id_batch_size = sync_parent_id_batch_size // 2
+                if sync_parent_id_batch_size < 1:
+                    LOGGER.error("Failed to sync records even with batch size 1")
+                    raise
 
         LOGGER.info(f"Incremental sync completed: {total_synced} total records synced to {target_namespace_id}")