diff --git a/clients/python/src/examples/cli.py b/clients/python/src/examples/cli.py index 5e9d4e46..64f93d0b 100644 --- a/clients/python/src/examples/cli.py +++ b/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, diff --git a/clients/python/src/taskbroker_client/worker/__init__.py b/clients/python/src/taskbroker_client/worker/__init__.py index 79f70886..d94f2b62 100644 --- a/clients/python/src/taskbroker_client/worker/__init__.py +++ b/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") diff --git a/clients/python/src/taskbroker_client/worker/push_clients.py b/clients/python/src/taskbroker_client/worker/push_clients.py index b0f95e0e..f41fb36c 100644 --- a/clients/python/src/taskbroker_client/worker/push_clients.py +++ b/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], diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 7f89d5d1..0cce2164 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/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,21 +274,22 @@ 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, }, ) @@ -297,25 +297,25 @@ def _send_results( 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: diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 560fbad4..7253a74e 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -44,7 +44,6 @@ from taskbroker_client.types import InflightTaskActivation, ProcessingResult from taskbroker_client.worker.producer import TaskProducer, _pending_futures from taskbroker_client.worker.worker import ( - BatchPushTaskWorker, PushTaskWorker, TaskWorker, TaskWorkerProcessingPool, @@ -339,7 +338,7 @@ def _make_result_thread_pool( *, concurrency: int = 3, result_queue_maxsize: int = 3, - update_in_batches: bool, + update_in_batches: bool = False, ) -> TaskWorkerProcessingPool: return TaskWorkerProcessingPool( app_module="examples.app:app", @@ -349,8 +348,8 @@ def _make_result_thread_pool( concurrency=concurrency, result_queue_maxsize=result_queue_maxsize, processing_pool_name="test", - process_type="fork", update_in_batches=update_in_batches, + process_type="fork", ) @@ -648,7 +647,7 @@ def test_result_thread_flushes_partial_batch_on_queue_empty(self) -> None: def test_result_thread_sends_results_individually_without_batching(self) -> None: capture = _SendResultCapture() - pool = _make_result_thread_pool(capture, update_in_batches=False) + pool = _make_result_thread_pool(capture) try: pool.start_result_thread() @@ -723,19 +722,6 @@ def test_constructor_push_mode(self) -> None: self.assertTrue(taskworker.client is not None) self.assertEqual(taskworker._grpc_port, 50099) - def test_constructor_batch_push_mode(self) -> None: - taskworker = BatchPushTaskWorker( - app_module="examples.app:app", - broker_service="127.0.0.1:50051", - max_child_task_count=100, - process_type="fork", - grpc_port=50099, - update_in_batches=True, - ) - - self.assertTrue(taskworker.client is not None) - self.assertEqual(taskworker._grpc_port, 50099) - def test_push_worker_health_check_touches_while_idle(tmp_path: Path) -> None: taskworker = PushTaskWorker( @@ -760,30 +746,6 @@ def test_push_worker_health_check_touches_while_idle(tmp_path: Path) -> None: assert taskworker._health_check_thread is None -def test_batch_push_worker_health_check_touches_while_idle(tmp_path: Path) -> None: - taskworker = BatchPushTaskWorker( - app_module="examples.app:app", - broker_service="127.0.0.1:50051", - max_child_task_count=100, - process_type="fork", - health_check_file_path=str(tmp_path / "health"), - health_check_sec_per_touch=0.01, - update_in_batches=True, - ) - - with mock.patch.object(taskworker.client, "emit_health_check") as mock_emit: - taskworker._start_health_check_thread() - try: - start = time.time() - while mock_emit.call_count < 2 and time.time() - start < 1: - time.sleep(0.01) - finally: - taskworker._stop_health_check_thread() - - assert mock_emit.call_count >= 2 - assert taskworker._health_check_thread is None - - def _make_push_worker(**kwargs: Any) -> PushTaskWorker: return PushTaskWorker( app_module="examples.app:app", @@ -1036,33 +998,6 @@ def test_push_task_success(self) -> None: self.assertEqual(inflight.activation.id, SIMPLE_TASK.activation.id) self.assertEqual(inflight.host, "broker-host:50051") - def test_batch_push_task_success(self) -> None: - taskworker = BatchPushTaskWorker( - app_module="examples.app:app", - broker_service="127.0.0.1:50051", - max_child_task_count=100, - process_type="fork", - update_in_batches=True, - ) - with mock.patch.object( - taskworker.worker_pool, "push_task", return_value=True - ) as mock_push_task: - request = PushTaskRequest( - task=SIMPLE_TASK.activation, - callback_url="broker-host:50051", - ) - mock_context = mock.MagicMock() - servicer = WorkerServicer(taskworker.worker_pool) - - response = servicer.PushTask(request, mock_context) - - self.assertIsInstance(response, PushTaskResponse) - mock_context.abort.assert_not_called() - mock_push_task.assert_called_once_with(mock.ANY, timeout=5) - (inflight,) = mock_push_task.call_args[0] - self.assertEqual(inflight.activation.id, SIMPLE_TASK.activation.id) - self.assertEqual(inflight.host, "broker-host:50051") - def test_push_task_worker_busy(self) -> None: taskworker = PushTaskWorker( app_module="examples.app:app", @@ -1085,29 +1020,6 @@ def test_push_task_worker_busy(self) -> None: grpc.StatusCode.RESOURCE_EXHAUSTED, "worker busy" ) - def test_batch_push_task_worker_busy(self) -> None: - taskworker = BatchPushTaskWorker( - app_module="examples.app:app", - broker_service="127.0.0.1:50051", - max_child_task_count=100, - process_type="fork", - child_tasks_queue_maxsize=1, - update_in_batches=True, - ) - with mock.patch.object(taskworker.worker_pool, "push_task", return_value=False): - request = PushTaskRequest( - task=SIMPLE_TASK.activation, - callback_url="broker-host:50051", - ) - mock_context = mock.MagicMock() - servicer = WorkerServicer(taskworker.worker_pool) - - servicer.PushTask(request, mock_context) - - mock_context.abort.assert_called_once_with( - grpc.StatusCode.RESOURCE_EXHAUSTED, "worker busy" - ) - @mock.patch("taskbroker_client.worker.workerchild.capture_checkin") def test_child_process_complete(mock_capture_checkin: mock.MagicMock) -> None: