feat: support automatic worker sizing

docs/running-pipelines/cloud-launcher.md

@@ -17,13 +17,20 @@ shards, run transforms, write outputs, and report logs and metrics.
 ```python
 pipeline.launch_cloud(
     name="aloha-trim",
-    num_workers=8,
+    num_workers="auto",
     cpus_per_worker=4,
     mem_mb_per_worker=8192,
     secrets={"HF_TOKEN": None},
 )
 ```
 
+Set `num_workers="auto"` to request one worker for every shard in each stage.
+The cloud registration runtime discovers shards after secrets and environment
+variables are mounted, then Macrodata Cloud applies its normal worker and GPU
+limits to the resulting count. The submitting machine does not need access to
+private inputs. An empty stage starts no worker containers and completes after
+shard registration. Pass a positive integer when you want a fixed worker count.
+
 ## What gets submitted
 
 A cloud submission includes:

docs/running-pipelines/local-launcher.md

@@ -11,10 +11,14 @@ to the Macrodata Cloud.
 ```python
 pipeline.launch_local(
     name="debug-local",
-    num_workers=2,
+    num_workers="auto",
 )
 ```
 
+With `num_workers="auto"`, Refiner starts one local worker for each shard in
+the current stage. An empty stage starts no workers. Pass a positive
+integer to cap execution at a fixed number of worker processes.
+
 The local launcher is useful for:
 
 - verifying a writer on a small dataset

src/refiner/launchers/base.py

@@ -12,6 +12,7 @@
     StageComputeRequirements,
     compile_planned_stages,
     plan_pipeline_stages,
+    WorkerCount,
 )
 from refiner.pipeline.resources import GPU
 
@@ -25,16 +26,18 @@ def __init__(
         *,
         pipeline: RefinerPipeline,
         name: str,
-        num_workers: int = 1,
+        num_workers: WorkerCount = 1,
         cpus_per_worker: int | None = None,
         gpu: GPU | None = None,
     ):
         if not name.strip():
             raise ValueError("name must be non-empty")
         self.pipeline = pipeline
         self.name = name
-        if num_workers <= 0:
-            raise ValueError("num_workers must be > 0")
+        if num_workers != "auto" and (
+            not isinstance(num_workers, int) or num_workers <= 0
+        ):
+            raise ValueError("num_workers must be > 0 or 'auto'")
         self.num_workers = num_workers
         if cpus_per_worker is not None and cpus_per_worker <= 0:
             raise ValueError("cpus_per_worker must be > 0")
@@ -47,10 +50,7 @@ def _build_local_job_id(name: str) -> str:
         return f"{slug}-{int(time.time())}-{uuid4().hex[:8]}"
 
     def _planned_stages(self) -> list[PlannedStage]:
