Skip to content

feat(taskbroker): Better Child Management #732

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

Closed
george-sentry wants to merge 13 commits into from
Closed

feat(taskbroker): Better Child Management #732

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
024240e
Don't Quit Children Until Replaced
george-sentry Jun 29, 2026
78f0fce
Add Fake Sleep to Child Process
george-sentry Jun 29, 2026
8904e3c
Remove Fake Sleep
george-sentry Jun 29, 2026
4d8a100
Handle Process Start Failure
george-sentry Jun 29, 2026
01b3550
Formatting
george-sentry Jun 29, 2026
70eddc6
Readability
george-sentry Jun 29, 2026
60f805a
Try State Management Again
george-sentry Jun 29, 2026
12b2d19
Fix Child Replacement Logic
george-sentry Jun 29, 2026
dadff17
Newly Exiting Children Should Spawn Pending Children
george-sentry Jun 29, 2026
f82d5ca
Change Exiting Children Comment
george-sentry Jun 29, 2026
f7a356e
Fix Needed Count Again
george-sentry Jun 29, 2026
82dae9a
Move Counting Logic into Separate Method
george-sentry Jun 29, 2026
62ad082
Minor Edits
george-sentry Jun 29, 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
206 changes: 188 additions & 18 deletions clients/python/src/taskbroker_client/worker/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,13 @@
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from multiprocessing.context import ForkContext, ForkServerContext, SpawnContext
from multiprocessing.process import BaseProcess
from multiprocessing.synchronize import Event
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, List
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal
from uuid import UUID, uuid4

import grpc
from grpc_health.v1 import health, health_pb2, health_pb2_grpc
Expand Down Expand Up @@ -42,7 +45,7 @@
parse_rpc_secret_list,
)
from taskbroker_client.worker.push_clients import BatchPushTaskbrokerClient, PushTaskbrokerClient
from taskbroker_client.worker.workerchild import child_process
from taskbroker_client.worker.workerchild import ChildMessage, child_process

if TYPE_CHECKING:
ServerInterceptor = grpc.ServerInterceptor[Any, Any]
Expand Down Expand Up @@ -122,6 +125,16 @@ class RequeueException(Exception):
pass


ChildState = Literal["pending", "ready", "exiting"]


@dataclass
class TrackedChild:
process: BaseProcess
state: ChildState
release: Event


class PushTaskWorker:
_mp_context: ForkContext | SpawnContext | ForkServerContext

