forked from goodboy/tractor
1
0
Fork 0
tractor/tractor/_clustering.py

54 lines
1.3 KiB
Python
Raw Normal View History

'''
Actor cluster helpers.
'''
2021-10-16 15:33:39 +00:00
from __future__ import annotations
from contextlib import asynccontextmanager as acm
from multiprocessing import cpu_count
from typing import AsyncGenerator, Optional
import trio
import tractor
@acm
async def open_actor_cluster(
modules: list[str],
count: int = cpu_count(),
names: Optional[list[str]] = None,
start_method: Optional[str] = None,
hard_kill: bool = False,
2021-10-04 20:01:09 +00:00
) -> AsyncGenerator[
list[str],
dict[str, tractor.Portal]
]:
portals: dict[str, tractor.Portal] = {}
if not names:
names = [f'worker_{i}' for i in range(count)]
if not len(names) == count:
raise ValueError(
'Number of names is {len(names)} but count it {count}')
async with tractor.open_nursery(start_method=start_method) as an:
2021-10-04 20:01:09 +00:00
async with trio.open_nursery() as n:
uid = tractor.current_actor().uid
async def _start(name: str) -> None:
name = f'{name}.{uid}'
portals[name] = await an.start_actor(
enable_modules=modules,
name=name,
)
for name in names:
n.start_soon(_start, name)
assert len(portals) == count
yield portals
2021-10-16 15:50:55 +00:00
await an.cancel(hard_kill=hard_kill)