-        requested_workers = getattr(self, "num_workers", None)
-        default_num_workers = (
-            requested_workers if isinstance(requested_workers, int) else 1
-        )
+        default_num_workers = getattr(self, "num_workers", 1)
         return plan_pipeline_stages(
             self.pipeline,
             default_num_workers=default_num_workers,

src/refiner/launchers/cloud.py

@@ -5,7 +5,7 @@
 import re
 import sys
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Literal, cast
 
 from refiner.cli.run.modes import (
     CloudAttachContext,
@@ -97,7 +97,8 @@ class CloudLauncher(BaseLauncher):
     Args:
         pipeline: Pipeline to execute.
         name: Human-readable run name.
-        num_workers: Requested logical worker count for cloud execution.
+        num_workers: Requested logical worker count for cloud execution, or
+            ``"auto"`` to launch one worker per stage shard.
         cpus_per_worker: Optional requested CPU cores per worker.
         mem_mb_per_worker: Optional requested memory in MB per worker for cloud scheduling.
         gpu: Optional GPU runtime request for cloud scheduling.
@@ -117,7 +118,7 @@ def __init__(
         *,
         pipeline: "RefinerPipeline",
         name: str,
-        num_workers: int = 1,
+        num_workers: int | Literal["auto"] = 1,
         cpus_per_worker: int | None = None,
         mem_mb_per_worker: int | None = None,
         gpu: GPU | None = None,

src/refiner/launchers/local.py

@@ -4,8 +4,9 @@
 import subprocess
 import sys
 import threading
+from dataclasses import replace
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 from uuid import uuid4
 
 import cloudpickle
@@ -52,7 +53,7 @@ def __init__(
         *,
         pipeline: RefinerPipeline,
         name: str,
-        num_workers: int = 1,
+        num_workers: int | Literal["auto"] = 1,
         rundir: str | None = None,
         gpu: GPU | None = None,
     ):
@@ -69,6 +70,24 @@ def __init__(
         self.job_tracking_url: str | None = None
         self._total_stages = 1
 
+    def _resolved_stages(
+        self,
+        stages: list[PlannedStage] | None = None,
+    ) -> list[PlannedStage]:
+        resolved = super()._resolved_stages(stages)
+        return [
+            replace(
+                stage,
+                compute=replace(
+                    stage.compute,
+                    num_workers=len(stage.pipeline.list_shards()),
+                ),
+            )
+            if stage.compute.num_workers == "auto"
+            else stage
+            for stage in resolved
+        ]
+
     def _collect_worker_results(
         self,
         *,
@@ -283,23 +302,12 @@ def _launch_stage(
     ) -> LaunchStats:
         # Resolve worker capacity and remaining stage shards.
         stage_workers = stage.compute.num_workers
+        if not isinstance(stage_workers, int):
+            raise RuntimeError("local stage worker count was not resolved")
         if self.job_id is None or self.rundir is None:
             raise RuntimeError(
                 "local launcher must be initialized in launch() before running stages"
             )
-        available_cpus = len(available_cpu_ids())
-        if stage_workers > available_cpus:
-            logger.warning(
-                f"stage {stage.index} requested {stage_workers} workers, but only {available_cpus} CPUs are available on this machine."
-            )
-        gpu_sets = (
-            build_gpu_sets(
-                num_workers=stage_workers,
-                gpu_count_per_worker=stage.compute.gpu.count,
-            )
-            if stage.compute.gpu is not None
-            else [[] for _ in range(stage_workers)]
-        )
         completed_shard_ids = {
             row.shard_id
             for row in read_finalized_workers(
@@ -325,6 +333,23 @@ def _launch_stage(
                 output_rows=0,
             )
 
+        if self.num_workers == "auto" and stage.compute.inherit_launcher_resources:
+            stage_workers = len(shards)
+        available_cpus = len(available_cpu_ids())
+        if stage_workers > available_cpus:
+            logger.warning(
+                f"stage {stage.index} requested {stage_workers} workers, but only {available_cpus} CPUs are available on this machine."
+            )
+
+        gpu_sets = (
+            build_gpu_sets(
+                num_workers=stage_workers,
+                gpu_count_per_worker=stage.compute.gpu.count,
+            )
+            if stage.compute.gpu is not None
+            else [[] for _ in range(stage_workers)]
+        )
+
         # Persist the stage payload and worker assignments under the rundir.
         stage_run_dir = Path(self.rundir) / f"stage-{stage.index}"
         stage_run_dir.mkdir(parents=True, exist_ok=True)
@@ -381,7 +406,7 @@ def launch(self) -> LaunchStats:
         if attach_mode_override() == "detach":
             raise SystemExit("--detach is only supported for cloud launches.")
         available_cpus = len(available_cpu_ids())
-        if self.num_workers > available_cpus:
+        if isinstance(self.num_workers, int) and self.num_workers > available_cpus:
             logger.warning(
                 f"launch requested {self.num_workers} workers, but only {available_cpus} CPUs are available on this machine."
             )

src/refiner/pipeline/pipeline.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Iterator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Callable, cast
+from typing import TYPE_CHECKING, Any, Callable, Literal, cast
 
 from fsspec import AbstractFileSystem
 
@@ -676,15 +676,16 @@ def launch_local(
         self,
         *,
         name: str,
-        num_workers: int = 1,
+        num_workers: int | Literal["auto"] = 1,
         rundir: str | None = None,
         gpu: GPU | None = None,
     ) -> "LaunchStats":
         """Launch the pipeline locally.
 
         Args:
             name: Human-readable run name.
-            num_workers: Number of local worker processes.
+            num_workers: Number of local worker processes, or ``"auto"`` to
+                launch one worker per stage shard.
             rundir: Optional explicit local run directory. Reuse it to resume a prior local run.
             gpu: Optional GPU devices exposed per worker. `cuda_version` is accepted
                 for API consistency but ignored by local launch.
@@ -704,7 +705,7 @@ def launch_cloud(
         self,
         *,
         name: str,
-        num_workers: int = 1,
+        num_workers: int | Literal["auto"] = 1,
         cpus_per_worker: int | None = None,
         mem_mb_per_worker: int | None = None,
         gpu: GPU | None = None,
@@ -720,7 +721,8 @@ def launch_cloud(
 
         Args:
             name: Human-readable run name.
-            num_workers: Requested logical worker count.
+            num_workers: Requested logical worker count, or ``"auto"`` to
+                launch one worker per stage shard.
             cpus_per_worker: Optional requested CPU cores per worker.
             mem_mb_per_worker: Optional requested memory in MB per worker for cloud scheduling.
             gpu: Optional structured GPU request.

src/refiner/pipeline/planning.py

@@ -5,7 +5,7 @@
 import textwrap
 from dataclasses import dataclass
 from types import CodeType
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias
 
 from refiner.pipeline.steps import (
     CastStep,
@@ -37,10 +37,12 @@
 
 _REFINER_BUILTIN_CALL_ATTR = "__refiner_builtin_call__"
 
+WorkerCount: TypeAlias = int | Literal["auto"]
+
 
 @dataclass(frozen=True, slots=True)
 class StageComputeRequirements:
-    num_workers: int
+    num_workers: WorkerCount
     cpus_per_worker: int | None = None
     memory_mb_per_worker: int | None = None
     gpu: GPU | None = None
@@ -462,14 +464,16 @@ def _unique_name(base: str) -> str:
 
 
 def plan_pipeline_stages(
-    pipeline: "RefinerPipeline", *, default_num_workers: int
+    pipeline: "RefinerPipeline", *, default_num_workers: WorkerCount
 ) -> list[PlannedStage]:
     """Return the ordered execution stages for a pipeline.
 
     This is currently a placeholder splitter that yields a single stage. Future
     multi-stage planning logic should live here.
     """
-    if default_num_workers <= 0:
+    if default_num_workers != "auto" and (
+        not isinstance(default_num_workers, int) or default_num_workers <= 0
+    ):
         raise ValueError("default_num_workers must be > 0")
 
     from refiner.pipeline.pipeline import RefinerPipeline

src/refiner/platform/client/models.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
-from typing import Any
+from typing import Any, Literal
 
 import msgspec
 
@@ -123,7 +123,7 @@ class StageLifecycleResponse(msgspec.Struct, frozen=True):
 
 @dataclass(frozen=True, slots=True)
 class CloudRuntimeConfig:
-    num_workers: int
+    num_workers: int | Literal["auto"]
     cpus_per_worker: int | None = None
     mem_mb_per_worker: int | None = None
     gpu: GPU | None = None

tests/launchers/test_base_launcher.py

@@ -2,6 +2,8 @@
 
 from typing import cast
 
+import pytest
+
 from refiner.job_urls import build_job_tracking_url
 from refiner.platform.client import MacrodataClient
 from refiner.pipeline import RefinerPipeline
@@ -29,6 +31,16 @@ def _stage_compute_requirements(
         )
 
 
+@pytest.mark.parametrize("num_workers", [0, -1, "AUTO"])
+def test_launcher_rejects_invalid_worker_count(num_workers) -> None:
+    with pytest.raises(ValueError, match="num_workers must be > 0 or 'auto'"):
+        _DummyLauncher(
+            pipeline=cast(RefinerPipeline, object()),
+            name="unit-test",
+            num_workers=num_workers,
+        )
+
+
 def test_job_tracking_url_sanitizes_terminal_control_characters() -> None:
     client = MacrodataClient(api_key="md_test", base_url="https://app.\x9bexample.com")
 

tests/launchers/test_cloud_launcher.py

@@ -258,6 +258,29 @@ def test_pipeline_launch_cloud_submits_compiled_plan(monkeypatch) -> None:
     assert captured["events"] == ["upload-urls", "upload", "complete", "submit"]
 
 
+def test_pipeline_launch_cloud_preserves_auto_workers_without_listing_shards(
+    monkeypatch,
+) -> None:
+    captured = _stub_cloud_submit(monkeypatch)
+    pipeline = mdr.from_items([1, 2, 3], items_per_shard=2)
+    monkeypatch.setattr(
+        type(pipeline.source),
+        "list_shards",
+        lambda _: pytest.fail("cloud submission must not list shards locally"),
+    )
+
+    pipeline.launch_cloud(
+        name="auto workers",
+        num_workers="auto",
+    )
+
+    request = cast(CloudRunCreateRequest, captured["submit_request"])
+    stage = request.stage_payloads[0]
+    assert request.plan["stages"][0]["requested_num_workers"] == "auto"
+    assert stage.runtime.num_workers == "auto"
+    assert "num_shards" not in stage.to_dict()
+
+
 def test_pipeline_launch_cloud_embeds_runtime_services(monkeypatch) -> None:
     captured = _stub_cloud_submit(monkeypatch)
     monkeypatch.setattr(