Expand Down Expand Up @@ -800,17 +813,51 @@ def __init__(
self._processed_tasks: multiprocessing.Queue[ProcessingResult] = self._mp_context.Queue(
maxsize=result_queue_maxsize
)
self._children: list[BaseProcess] = []
self._children: Dict[UUID, TrackedChild] = {}
self._children_lock = threading.Lock()
self._shutdown_event = self._mp_context.Event()
self._ready_counter = self._mp_context.Value("i", 0)
self._result_thread: threading.Thread | None = None
self._metrics_thread: threading.Thread | None = None
self._spawn_children_thread: threading.Thread | None = None

@property
def ready_count(self) -> int:
"""Number of children that have finished warming up and are consuming."""
return self._ready_counter.value
with self._children_lock:
return sum(1 for c in self._children.values() if c.state == "ready")

def _plan_child_pool_changes(self) -> tuple[int, list[TrackedChild]]:
"""
Determine how to move the child pool toward the desired concurrency.

The caller must hold `_children_lock`.

Exiting children are still consuming tasks until released, so they count
as serving capacity. Pending children are not serving yet, but they count
as replacement capacity that has already been requested.
"""
pending_children = [c for c in self._children.values() if c.state == "pending"]
ready_children = [c for c in self._children.values() if c.state == "ready"]
exiting_children = [c for c in self._children.values() if c.state == "exiting"]

desired_serving_capacity = self._concurrency
current_serving_capacity = len(ready_children) + len(exiting_children)

excess_serving_capacity = max(current_serving_capacity - desired_serving_capacity, 0)
children_to_release = exiting_children[:excess_serving_capacity]

kept_exiting_count = len(exiting_children) - len(children_to_release)

# Keep enough pending children to both reach desired concurrency and replace
# exiting children that are still serving during their grace period
target_pending_count = max(
desired_serving_capacity - len(ready_children),
kept_exiting_count,
)

children_to_spawn = max(target_pending_count - len(pending_children), 0)

return children_to_spawn, children_to_release

def send_results(self, results: list[ProcessingResult], is_draining: bool = False) -> None:
"""
Expand Down Expand Up @@ -917,16 +964,97 @@ def result_thread() -> None:
def start_spawn_children_thread(self) -> None:
def spawn_children_thread() -> None:
logger.debug("taskworker.worker.spawn_children_thread.started")

# Queue of incoming message from children
messages: multiprocessing.Queue[ChildMessage] = self._mp_context.Queue()

@sentry sentry Bot Jun 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: The messages queue is a local variable in spawn_children_thread and can be garbage-collected upon thread exit, causing child processes to fail when they try to use it.
Severity: MEDIUM

Suggested Fix

Move the messages queue to a class attribute to ensure its lifetime exceeds that of the child processes that use it. Alternatively, explicitly manage the queue's lifecycle by calling messages.close() and messages.join_thread() within a try/finally block in the shutdown flow to guarantee proper cleanup.

Prompt for AI Agent 
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.

Location: clients/python/src/taskbroker_client/worker/worker.py#L969

Potential issue: The `messages` queue is created as a local variable within the
`spawn_children_thread` function. When this thread terminates during shutdown, the
`messages` queue object goes out of scope and can be garbage collected by Python.
However, child processes, which are terminated shortly after, may still hold references
to this queue. If a child attempts to send a message via `messages.put_nowait()` after
the queue has been garbage collected, it could lead to a `BrokenPipeError` or a
deadlock, creating a race condition during the shutdown sequence.


while not self._shutdown_event.is_set():
self._children = [child for child in self._children if child.is_alive()]
if len(self._children) >= self._concurrency:
time.sleep(0.1)
continue
for i in range(self._concurrency - len(self._children)):
# Read any events that may have come in since the last loop iteration
received: List[ChildMessage] = []

while True:
try:
message = messages.get(block=False)
received.append(message)
except queue.Empty:
break

with self._children_lock:
children = list(self._children.items())

for cid, c in children:
if c.process.is_alive():
continue

c.process.join(timeout=0)
self._children.pop(cid)

logger.info(
"taskworker.child.exited",
extra={
"pid": c.process.pid,
"cid": str(cid),
"exitcode": c.process.exitcode,
"state": c.state,
"processing_pool": self._processing_pool_name,
},
)

for message in received:
child = self._children.get(message.child_id)

if child is None:
logger.warning(
"taskworker.child_message.unknown_child",
extra={
"cid": str(message.child_id),
"event": message.event,
"processing_pool": self._processing_pool_name,
},
)

continue

if message.event == "ready":
if child.state != "pending":
logger.warning(
"taskworker.child.ready_not_pending",
extra={
"cid": str(message.child_id),
"state": child.state,
"processing_pool": self._processing_pool_name,
},
)

child.state = "ready"

elif message.event == "exiting":
if child.state != "ready":
logger.warning(
"taskworker.child.exiting_not_ready",
extra={
"cid": str(message.child_id),
"state": child.state,
"processing_pool": self._processing_pool_name,
},
)

child.state = "exiting"

children_to_spawn, children_to_release = self._plan_child_pool_changes()

for child in children_to_release:
child.release.set()

for _ in range(children_to_spawn):
child_id = uuid4()
release = self._mp_context.Event()

process = self._mp_context.Process(
name=f"taskworker-child-{i}",
name=f"taskworker-child-{child_id}",
target=child_process,
args=(
child_id,
self._app_module,
self._child_tasks,
self._processed_tasks,
Expand All @@ -935,20 +1063,59 @@ def spawn_children_thread() -> None:
self._processing_pool_name,
self._process_type,
self._skip_awaiting_futures,
self._ready_counter,
messages,
release,
),
)
process.start()
self._children.append(process)

try:
process.start()
except Exception as e:
logger.exception(
"taskworker.child.spawn.failed",
extra={
"cid": str(child_id),
"error": e,
"processing_pool": self._processing_pool_name,
},
)

self._metrics.incr(
"taskworker.worker.child.spawn",
tags={
"processing_pool": self._processing_pool_name,
"result": "failure",
},
)

continue

with self._children_lock:
self._children[child_id] = TrackedChild(
process=process,
state="pending",
release=release,
)

logger.info(
"taskworker.spawn_child",
extra={"pid": process.pid, "processing_pool": self._processing_pool_name},
extra={
"pid": process.pid,
"cid": str(child_id),
"processing_pool": self._processing_pool_name,
},
)

self._metrics.incr(
"taskworker.worker.spawn_child",
tags={"processing_pool": self._processing_pool_name},
tags={
"processing_pool": self._processing_pool_name,
"result": "success",
},
)

time.sleep(0.1)

self._spawn_children_thread = threading.Thread(
name="spawn-children", target=spawn_children_thread, daemon=True
)
Expand Down Expand Up @@ -1005,9 +1172,12 @@ def shutdown(self) -> None:
self._spawn_children_thread.join()

logger.info("taskworker.worker.shutdown.children")
for child in self._children:
with self._children_lock:
children = [tracked_child.process for tracked_child in self._children.values()]

for child in children:
child.terminate()
for child in self._children:
for child in children:
child.join(WORKER_CHILD_JOIN_TIMEOUT_SEC)
if child.is_alive():
child.kill()
Expand Down
56 changes: 36 additions & 20 deletions clients/python/src/taskbroker_client/worker/workerchild.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import contextlib
import logging
import multiprocessing
import queue
import signal
import threading
Expand All @@ -11,10 +12,8 @@
from functools import partial
from multiprocessing.synchronize import Event
from types import FrameType
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
from multiprocessing.sharedctypes import Synchronized
from typing import Any, Literal
from uuid import UUID

# XXX: Don't import any modules that will import django here, do those within child_process
import msgpack
Expand Down Expand Up @@ -166,7 +165,14 @@ def _log_task_retry_exhausted(
)


@dataclass(frozen=True)
class ChildMessage:
child_id: UUID
event: Literal["ready"] | Literal["exiting"]


def child_process(
child_id: UUID,
app_module: str,
child_tasks: queue.Queue[InflightTaskActivation],
processed_tasks: queue.Queue[ProcessingResult],
Expand All @@ -175,7 +181,8 @@ def child_process(
processing_pool_name: str,
process_type: str,
skip_awaiting_futures: bool,
ready_counter: "Synchronized[int] | None" = None,
messages: multiprocessing.Queue[ChildMessage],
parent_release: Event,
) -> None:
"""
The entrypoint for spawned worker children.
Expand Down Expand Up @@ -351,18 +358,30 @@ def check_task_future_completion(
)
_future_completion_thread.start()

waiting_for_parent_release = False

while not shutdown_event.is_set() and not local_shutdown.is_set():
if max_task_count and processed_task_count >= max_task_count:
metrics.incr(
"taskworker.worker.max_task_count_reached",
tags={"count": processed_task_count, "processing_pool": processing_pool_name},
)
logger.info(
"taskworker.max_task_count_reached", extra={"count": processed_task_count}
)
# Still set the shutdown signal to trigger shutdown of the future checker thread
local_shutdown.set()
break
if not waiting_for_parent_release:
metrics.incr(
"taskworker.worker.max_task_count_reached",
tags={
"count": processed_task_count,
"processing_pool": processing_pool_name,
},
)

logger.info(
"taskworker.max_task_count_reached", extra={"count": processed_task_count}
)

# Tell the parent this child can be replaced.
messages.put_nowait(ChildMessage(child_id, "exiting"))
waiting_for_parent_release = True

if parent_release.is_set():
local_shutdown.set()
break

try:
inflight = child_tasks.get(timeout=1.0)
Expand Down Expand Up @@ -806,11 +825,8 @@ def _task_execution_complete(
futures_start_time,
)

# Signal that this child has finished warmup and ready to consume tasks. The parent uses this
# to gate the gRPC SERVING health signal. Monotonic by design
if ready_counter is not None:
with ready_counter.get_lock():
ready_counter.value += 1
# Tell the parent that this child has warmed up and is ready to consume tasks
messages.put_nowait(ChildMessage(child_id, "ready"))

# Run the worker loop
run_worker(
Expand Down
Loading
Loading