← back to Exo
add lazy task group (#1569)
61d2a2b6cffbb53e51f0ec90a348c57c8f8f7f44 · 2026-02-23 16:05:59 +0000 · Evan Quiney
in a few places we instantiate task groups early, in others we have an
optional task group.
this standardizes these patterns into a lazy task group, which queues
tasks until it is entered, at which point it enters its inner task group
and starts all its queued tasks. it should be noted that no queued tasks
will be started until the group itself is started.
Files touched
M src/exo/download/coordinator.pyM src/exo/main.pyM src/exo/master/api.pyM src/exo/master/main.pyM src/exo/routing/router.pyM src/exo/shared/election.pyM src/exo/utils/info_gatherer/info_gatherer.pyA src/exo/utils/task_group.pyM src/exo/worker/main.py
Diff
commit 61d2a2b6cffbb53e51f0ec90a348c57c8f8f7f44
Author: Evan Quiney <evanev7@gmail.com>
Date: Mon Feb 23 16:05:59 2026 +0000
add lazy task group (#1569)
in a few places we instantiate task groups early, in others we have an
optional task group.
this standardizes these patterns into a lazy task group, which queues
tasks until it is entered, at which point it enters its inner task group
and starts all its queued tasks. it should be noted that no queued tasks
will be started until the group itself is started.
---
src/exo/download/coordinator.py | 6 +--
src/exo/main.py | 8 ++--
src/exo/master/api.py | 9 ++--
src/exo/master/main.py | 6 +--
src/exo/routing/router.py | 13 ++----
src/exo/shared/election.py | 17 ++------
src/exo/utils/info_gatherer/info_gatherer.py | 8 ++--
src/exo/utils/task_group.py | 65 ++++++++++++++++++++++++++++
src/exo/worker/main.py | 8 ++--
9 files changed, 95 insertions(+), 45 deletions(-)
diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py
index 5b6026be..f1988135 100644
--- a/src/exo/download/coordinator.py
+++ b/src/exo/download/coordinator.py
@@ -5,7 +5,6 @@ from random import random
import anyio
from anyio import current_time
-from anyio.abc import TaskGroup
from loguru import logger
from exo.download.download_utils import (
@@ -41,6 +40,7 @@ from exo.shared.types.worker.downloads import (
)
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.task_group import TaskGroup
@dataclass
@@ -66,7 +66,7 @@ class DownloadCoordinator:
# Internal event channel for forwarding (initialized in __post_init__)
event_sender: Sender[Event] = field(init=False)
event_receiver: Receiver[Event] = field(init=False)
- _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
+ _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
# Per-model throttle for download progress events
_last_progress_time: dict[ModelId, float] = field(default_factory=dict)
@@ -167,7 +167,7 @@ class DownloadCoordinator:
self._tg.start_soon(self._emit_existing_download_progress)
def shutdown(self) -> None:
- self._tg.cancel_scope.cancel()
+ self._tg.cancel_tasks()
# directly copied from worker
async def _resend_out_for_delivery(self) -> None:
diff --git a/src/exo/main.py b/src/exo/main.py
index 60729cc6..8a56ab19 100644
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -7,7 +7,6 @@ from dataclasses import dataclass, field
from typing import Self
import anyio
-from anyio.abc import TaskGroup
from loguru import logger
from pydantic import PositiveInt
@@ -23,6 +22,7 @@ from exo.shared.logging import logger_cleanup, logger_setup
from exo.shared.types.common import NodeId, SessionId
from exo.utils.channels import Receiver, channel
from exo.utils.pydantic_ext import CamelCaseModel
+from exo.utils.task_group import TaskGroup
from exo.worker.main import Worker
@@ -38,7 +38,7 @@ class Node:
node_id: NodeId
offline: bool
- _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
+ _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
@classmethod
async def create(cls, args: "Args") -> Self:
@@ -149,11 +149,11 @@ class Node:
def shutdown(self):
# if this is our second call to shutdown, just sys.exit
- if self._tg.cancel_scope.cancel_called:
+ if self._tg.cancel_called():
import sys
sys.exit(1)
- self._tg.cancel_scope.cancel()
+ self._tg.cancel_tasks()
async def _elect_loop(self):
with self.election_result_receiver as results:
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index 5ecc1575..85d964b7 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -11,8 +11,7 @@ from typing import Annotated, Literal, cast
from uuid import uuid4
import anyio
-from anyio import BrokenResourceError, create_task_group
-from anyio.abc import TaskGroup
+from anyio import BrokenResourceError
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -174,6 +173,7 @@ from exo.shared.types.worker.shards import Sharding
from exo.utils.banner import print_startup_banner
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import OrderedBuffer
+from exo.utils.task_group import TaskGroup
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
@@ -252,7 +252,7 @@ class API:
CommandId, Sender[ImageChunk | ErrorChunk]
] = {}
self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
- self._tg: TaskGroup | None = None
+ self._tg: TaskGroup = TaskGroup()
def reset(self, new_session_id: SessionId, result_clock: int):
logger.info("Resetting API State")
@@ -1591,8 +1591,7 @@ class API:
shutdown_ev = anyio.Event()
try:
- async with create_task_group() as tg:
- self._tg = tg
+ async with self._tg as tg:
logger.info("Starting API")
tg.start_soon(self._apply_state)
tg.start_soon(self._pause_on_new_election)
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index 196fae99..b798a96d 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -1,7 +1,6 @@
from datetime import datetime, timedelta, timezone
import anyio
-from anyio.abc import TaskGroup
from loguru import logger
from exo.master.event_log import DiskEventLog
@@ -63,6 +62,7 @@ from exo.shared.types.tasks import (
from exo.shared.types.worker.instances import InstanceId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import MultiSourceBuffer
+from exo.utils.task_group import TaskGroup
class Master:
@@ -77,7 +77,7 @@ class Master:
download_command_sender: Sender[ForwarderDownloadCommand],
):
self.state = State()
- self._tg: TaskGroup = anyio.create_task_group()
+ self._tg: TaskGroup = TaskGroup()
self.node_id = node_id
self.session_id = session_id
self.command_task_mapping: dict[CommandId, TaskId] = {}
@@ -116,7 +116,7 @@ class Master:
async def shutdown(self):
logger.info("Stopping Master")
- self._tg.cancel_scope.cancel()
+ self._tg.cancel_tasks()
async def _command_processor(self) -> None:
with self.command_receiver as commands:
diff --git a/src/exo/routing/router.py b/src/exo/routing/router.py
index d71275b7..cebfee82 100644
--- a/src/exo/routing/router.py
+++ b/src/exo/routing/router.py
@@ -8,11 +8,9 @@ from typing import cast
from anyio import (
BrokenResourceError,
ClosedResourceError,
- create_task_group,
move_on_after,
sleep_forever,
)
-from anyio.abc import TaskGroup
from exo_pyo3_bindings import (
AllQueuesFullError,
Keypair,
@@ -25,6 +23,7 @@ from loguru import logger
from exo.shared.constants import EXO_NODE_ID_KEYPAIR
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.pydantic_ext import CamelCaseModel
+from exo.utils.task_group import TaskGroup
from .connection_message import ConnectionMessage
from .topics import CONNECTION_MESSAGES, PublishPolicy, TypedTopic
@@ -111,10 +110,9 @@ class Router:
self._net: NetworkingHandle = handle
self._tmp_networking_sender: Sender[tuple[str, bytes]] | None = send
self._id_count = count()
- self._tg: TaskGroup | None = None
+ self._tg: TaskGroup = TaskGroup()
async def register_topic[T: CamelCaseModel](self, topic: TypedTopic[T]):
- assert self._tg is None, "Attempted to register topic after setup time"
send = self._tmp_networking_sender
if send:
self._tmp_networking_sender = None
@@ -148,8 +146,7 @@ class Router:
async def run(self):
logger.debug("Starting Router")
try:
- async with create_task_group() as tg:
- self._tg = tg
+ async with self._tg as tg:
for topic in self.topic_routers:
router = self.topic_routers[topic]
tg.start_soon(router.run)
@@ -165,9 +162,7 @@ class Router:
async def shutdown(self):
logger.debug("Shutting down Router")
- if not self._tg:
- return
- self._tg.cancel_scope.cancel()
+ self._tg.cancel_tasks()
async def _networking_subscribe(self, topic: str):
await self._net.gossipsub_subscribe(topic)
diff --git a/src/exo/shared/election.py b/src/exo/shared/election.py
index 9be7054d..6f6e1f8f 100644
--- a/src/exo/shared/election.py
+++ b/src/exo/shared/election.py
@@ -4,10 +4,8 @@ import anyio
from anyio import (
CancelScope,
Event,
- create_task_group,
get_cancelled_exc_class,
)
-from anyio.abc import TaskGroup
from loguru import logger
from exo.routing.connection_message import ConnectionMessage
@@ -15,6 +13,7 @@ from exo.shared.types.commands import ForwarderCommand
from exo.shared.types.common import NodeId, SessionId
from exo.utils.channels import Receiver, Sender
from exo.utils.pydantic_ext import CamelCaseModel
+from exo.utils.task_group import TaskGroup
DEFAULT_ELECTION_TIMEOUT = 3.0
@@ -82,13 +81,12 @@ class Election:
self._candidates: list[ElectionMessage] = []
self._campaign_cancel_scope: CancelScope | None = None
self._campaign_done: Event | None = None
- self._tg: TaskGroup | None = None
+ self._tg = TaskGroup()
async def run(self):
logger.info("Starting Election")
try:
- async with create_task_group() as tg:
- self._tg = tg
+ async with self._tg as tg:
tg.start_soon(self._election_receiver)
tg.start_soon(self._connection_receiver)
tg.start_soon(self._command_counter)
@@ -124,12 +122,7 @@ class Election:
)
async def shutdown(self) -> None:
- if not self._tg:
- logger.warning(
- "Attempted to shutdown election service that was not running"
- )
- return
- self._tg.cancel_scope.cancel()
+ self._tg.cancel_tasks()
async def _election_receiver(self) -> None:
with self._em_receiver as election_messages:
@@ -143,7 +136,6 @@ class Election:
if message.clock > self.clock:
self.clock = message.clock
logger.debug(f"New clock: {self.clock}")
- assert self._tg is not None
logger.debug("Starting new campaign")
candidates: list[ElectionMessage] = [message]
logger.debug(f"Candidates: {candidates}")
@@ -178,7 +170,6 @@ class Election:
# These messages are strictly peer to peer
self.clock += 1
logger.debug(f"New clock: {self.clock}")
- assert self._tg is not None
candidates: list[ElectionMessage] = []
self._candidates = candidates
logger.debug("Starting new campaign")
diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py
index a43e54f1..9cb04c06 100644
--- a/src/exo/utils/info_gatherer/info_gatherer.py
+++ b/src/exo/utils/info_gatherer/info_gatherer.py
@@ -8,8 +8,7 @@ from subprocess import CalledProcessError
from typing import Self, cast
import anyio
-from anyio import create_task_group, fail_after, open_process, to_thread
-from anyio.abc import TaskGroup
+from anyio import fail_after, open_process, to_thread
from anyio.streams.buffered import BufferedByteReceiveStream
from anyio.streams.text import TextReceiveStream
from loguru import logger
@@ -30,6 +29,7 @@ from exo.shared.types.thunderbolt import (
)
from exo.utils.channels import Sender
from exo.utils.pydantic_ext import TaggedModel
+from exo.utils.task_group import TaskGroup
from .macmon import MacmonMetrics
from .system_info import (
@@ -381,7 +381,7 @@ class InfoGatherer:
static_info_poll_interval: float | None = 60
rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
disk_poll_interval: float | None = 30
- _tg: TaskGroup = field(init=False, default_factory=create_task_group)
+ _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
async def run(self):
async with self._tg as tg:
@@ -408,7 +408,7 @@ class InfoGatherer:
await self.info_sender.send(nc)
def shutdown(self):
- self._tg.cancel_scope.cancel()
+ self._tg.cancel_tasks()
async def _monitor_static_info(self):
if self.static_info_poll_interval is None:
diff --git a/src/exo/utils/task_group.py b/src/exo/utils/task_group.py
new file mode 100644
index 00000000..41e7f0d6
--- /dev/null
+++ b/src/exo/utils/task_group.py
@@ -0,0 +1,65 @@
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from types import TracebackType
+from typing import Any, Unpack
+
+from anyio import create_task_group
+from anyio.abc import TaskGroup as TaskGroupABC
+
+
+@dataclass
+class TaskGroup:
+ _tg: TaskGroupABC | None = field(default=None, init=False)
+ _queued: list[tuple[Any, Any, Any]] | None = field(default_factory=list, init=False)
+
+ def is_running(self) -> bool:
+ return self._tg is not None
+
+ def cancel_tasks(self):
+ assert self._tg
+ self._tg.cancel_scope.cancel()
+
+ def cancel_called(self) -> bool:
+ assert self._tg
+ return self._tg.cancel_scope.cancel_called
+
+ def start_soon[*T](
+ self,
+ func: Callable[[Unpack[T]], Awaitable[Any]],
+ *args: Unpack[T],
+ name: object = None,
+ ) -> None:
+ assert self._tg is not None
+ assert self._queued is None
+ self._tg.start_soon(func, *args, name=name)
+
+ def queue[*T](
+ self,
+ func: Callable[[Unpack[T]], Awaitable[Any]],
+ *args: Unpack[T],
+ name: object = None,
+ ) -> None:
+ assert self._tg is None
+ assert self._queued is not None
+ self._queued.append((func, args, name))
+
+ async def __aenter__(self) -> TaskGroupABC:
+ assert self._tg is None
+ assert self._queued is not None
+ self._tg = create_task_group()
+ r = await self._tg.__aenter__()
+ for func, args, name in self._queued: # pyright: ignore[reportAny]
+ self._tg.start_soon(func, *args, name=name) # pyright: ignore[reportAny]
+ self._queued = None
+ return r
+
+ async def __aexit__(
+ self,
+ exc_type: type[BaseException] | None,
+ exc_val: BaseException | None,
+ exc_tb: TracebackType | None,
+ ) -> bool:
+ """Exit the task group context waiting for all tasks to finish."""
+ assert self._tg is not None, "aenter sets self.lazy, so it exists when we aexit"
+ assert self._queued is None
+ return await self._tg.__aexit__(exc_type, exc_val, exc_tb)
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 3d3d30a1..1ebcc933 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -3,8 +3,7 @@ from datetime import datetime, timezone
from random import random
import anyio
-from anyio import CancelScope, create_task_group, fail_after
-from anyio.abc import TaskGroup
+from anyio import CancelScope, fail_after
from loguru import logger
from exo.download.download_utils import resolve_model_in_path
@@ -51,6 +50,7 @@ from exo.utils.event_buffer import OrderedBuffer
from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
from exo.utils.info_gatherer.net_profile import check_reachable
from exo.utils.keyed_backoff import KeyedBackoff
+from exo.utils.task_group import TaskGroup
from exo.worker.plan import plan
from exo.worker.runner.runner_supervisor import RunnerSupervisor
@@ -80,7 +80,7 @@ class Worker:
self.state: State = State()
self.runners: dict[RunnerId, RunnerSupervisor] = {}
- self._tg: TaskGroup = create_task_group()
+ self._tg: TaskGroup = TaskGroup()
self._nack_cancel_scope: CancelScope | None = None
self._nack_attempts: int = 0
@@ -317,7 +317,7 @@ class Worker:
await self._start_runner_task(task)
def shutdown(self):
- self._tg.cancel_scope.cancel()
+ self._tg.cancel_tasks()
async def _start_runner_task(self, task: Task):
if (instance := self.state.instances.get(task.instance_id)) is not None:
← 0ff99a2c fix isinstance for qwen3Moe (#1595)
·
back to Exo
·
runner process checks (#1592) 22610147 →