Merge pull request #145 from weaviate/improve-batch-logic

Improve batch logic

Author: jfrancoa

weaviate_cli/commands/create.py

@@ -26,6 +26,7 @@
     CreateDataDefaults,
     CreateRoleDefaults,
     PERMISSION_HELP_STRING,
+    MAX_WORKERS,
 )
 
 
@@ -89,7 +90,7 @@ def create() -> None:
 @click.option(
     "--training_limit",
     default=CreateCollectionDefaults.training_limit,
-    help="Training limit for PQ and SQ (default: 10000).",
+    help=f"Training limit for PQ and SQ (default: {CreateCollectionDefaults.training_limit}).",
 )
 @click.option(
     "--multitenant", is_flag=True, help="Enable multitenancy (default: False)."
@@ -228,7 +229,7 @@ def create_collection_cli(
 @click.option(
     "--number_tenants",
     default=CreateTenantsDefaults.number_tenants,
-    help="Number of tenants to create (default: 100).",
+    help=f"Number of tenants to create (default: {CreateTenantsDefaults.number_tenants}).",
 )
 @click.option(
     "--tenant_batch_size",
@@ -336,7 +337,8 @@ def create_backup_cli(ctx, backend, backup_id, include, exclude, wait, cpu_for_b
 @click.option(
     "--limit",
     default=CreateDataDefaults.limit,
-    help="Number of objects to import (default: 1000).",
+    help=f"Number of objects to import (default: {CreateDataDefaults.limit}).",
+    type=int,
 )
 @click.option(
     "--consistency_level",
@@ -390,6 +392,23 @@ def create_backup_cli(ctx, backend, backup_id, include, exclude, wait, cpu_for_b
     is_flag=True,
     help="Enable multi-vector (default: False).",
 )
+@click.option(
+    "--dynamic_batch",
+    is_flag=True,
+    help="Enable dynamic batching (default: False).",
+)
+@click.option(
+    "--batch_size",
+    default=CreateDataDefaults.batch_size,
+    help=f"Number of objects to ingest in each batch (default: {CreateDataDefaults.batch_size}).",
+    type=int,
+)
+@click.option(
+    "--concurrent_requests",
+    default=MAX_WORKERS,
+    type=int,
+    help=f"Number of concurrent requests to send to the server (default: {MAX_WORKERS}).",
+)
 @click.pass_context
 def create_data_cli(
     ctx,
@@ -406,6 +425,9 @@ def create_data_cli(
     wait_for_indexing,
     verbose,
     multi_vector,
+    dynamic_batch,
+    batch_size,
+    concurrent_requests,
 ):
     """Ingest data into a collection in Weaviate."""
 
@@ -423,6 +445,12 @@ def create_data_cli(
         click.echo("Error: --uuid has no effect unless --limit=1 is enabled.")
         sys.exit(1)
 
+    if dynamic_batch and not randomize:
+        click.echo(
+            "Error: --dynamic_batch has no effect unless --randomize is enabled."
+        )
+        sys.exit(1)
+
     client: Optional[WeaviateClient] = None
     try:
         client = get_client_from_context(ctx)
@@ -442,6 +470,9 @@ def create_data_cli(
             wait_for_indexing=wait_for_indexing,
             verbose=verbose,
             multi_vector=multi_vector,
+            dynamic_batch=dynamic_batch,
+            batch_size=batch_size,
+            concurrent_requests=concurrent_requests,
         )
     except Exception as e:
         click.echo(f"Error: {e}")

weaviate_cli/defaults.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+import multiprocessing
 from typing import Optional, List, Dict
 
 
@@ -50,6 +51,14 @@
 QUERY_MAXIMUM_RESULTS = 10000
 MAX_OBJECTS_PER_BATCH = 5000
 
+try:
+    _CPU_COUNT = multiprocessing.cpu_count()
+except (NotImplementedError, OSError):
+    _CPU_COUNT = None
+# Fallback to 1 worker if CPU count is unavailable or invalid
+_SAFE_CPU_COUNT = _CPU_COUNT if isinstance(_CPU_COUNT, int) and _CPU_COUNT > 0 else 1
+MAX_WORKERS = min(32, _SAFE_CPU_COUNT + 4)
+
 
 @dataclass
 class CreateCollectionDefaults:
@@ -102,6 +111,8 @@ class CreateDataDefaults:
     wait_for_indexing: bool = False
     verbose: bool = False
     multi_vector: bool = False
+    batch_size: int = 1000
+    dynamic_batch: bool = False
 
 
 @dataclass

weaviate_cli/managers/config_manager.py

@@ -9,6 +9,7 @@
 from weaviate.exceptions import WeaviateGRPCUnavailableError
 from pathlib import Path
 from typing import Dict, Optional, Union
+from weaviate.config import AdditionalConfig, Timeout as WeaviateTimeout
 
 
 class ConfigManager:
@@ -107,6 +108,9 @@ def _create_client(
     ) -> Union[weaviate.WeaviateClient, weaviate.WeaviateAsyncClient]:
         """Internal method to create client based on async flag"""
         auth_config: Optional[weaviate.auth.AuthCredentials] = None
+        additional_config: Optional[AdditionalConfig] = (
+            self.__additional_config_if_slow()
+        )
 
         if "auth" in self.config:
             if self.config["auth"].get("type") == "user":
@@ -133,6 +137,7 @@ def _create_client(
                 grpc_port=self.config["grpc_port"],
                 auth_credentials=auth_config,
                 headers=headers,
+                additional_config=additional_config,
             )
             if async_client:
                 return weaviate.use_async_with_local(**common_kwargs)
@@ -144,6 +149,7 @@ def _create_client(
                 cluster_url=self.config["host"],
                 auth_credentials=auth_config,
                 headers=headers,
+                additional_config=additional_config,
             )
             if async_client:
                 try:
@@ -189,6 +195,7 @@ def _create_client(
                 grpc_port=grpc_port,
                 auth_credentials=auth_config,
                 headers=headers,
+                additional_config=additional_config,
             )
         else:
             common_kwargs = dict(
@@ -200,8 +207,25 @@ def _create_client(
                 grpc_port=self.config["grpc_port"],
                 auth_credentials=auth_config,
                 headers=headers,
+                additional_config=additional_config,
             )
 
         if async_client:
             return weaviate.use_async_with_custom(**common_kwargs)
         return weaviate.connect_to_custom(**common_kwargs)
+
+    def __additional_config_if_slow(self) -> Optional[AdditionalConfig]:
+        """Return AdditionalConfig with doubled timeouts when SLOW_CONNECTION is set."""
+        env_flag = os.getenv("SLOW_CONNECTION")
+        if not env_flag:
+            return None
+        if str(env_flag).strip().lower() not in {"1", "true", "yes", "on"}:
+            return None
+        defaults = AdditionalConfig().timeout
+        return AdditionalConfig(
+            timeout=WeaviateTimeout(
+                query=defaults.query * 2,
+                insert=defaults.insert * 2,
+                init=defaults.init * 2,
+            )
+        )