Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
24 changes: 23 additions & 1 deletion cli/stack/src/flowmesh_cli_stack/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,13 @@
select_worker_images,
)

from .utils import DEFAULT_ENV_FILE, STACK_PATH_KEYS, stack_node_client
from .utils import (
DEFAULT_ENV_FILE,
STACK_PATH_KEYS,
STACK_SLUG_ENV,
stack_node_client,
stack_resource_env_overrides,
)

app = get_typer(help="Create and manage workers on the local node.")

Expand Down Expand Up @@ -74,6 +80,15 @@ def worker_up(
"Repeat --config-raw to provide multiple configs."
),
),
name_template: str | None = typer.Option(
None,
"--name-template",
help=(
"Worker name template for cpu/gpu presets. Placeholders: "
"{slug}, {kind}, {idx}, {gpu}. "
"Default: '{slug}_worker_{kind}_{idx|gpu}'."
),
),
env_file: Path = typer.Option(
DEFAULT_ENV_FILE, "--env-file", help="Env file to load defaults"
),
Expand All @@ -84,6 +99,11 @@ def worker_up(
) -> None:
"""Create and start one or more workers from presets or a custom config file."""
client = stack_node_client(env_file, base_url, token or None)
try:
slug = stack_resource_env_overrides(os.environ)[STACK_SLUG_ENV]
except ValueError as exc:
logging.error(str(exc))
raise typer.Exit(code=1)
logging.info("Creating workers...")
try:
created = create_workers(
Expand All @@ -93,6 +113,8 @@ def worker_up(
targets=targets,
config_paths=config,
config_raw=config_raw,
name_template=name_template,
slug=slug,
)
except FlowMeshError as exc:
logging.error(str(exc))
Expand Down
14 changes: 14 additions & 0 deletions docs/CLI.md
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,20 @@ flowmesh stack worker up gpu --targets 0 # 1 GPU worker pinned to GPU 0
flowmesh stack down
```

Preset workers (`cpu` / `gpu`) are named after the stack slug, so they are
isolated across stacks sharing a host the same way other stack objects are:
`flowmesh_node_worker_cpu_0` by default, `flowmesh_node_<suffix>_worker_cpu_0`
when `FLOWMESH_STACK_SUFFIX` is set. Override the naming with
`--name-template`, which accepts the placeholders `{slug}`, `{kind}`, `{idx}`,
and `{gpu}`:

```bash
flowmesh stack worker up cpu 2 --name-template '{slug}-run-{idx}' # {slug}-run-0, {slug}-run-1
```

The template must keep names unique within one `up` invocation — include
`{idx}` or `{gpu}` when creating more than one worker.

`flowmesh stack up` reads `NODE_ROLE` from the env file (default `root`). On a
root node, both local Redis services are deployed alongside the server. On a
worker node (`NODE_ROLE=worker`), Redis services are skipped — the worker
Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ all = [
[dependency-groups]
# Runtime/container dependencies.
runtime-server = [
"aiohttp>=3.13.4",
"aiohttp>=3.14.0",
"cryptography>=46.0.7",
"docker>=7.1.0",
"fastapi>=0.135.0",
Expand Down Expand Up @@ -78,7 +78,7 @@ runtime-training-gpu = [
]
runtime-agent = [
"aiofiles>=24.1.0",
"aiohttp>=3.13.4",
"aiohttp>=3.14.0",
"art>=6.5",
"arxiv>=2.2.0",
"beautifulsoup4>=4.14.3",
Expand Down
134 changes: 93 additions & 41 deletions sdk/stack/src/flowmesh_stack/workers.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,19 +65,29 @@ def create_workers(
targets: str = "all",
config_paths: list[Path] | None = None,
config_raw: list[str] | None = None,
name_template: str | None = None,
slug: str | None = None,
) -> list[tuple[str, dict[str, Any]]]:
"""Create node workers from configs or built-in cpu/gpu presets.

When ``kind`` is "gpu", if ``count`` is equal to 1, a single worker with the
specified GPU targets will be created. If ``count`` is greater than 1, one worker
will be created per GPU target, and the number of targets must match ``count``.

For the built-in cpu/gpu presets, ``slug`` prefixes the generated worker
aliases so names are scoped to a stack, and ``name_template`` overrides the
naming entirely (placeholders ``{slug}``, ``{kind}``, ``{idx}``, ``{gpu}``).
Both are ignored when ``config_paths`` / ``config_raw`` are provided, since
those payloads carry their own aliases.
"""
payloads = _payloads_for_worker_create(
kind=kind,
count=count,
targets=targets,
config_paths=config_paths,
config_raw=config_raw,
name_template=name_template,
slug=slug,
)
if not payloads:
return []
Expand Down Expand Up @@ -162,12 +172,47 @@ def detect_gpu_targets(targets: str) -> list[str]:
return [item.strip() for item in result.stdout.splitlines() if item.strip()]


_ALIAS_FIELDS = "{slug}, {kind}, {idx}, {gpu}"


class _StrictFormatDict(dict[str, Any]):
def __missing__(self, key: str) -> Any:
raise KeyError(key)


def _worker_alias(
kind: str,
idx: int,
slug: str | None,
template: str | None,
gpu: str | None = None,
) -> str:
if template is None:
template = (
("{slug}_" if slug else "")
+ f"worker_{kind}_"
+ ("{idx}" if gpu is None else "{gpu}")
)
fields: dict[str, Any] = {"slug": slug or "", "kind": kind, "idx": idx}
if gpu is not None:
fields["gpu"] = gpu
try:
return template.format_map(_StrictFormatDict(fields))
except (KeyError, IndexError, ValueError) as exc:
raise FlowMeshError(
f"invalid --name-template ({exc}); "
f"available placeholders: {_ALIAS_FIELDS}"
) from exc


def _payloads_for_worker_create(
kind: str,
count: int,
targets: str,
config_paths: list[Path] | None,
config_raw: list[str] | None,
name_template: str | None = None,
slug: str | None = None,
) -> list[tuple[str, str]]:
if count < 1:
raise FlowMeshError("Worker count must be at least 1.")
Expand All @@ -183,25 +228,17 @@ def _payloads_for_worker_create(
_extend_payloads(payloads, f"worker from raw#{idx}", raw)
return payloads

specs: list[tuple[str, dict[str, Any], str]] # alias, worker_config, label
if kind == "cpu":
return [
specs = [
(
json.dumps(
{
"provider": "docker",
"init_on_start": True,
"worker_config": {
"worker_type": "cpu",
"worker_alias": f"worker_cpu_{idx}",
},
}
),
(alias := _worker_alias("cpu", idx, slug, name_template)),
{"worker_type": "cpu", "worker_alias": alias},
"CPU worker",
)
for idx in range(count)
]

if kind == "gpu":
elif kind == "gpu":
raw_gpu_ids = detect_gpu_targets(targets)
gpu_ids: list[int] = []
for raw_gpu_id in raw_gpu_ids:
Expand All @@ -217,47 +254,62 @@ def _payloads_for_worker_create(
f"Consider setting count={len(gpu_ids)} or specifying exactly "
f"{count} GPU targets."
)
gpu_payloads = [
specs = [
(
json.dumps(
{
"provider": "docker",
"init_on_start": True,
"worker_config": {
"worker_type": "gpu",
"cuda_devices": [gpu_id],
"worker_alias": f"worker_gpu_{gpu_id}",
},
}
(
alias := _worker_alias(
"gpu", idx, slug, name_template, gpu=str(gpu_id)
)
),
{
"worker_type": "gpu",
"cuda_devices": [gpu_id],
"worker_alias": alias,
},
f"GPU worker for GPU {gpu_id}",
)
for gpu_id in gpu_ids
for idx, gpu_id in enumerate(gpu_ids)
]
else:
worker_suffix = "all" if targets == "all" else "_".join(raw_gpu_ids)
gpu_payloads = [
specs = [
(
json.dumps(
{
"provider": "docker",
"init_on_start": True,
"worker_config": {
"worker_type": "gpu",
"cuda_devices": gpu_ids,
"worker_alias": f"worker_gpu_{worker_suffix}",
},
}
(
alias := _worker_alias(
"gpu", 0, slug, name_template, gpu=worker_suffix
)
),
{
"worker_type": "gpu",
"cuda_devices": gpu_ids,
"worker_alias": alias,
},
f"GPU worker for GPUs {', '.join(raw_gpu_ids)}",
)
]
else:
raise FlowMeshError("worker up expects kind cpu|gpu or use --config")

aliases = [alias for alias, _, _ in specs]
if len(set(aliases)) != len(aliases):
raise FlowMeshError(
"--name-template produced duplicate worker names; "
"include {idx} or {gpu} to disambiguate"
)

if not gpu_payloads:
raise FlowMeshError("No GPUs detected or specified.")
return gpu_payloads

raise FlowMeshError("worker up expects kind cpu|gpu or use --config")
return [
(
json.dumps(
{
"provider": "docker",
"init_on_start": True,
"worker_config": worker_config,
}
),
label,
)
for _, worker_config, label in specs
]


def _extend_payloads(
Expand Down
2 changes: 1 addition & 1 deletion src/server/requirements.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# AUTO-GENERATED by scripts/dev/sync_requirements.py — do not edit by hand.
# Regenerate with: uv run scripts/dev/sync_requirements.py --write
aiohttp==3.13.5
aiohttp==3.14.1
cryptography==47.0.0
docker==7.1.0
fastapi==0.136.1
Expand Down
2 changes: 1 addition & 1 deletion src/worker/requirements/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
# Regenerate with: uv run scripts/dev/sync_requirements.py --write
accelerate==1.12.0
aiofiles==24.1.0
aiohttp==3.13.5
aiohttp==3.14.1
arize-phoenix-client==1.21.0
arize-phoenix-otel==0.13.1
art==6.5
Expand Down
Loading