ref(client): Make BatchPushTaskWorker the default

clients/python/src/examples/cli.py

@@ -82,9 +82,6 @@ def scheduler() -> None:
 @click.option(
     "--push-mode", help="Whether to run in PUSH or PULL mode.", default=False, is_flag=True
 )
-@click.option(
-    "--batch-push-mode", help="Whether to run in BATCH PUSH mode.", default=False, is_flag=True
-)
 @click.option(
     "--grpc-port",
     help="Port for the gRPC server to listen on.",
@@ -101,30 +98,14 @@ def worker(
     rpc_host: str,
     concurrency: int,
     push_mode: bool,
-    batch_push_mode: bool,
     grpc_port: int,
     prometheus_port: int | None,
 ) -> None:
-    from taskbroker_client.worker import BatchPushTaskWorker, PushTaskWorker, TaskWorker
+    from taskbroker_client.worker import PushTaskWorker, TaskWorker
 
     click.echo("Starting worker")
-    if batch_push_mode:
-        worker: PushTaskWorker | TaskWorker = BatchPushTaskWorker(
-            app_module="examples.app:app",
-            broker_service=rpc_host,
-            max_child_task_count=100,
-            concurrency=concurrency,
-            child_tasks_queue_maxsize=concurrency * 2,
-            result_queue_maxsize=concurrency * 2,
-            rebalance_after=32,
-            processing_pool_name="examples",
-            process_type="forkserver",
-            grpc_port=grpc_port,
-            update_in_batches=True,
-            prometheus_port=prometheus_port,
-        )
-    elif push_mode:
-        worker = PushTaskWorker(
+    if push_mode:
+        worker: PushTaskWorker | TaskWorker = PushTaskWorker(
             app_module="examples.app:app",
             broker_service=rpc_host,
             max_child_task_count=100,

clients/python/src/taskbroker_client/worker/__init__.py

@@ -1,3 +1,3 @@
-from .worker import BatchPushTaskWorker, PushTaskWorker, TaskWorker
+from .worker import PushTaskWorker, TaskWorker
 
-__all__ = ("TaskWorker", "PushTaskWorker", "BatchPushTaskWorker")
+__all__ = ("TaskWorker", "PushTaskWorker")

clients/python/src/taskbroker_client/worker/push_clients.py

@@ -98,66 +98,6 @@ def _connect_to_host(self, host: str) -> ConsumerServiceStub:
     def emit_health_check(self) -> None:
         self._emit_health_check()
 
-    def update_tasks(self, processing_results: list[ProcessingResult]) -> None:
-        for processing_result in processing_results:
-            self._update_task_single(processing_result)
-
-    def _update_task_single(
-        self,
-        processing_result: ProcessingResult,
-    ) -> None:
-        """
-        Update the status for a given task activation.
-        """
-        self._emit_health_check()
-
-        request = SetTaskStatusRequest(
-            id=processing_result.task_id,
-            status=processing_result.status,
-            fetch_next_task=None,
-            max_attempts=processing_result.max_attempts,
-            delay_on_retry=processing_result.delay_on_retry,
-        )
-
-        retries = 0
-        exception = None
-        while retries < 3:
-            try:
-                with self._metrics.timer(
-                    "taskworker.update_task.rpc",
-                    tags={
-                        "service": self._service,
-                        "processing_pool": self._processing_pool_name,
-                    },
-                ):
-                    self._stub.SetTaskStatus(request)
-                exception = None
-                break
-            except grpc.RpcError as err:
-                exception = err
-                self._metrics.incr(
-                    "taskworker.client.rpc_error",
-                    tags={
-                        "method": "SetTaskStatus",
-                        "status": err.code().name,
-                        "processing_pool": self._processing_pool_name,
-                    },
-                )
-            finally:
-                retries += 1
-
-        if exception:
-            raise exception
-
-
-class BatchPushTaskbrokerClient(PushTaskbrokerClient):
-    """
-    Taskworker RPC client wrapper
-
-    Push brokers are a deployment so they don't need to be connected to individually. There is one service provided
-    that works for all the brokers. This client pushes batches of activation updates.
-    """
-
     def update_tasks(
         self,
         processing_results: list[ProcessingResult],

clients/python/src/taskbroker_client/worker/worker.py

@@ -46,7 +46,7 @@
     TaskbrokerClient,
     parse_rpc_secret_list,
 )
-from taskbroker_client.worker.push_clients import BatchPushTaskbrokerClient, PushTaskbrokerClient
+from taskbroker_client.worker.push_clients import PushTaskbrokerClient
 from taskbroker_client.worker.workerchild import ChildMessage, child_process
 
 if TYPE_CHECKING:
@@ -180,7 +180,6 @@ def __init__(
         health_check_sec_per_touch: float = DEFAULT_WORKER_HEALTH_CHECK_SEC_PER_TOUCH,
         grpc_port: int = 50052,
         push_task_timeout: float = 5,
-        update_in_batches: bool = False,
         skip_awaiting_futures: bool = True,
         warmup_timeout: float = DEFAULT_WORKER_WARMUP_TIMEOUT_SEC,
         prometheus_port: int | None = None,
@@ -209,7 +208,7 @@ def __init__(
             processing_pool_name=processing_pool_name,
             pod_name=pod_name,
             process_type=process_type,
-            update_in_batches=update_in_batches,
+            update_in_batches=True,
             skip_awaiting_futures=skip_awaiting_futures,
             prometheus_port=prometheus_port,
             future_checking_frequency=future_checking_frequency,
@@ -275,47 +274,48 @@ def _send_results(
         """
         Send a result to the broker. If the set has failed before, sleep briefly before retrying.
         """
-        assert (
-            len(results) == 1
-        ), "Only one result can be sent at a time with the regular push client"
-        result = results[0]
+        for result in results:
+            self._metrics.distribution(
+                "taskworker.worker.complete_duration",
+                time.monotonic() - result.receive_timestamp,
+                tags={"processing_pool": self._processing_pool_name},
+            )
         self._metrics.distribution(
-            "taskworker.worker.complete_duration",
-            time.monotonic() - result.receive_timestamp,
+            "taskworker.worker.update_status_batch_size",
+            len(results),
             tags={"processing_pool": self._processing_pool_name},
         )
 
         logger.debug(
-            "taskworker.workers._send_result",
+            "taskworker.send_update_task_batch.batch_sent",
             extra={
-                "task_id": result.task_id,
-                "next": False,  # Push mode doesn't support fetching next tasks
+                "results": [result.task_id for result in results],
                 "processing_pool": self._processing_pool_name,
             },
         )
         # Use the shutdown_event as a sleep mechanism
         self._grpc_sync_event.wait(self._setstatus_backoff_seconds)
 
         try:
-            self.client.update_tasks([result])
+            self.client.update_tasks(results)
             self._setstatus_backoff_seconds = 0
             return None
         except grpc.RpcError as e:
             self._setstatus_backoff_seconds = min(self._setstatus_backoff_seconds + 1, 10)
             logger.warning(
-                "taskworker.send_update_task.failed",
-                extra={"task_id": result.task_id, "error": e},
+                "taskworker.send_update_task_batch.failed",
+                extra={"results": [result.task_id for result in results], "error": e},
             )
             if e.code() != grpc.StatusCode.NOT_FOUND:
                 # If the task was not found, we can't update it, so we should just return None
-                raise RequeueException(f"Failed to update task: {e}")
+                raise RequeueException(f"Failed to update task batch: {e}")
         except HostTemporarilyUnavailable as e:
             self._setstatus_backoff_seconds = min(
                 self._setstatus_backoff_seconds + 4, MAX_BACKOFF_SECONDS_WHEN_HOST_UNAVAILABLE
             )
             logger.info(
-                "taskworker.send_update_task.temporarily_unavailable",
-                extra={"task_id": result.task_id, "error": str(e)},
+                "taskworker.send_update_task_batch.temporarily_unavailable",
+                extra={"task_ids": [result.task_id for result in results], "error": str(e)},
             )
             raise RequeueException(f"Failed to update task: {e}")
 
@@ -505,87 +505,6 @@ def shutdown(self) -> None:
         self.worker_pool.shutdown()
 
 
-class BatchPushTaskWorker(PushTaskWorker):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        assert (
-            kwargs["update_in_batches"] is True
-        ), "BatchPushTaskWorker must be initialized with update_in_batches=True"
-        super().__init__(*args, **kwargs)
-
-    def _create_client(
-        self,
-        service: str,
-        application: str,
-        metrics: MetricsBackend,
-        health_check_settings: HealthCheckSettings | None = None,
-        rpc_secret: str | None = None,
-        grpc_config: str | None = None,
-        processing_pool_name: str | None = None,
-    ) -> PushTaskbrokerClient:
-        return BatchPushTaskbrokerClient(
-            service=service,
-            application=application,
-            metrics=metrics,
-            health_check_settings=health_check_settings,
-            rpc_secret=rpc_secret,
-            grpc_config=grpc_config,
-            processing_pool_name=processing_pool_name,
-        )
-
-    def _send_results(
-        self, results: list[ProcessingResult], is_draining: bool = False
-    ) -> InflightTaskActivation | None:
-        """
-        Send a result to the broker. If the set has failed before, sleep briefly before retrying.
-        """
-        for result in results:
-            self._metrics.distribution(
-                "taskworker.worker.complete_duration",
-                time.monotonic() - result.receive_timestamp,
-                tags={"processing_pool": self._processing_pool_name},
-            )
-        self._metrics.distribution(
-            "taskworker.worker.update_status_batch_size",
-            len(results),
-            tags={"processing_pool": self._processing_pool_name},
-        )
-
-        logger.debug(
-            "taskworker.send_update_task_batch.batch_sent",
-            extra={
-                "results": [result.task_id for result in results],
-                "processing_pool": self._processing_pool_name,
-            },
-        )
-        # Use the shutdown_event as a sleep mechanism
-        self._grpc_sync_event.wait(self._setstatus_backoff_seconds)
-
-        try:
-            self.client.update_tasks(results)
-            self._setstatus_backoff_seconds = 0
-            return None
-        except grpc.RpcError as e:
-            self._setstatus_backoff_seconds = min(self._setstatus_backoff_seconds + 1, 10)
-            logger.warning(
-                "taskworker.send_update_task_batch.failed",
-                extra={"results": [result.task_id for result in results], "error": e},
-            )
-            if e.code() != grpc.StatusCode.NOT_FOUND:
-                # If the task was not found, we can't update it, so we should just return None
-                raise RequeueException(f"Failed to update task batch: {e}")
-        except HostTemporarilyUnavailable as e:
-            self._setstatus_backoff_seconds = min(
-                self._setstatus_backoff_seconds + 4, MAX_BACKOFF_SECONDS_WHEN_HOST_UNAVAILABLE
-            )
-            logger.info(
-                "taskworker.send_update_task_batch.temporarily_unavailable",
-                extra={"task_ids": [result.task_id for result in results], "error": str(e)},
-            )
-            raise RequeueException(f"Failed to update task: {e}")
-
-        return None
-
-
 class TaskWorker:
     """
     A TaskWorker fetches tasks from a taskworker RPC host and handles executing task activations.
@@ -638,6 +557,7 @@ def __init__(
             result_queue_maxsize=result_queue_maxsize,
             processing_pool_name=processing_pool_name,
             process_type=process_type,
+            update_in_batches=False,
             skip_awaiting_futures=skip_awaiting_futures,
             future_checking_frequency=future_checking_frequency,
         )
@@ -837,7 +757,6 @@ def __init__(
 
         self._processing_pool_name = processing_pool_name or "unknown"
         self._pod_name = pod_name or "unknown"
-        self._update_in_batches = update_in_batches
 
         self._send_result_fn = send_result_fn
 
@@ -847,7 +766,7 @@ def __init__(
         self._metrics = app.metrics
         self._skip_awaiting_futures = skip_awaiting_futures
         self._future_checking_frequency = future_checking_frequency
-
+        self._update_in_batches = update_in_batches
         self._mp_context = mp_context
         self._process_type = process_type
 
@@ -1002,6 +921,7 @@ def result_thread() -> None:
                         try:
                             result = self._processed_tasks.get(timeout=1.0)
                             if not self._update_in_batches:
+                                # This needs to stay until the pull taskbroker is removed
                                 executor.submit(self.send_results, [result], False)
                                 break
                             else: