From 96bae923bdc912ad2b016109d85449c09f907e74 Mon Sep 17 00:00:00 2001 From: nichinichisou Date: Thu, 15 Oct 2026 03:06:08 +0800 Subject: [PATCH] feat(orchestrate): run the tasks of a batch as one job in one worker A stage can now name a batch per task (Stage.batch). The tasks of a stage that share a batch form one job: one worker runs them one after another, so what they share (such as a loaded bundle) is loaded once, and calls Stage.end_batch after the last of them to drop it. A task without a batch is a job alone. Scheduling, memory admission and the plan's estimates work on jobs: a job's CPU is the sum of its tasks', its peak the largest of theirs. When a worker dies inside a job, the task it was running fails and the job's remaining tasks run in another worker. The run log names each task's job. Co-Authored-By: Claude Opus 5.5 --- docs/assets.md | 17 +- src/nnnotes/pipeline/orchestrate.py | 278 ++++++++++++++++++---------- src/nnnotes/pipeline/plan.py | 13 +- src/nnnotes/pipeline/stages.py | 8 + tests/pipeline/test_orchestrate.py | 127 +++++++++++-- tests/pipeline/test_plan.py | 10 +- 6 files changed, 328 insertions(+), 125 deletions(-) diff --git a/docs/assets.md b/docs/assets.md index 7060085..c94bb36 100644 --- a/docs/assets.md +++ b/docs/assets.md @@ -58,12 +58,15 @@ stages that need those results (atlas sprites, views), writes the layouts and th (CRI audio and movies) not downloaded yet are fetched while the bundle stages run. The census also covers the bundles the selected ones depend on, so script classes and atlas textures in other bundles resolve. -Tasks run in spawned worker processes that stay alive between tasks (no more of them start than a stage has tasks to -run). The longest tasks (by estimated CPU time) start first; a task starts only while the estimated peak memory of -the running tasks and its own fits the budget, and a task estimated above the budget runs alone. A worker is -replaced after 500 tasks or when its resident memory after a task is above 3 GiB, or above 40 % of its part of the -memory budget (the budget divided by the workers) when that is less, but not below 256 MiB: a worker keeps much of -its largest task's memory, which the budget does not count (`--explain` prints the threshold). The estimates come +Tasks run in spawned worker processes that stay alive between tasks. The tasks of a stage that read the same +loaded data (a stage's batch, such as one bundle) form one job: one worker runs them one after another and loads the +data once. Every other task is a job alone. No more workers start than a stage has jobs to run. The longest jobs (by +estimated CPU time, the sum over their tasks) start first; a job starts only while the estimated peak memory of the +running jobs and its own (the largest of its tasks') fits the budget, and a job estimated above the budget runs +alone. A worker is replaced after a job once it has run 500 tasks, or when its resident memory at the end of a job +is above 3 GiB, or above 40 % of its part of the memory budget (the budget divided by the workers) when that is +less, but not below 256 MiB: a worker keeps much of its largest task's memory, which the budget does not count +(`--explain` prints the threshold). The estimates come from each stage's cost model, scaled per stage by the costs measured in earlier runs of the depot (`/costs/calibration.json`, updated after every export; a task's measured peak counts from its worker's memory at its start). None of this changes an output: a run with one worker and a run with many write the same depot @@ -71,7 +74,7 @@ and the same files. A failed object is reported with its task's result and stays failed until something in its key changes. A task that fails as a whole (an exception, an input that cannot be read, a worker that dies) leaves no result and runs again -next time. +next time; when a worker dies, the tasks of its job after the failed one run in another worker. ### depot plan diff --git a/src/nnnotes/pipeline/orchestrate.py b/src/nnnotes/pipeline/orchestrate.py index 6d54736..8006c1b 100644 --- a/src/nnnotes/pipeline/orchestrate.py +++ b/src/nnnotes/pipeline/orchestrate.py @@ -4,17 +4,23 @@ the caller plus the keys of the tasks already done), so every input is known when a task is described. Per stage: a task whose result is in the depot is a hit; the others run in parallel, then the next stage starts. -The pool: up to `workers` spawned processes (no more than a stage has tasks to run), long-lived (each imports -what its stages need once), one task at a time each. Tasks start in order of decreasing estimated CPU (longest -first), a task only when the estimated peak memory of the running tasks plus its own fits `memory` less -`worker_bytes` per worker (what the processes hold before any task); a task estimated above that budget runs alone -and is reported. A worker is replaced after `recycle_tasks` tasks, or when its resident memory after a task is -above `recycle_rss`; the workers are stopped together at the end. `workers=0` runs the tasks in this process, one -after another. +Jobs: the tasks of a stage that share a batch (Stage.batch, e.g. the tasks reading one bundle) form one job, run +one after another by one process, which loads what they share once and drops it at the job's end +(Stage.end_batch); a task without a batch is a job alone. A job's estimate is the sum of its tasks' CPU and the +largest of their peaks. + +The pool: up to `workers` spawned processes (no more than a stage has jobs to run), long-lived (each imports +what its stages need once), one job at a time each. Jobs start in order of decreasing estimated CPU (longest +first), a job only when the estimated peak memory of the running jobs plus its own fits `memory` less +`worker_bytes` per worker (what the processes hold before any task); a job estimated above that budget runs alone +and is reported. A worker is replaced after a job once it has run `recycle_tasks` tasks, or when its resident memory +at the job's end is above `recycle_rss`; the workers are stopped together at the end. `workers=0` runs the jobs in +this process, one after another. Failures: a failed object is an item of its task's result, cached with it. An exception out of a stage, an input that cannot be read, or a worker that dies fails the task: nothing of it is cached, the run records it (failures(), the run log) and goes on; a task that reads the result of a failed task fails as `upstream.failed`. +The tasks of a job after the one its worker died in run in another worker. A ConfigError (a setting or tool the run needs) stops the run: no new task starts, and run() raises it once the running tasks have ended (their results stay valid). Stage results do not depend on the number of workers, their order or their timing: a run with one worker and a run @@ -143,6 +149,8 @@ def run_one(task: Task, depot, stages: dict, force: bool) -> dict: def _worker_main(conn, depot, stages: dict, initializer, initargs) -> None: + """Run jobs ([task description], [force], batch) as they come: one reply per task, in order (none after a + ConfigError), then {"end": True, "rss"} once the batch is dropped; None stops the worker.""" if initializer is not None: initializer(depot, *initargs) while True: @@ -152,13 +160,20 @@ def _worker_main(conn, depot, stages: dict, initializer, initargs) -> None: return if msg is None: return - doc, force = msg - try: - task = Task.from_json(doc) - except Exception as e: - conn.send({"ok": False, "code": "task.incompatible", "message": str(e), "trace": "", "rss": rss()}) - continue - conn.send(run_one(task, depot, stages, force)) + docs, forces, batch = msg + for doc, force in zip(docs, forces): + try: + task = Task.from_json(doc) + except Exception as e: + reply = {"ok": False, "code": "task.incompatible", "message": str(e), "trace": "", "rss": rss()} + else: + reply = run_one(task, depot, stages, force) + conn.send(reply) + if reply.get("code") == "config": + break + if batch is not None: + stages[batch[0]].end_batch(batch[1]) + conn.send({"end": True, "rss": rss()}) class _Worker: @@ -173,12 +188,13 @@ def __init__(self, ctx, n: int, depot, stages, initializer, initargs): ) self.proc.start() child.close() - self.task: Task | None = None - self.done = 0 + self.job: Job | None = None + self.replied = 0 # replies received for the tasks of the job + self.done = 0 # tasks run - def send(self, task: Task, force: bool) -> None: - self.task = task - self.conn.send((task.to_json(), force)) + def send(self, job: Job, forces: list[bool]) -> None: + self.job, self.replied = job, 0 + self.conn.send(([t.to_json() for t in job.tasks], forces, job.batch)) def ask_stop(self) -> None: try: @@ -200,40 +216,86 @@ def stop(self, wait_s: float = 10.0) -> None: # ---------------------------------------------------------------- scheduling +@dataclass +class Job: + """Tasks one process runs one after another: those of one batch (`batch`: (stage name, Stage.batch)), or a + task alone. Its estimate: the sum of its tasks' CPU, the largest of their peaks (one runs at a time).""" + + tasks: list + batch: tuple | None = None + cpu: float = field(init=False) + peak: int = field(init=False) + + def __post_init__(self): + self.cpu = sum(_cpu(t) for t in self.tasks) + self.peak = max(_peak(t) for t in self.tasks) + + @property + def id(self) -> str: + return self.tasks[0].id + + +def jobs(tasks, stages: dict) -> list[Job]: + """The jobs of `tasks`: one per batch of a stage (its tasks by id), one per task without a batch.""" + alone, batches = [], {} + for t in sorted(tasks, key=lambda t: t.id): + b = stages[t.stage].batch(t) + if b is None: + alone.append(Job([t])) + else: + batches.setdefault((t.stage, b), []).append(t) + return alone + [Job(ts, k) for k, ts in batches.items()] + + class Scheduler: - """The start order of a stage's tasks: decreasing estimated CPU (then id), admitted while the estimated peak - memory of the running tasks plus the next one fits `budget` (bytes; None: no limit). A task estimated above the - budget waits until nothing runs, then runs alone; no task behind it in the order starts before it.""" + """The start order of a stage's jobs: decreasing estimated CPU (then id), admitted while the estimated peak + memory of the running jobs plus the next one fits `budget` (bytes; None: no limit). A job estimated above the + budget waits until nothing runs, then runs alone; no job behind it in the order starts before it.""" - def __init__(self, tasks, budget: int | None): - self.queue = sorted(tasks, key=lambda t: (-_cpu(t), t.id)) + def __init__(self, jobs, budget: int | None): + self.queue = sorted(jobs, key=self._order) self.budget = budget self.running: dict[str, int] = {} - self.oversize: list[str] = [] + self.oversize: list[str] = [] # task ids + self.stopped = False + + @staticmethod + def _order(j: Job) -> tuple: + return (-j.cpu, j.id) def used(self) -> int: return sum(self.running.values()) - def next(self) -> Task | None: + def next(self) -> Job | None: used = self.used() - for i, t in enumerate(self.queue): - need = _peak(t) - if self.budget is None or used + need <= self.budget: + for i, j in enumerate(self.queue): + if self.budget is None or used + j.peak <= self.budget: return self._take(i) - if need > self.budget: + if j.peak > self.budget: if self.running: return None - self.oversize.append(t.id) + self.oversize += [t.id for t in j.tasks] return self._take(i) return None - def _take(self, i: int) -> Task: - t = self.queue.pop(i) - self.running[t.id] = _peak(t) - return t + def _take(self, i: int) -> Job: + j = self.queue.pop(i) + self.running[j.id] = j.peak + return j + + def finished(self, jid: str) -> None: + self.running.pop(jid, None) + + def put(self, job: Job) -> None: + """Queue `job` (the rest of a job whose worker died) in its place in the order; nothing once stopped.""" + if not self.stopped: + self.queue.append(job) + self.queue.sort(key=self._order) - def finished(self, tid: str) -> None: - self.running.pop(tid, None) + def stop(self) -> None: + """Start nothing more.""" + self.queue.clear() + self.stopped = True def _cpu(t: Task) -> float: @@ -594,22 +656,26 @@ def _fail(run: Run, env: Env, tid: str, key, stage: str, code: str, message: str # ---------------------------------------------------------------- execution def _execute(self, tasks: list[Task], force, run: Run) -> dict[str, dict]: - """Run `tasks`; {task id: reply}.""" + """Run `tasks`; {task id: reply} (none for a task not started because the run stopped).""" if not tasks: return {} - sched = Scheduler(tasks, self.task_budget()) + sched = Scheduler(jobs(tasks, self.stages), self.task_budget()) replies: dict[str, dict] = {} if self.workers == 0: while True: - t = sched.next() - if t is None: + j = sched.next() + if j is None: break - self._record(run, t, sched, None) - replies[t.id] = run_one(t, self.depot, self.stages, force(t)) - self._done(run, t, replies[t.id], None) - sched.finished(t.id) - if replies[t.id].get("code") == "config": - sched.queue.clear() + self._record(run, j, sched, None) + for t in j.tasks: + replies[t.id] = run_one(t, self.depot, self.stages, force(t)) + self._done(run, t, replies[t.id], None) + if replies[t.id].get("code") == "config": + sched.stop() + break + if j.batch is not None: + self.stages[j.batch[0]].end_batch(j.batch[1]) + sched.finished(j.id) else: self._pool_run(sched, force, replies, run) run.oversize += sched.oversize @@ -620,59 +686,77 @@ def _spawn(self) -> _Worker: return _Worker(mp.get_context("spawn"), self._started, self.depot, self.stages, self.initializer, self.initargs) def _pool_run(self, sched: Scheduler, force, replies: dict, run: Run) -> None: - while len(self._pool) < min(self.workers, len(sched.queue)): # no more processes than tasks + while len(self._pool) < min(self.workers, len(sched.queue)): # no more processes than jobs self._pool.append(self._spawn()) try: while sched.queue or sched.running: for i in range(len(self._pool)): - if self._pool[i].task is None: - t = sched.next() - if t is None: + if self._pool[i].job is None: + j = sched.next() + if j is None: break if not self._pool[i].proc.is_alive(): # died while idle self._replace(self._pool[i]) w = self._pool[i] - self._record(run, t, sched, w) - w.send(t, force(t)) - busy = [w for w in self._pool if w.task is not None] + self._record(run, j, sched, w) + w.send(j, [force(t) for t in j.tasks]) + busy = [w for w in self._pool if w.job is not None] ready = wait([w.conn for w in busy] + [w.proc.sentinel for w in busy]) for w in busy: - if w.conn not in ready and w.proc.sentinel not in ready: - continue - t = w.task - assert t is not None # a busy worker has its task - reply = None - try: - if w.conn.poll(): - reply = w.conn.recv() - except (EOFError, OSError): - reply = None - if reply is None: - w.proc.join(5) - reply = { - "ok": False, - "code": "task.worker_died", - "message": f"worker exited with code {w.proc.exitcode} while running the task", - "trace": "", - "rss": None, - } - replies[t.id] = reply - self._done(run, t, reply, w) - sched.finished(t.id) - if reply.get("code") == "config": - sched.queue.clear() - w.task = None - w.done += 1 - if not w.proc.is_alive() or reply.get("code") == "task.worker_died" or self._recycle(w, reply): - self._replace(w) + if w.conn in ready or w.proc.sentinel in ready: + self._collect(w, sched, replies, run) except BaseException: self.close(kill=True) raise - def _recycle(self, w: _Worker, reply: dict) -> bool: + def _collect(self, w: _Worker, sched: Scheduler, replies: dict, run: Run) -> None: + """Take what worker `w` sent: the replies to its job's tasks, in order, then the job's end. A worker that + died fails the task it was running; the job's tasks after it are queued again.""" + job = w.job + assert job is not None # a busy worker has its job + while True: + try: + if not w.conn.poll(): + if w.proc.is_alive(): + return # more to come + break + msg = w.conn.recv() + except (EOFError, OSError): + break + if msg.get("end"): + sched.finished(job.id) + w.job = None + w.done += w.replied + if not w.proc.is_alive() or self._recycle(w, msg.get("rss")): + self._replace(w) + return + t = job.tasks[w.replied] + w.replied += 1 + replies[t.id] = msg + self._done(run, t, msg, w) + if msg.get("code") == "config": + sched.stop() + w.proc.join(5) + if w.replied < len(job.tasks): + t = job.tasks[w.replied] + replies[t.id] = { + "ok": False, + "code": "task.worker_died", + "message": f"worker exited with code {w.proc.exitcode} while running the task", + "trace": "", + "rss": None, + } + self._done(run, t, replies[t.id], w) + rest = job.tasks[w.replied + 1 :] + if rest: + sched.put(Job(rest, job.batch)) + sched.finished(job.id) + w.job = None + self._replace(w) + + def _recycle(self, w: _Worker, r: int | None) -> bool: if self.recycle_tasks and w.done >= self.recycle_tasks: return True - r = reply.get("rss") return bool(self.recycle_rss and r is not None and r > self.recycle_rss) def _replace(self, w: _Worker) -> None: @@ -683,18 +767,20 @@ def _replace(self, w: _Worker) -> None: w.conn.close() self._pool[self._pool.index(w)] = self._spawn() - def _record(self, run: Run, t: Task, sched: Scheduler, w) -> None: - run.log.append( - { - "event": "start", - "task": t.id, - "key": t.key, - "worker": None if w is None else w.n, - "estimate": t.cost.to_json() if t.cost is not None else None, - "runningBytes": sched.used() - _peak(t), - "budget": sched.budget, - } - ) + def _record(self, run: Run, j: Job, sched: Scheduler, w) -> None: + for t in j.tasks: + run.log.append( + { + "event": "start", + "task": t.id, + "key": t.key, + "job": j.id, + "worker": None if w is None else w.n, + "estimate": t.cost.to_json() if t.cost is not None else None, + "runningBytes": sched.used() - j.peak, + "budget": sched.budget, + } + ) def _done(self, run: Run, t: Task, reply: dict, w) -> None: e = { diff --git a/src/nnnotes/pipeline/plan.py b/src/nnnotes/pipeline/plan.py index 1cba6ee..e5184ea 100644 --- a/src/nnnotes/pipeline/plan.py +++ b/src/nnnotes/pipeline/plan.py @@ -28,7 +28,7 @@ from . import contract from .depot import write_file -from .orchestrate import WORKER_BYTES +from .orchestrate import WORKER_BYTES, jobs from .stages import Env, Pending, Stage, describe @@ -217,7 +217,7 @@ def plan( "cost": t.cost.to_json() if t.cost is not None else None, } ) - return Plan(_doc(nodes, edges, runnable, workers, memory, worker_bytes, layouts or {}), runnable) + return Plan(_doc(nodes, edges, jobs(runnable, stages), workers, memory, worker_bytes, layouts or {}), runnable) def _run_reasons(t, old: dict | None, depot, forced: bool) -> list[dict]: @@ -241,21 +241,22 @@ def _run_reasons(t, old: dict | None, depot, forced: bool) -> list[dict]: def _doc( - nodes: list, edges: set, runnable: list, workers: int, memory: int | None, worker_bytes: int, layouts: dict + nodes: list, edges: set, todo: list, workers: int, memory: int | None, worker_bytes: int, layouts: dict ) -> dict: + """The plan document; `todo`: the jobs (orchestrate.jobs) of the runnable tasks.""" count = {s: sum(n["status"] == s for n in nodes) for s in ("hit", "run", "unknown")} - peaks = sorted((t.cost.peak_bytes if t.cost else 0 for t in runnable), reverse=True) + peaks = sorted((j.peak for j in todo), reverse=True) procs = max(1, workers) budget = None if memory is None else max(0, memory - procs * worker_bytes) peak = sum(peaks[:procs]) if budget is not None and peaks: peak = min(peak, max(budget, peaks[0])) peak += procs * worker_bytes if peaks else 0 - oversize = sorted({t.id for t in runnable if budget is not None and t.cost and t.cost.peak_bytes > budget}) + oversize = sorted(t.id for j in todo if budget is not None and j.peak > budget for t in j.tasks) summary = { "tasks": len(nodes), **count, - "cpuSeconds": round(sum(t.cost.cpu_seconds for t in runnable if t.cost), 3), + "cpuSeconds": round(sum(j.cpu for j in todo), 3), "peakBytes": int(peak), "budget": memory, "oversize": oversize, diff --git a/src/nnnotes/pipeline/stages.py b/src/nnnotes/pipeline/stages.py index ced2e4b..1435c28 100644 --- a/src/nnnotes/pipeline/stages.py +++ b/src/nnnotes/pipeline/stages.py @@ -102,6 +102,14 @@ def estimate(self, subject: str, env: Env, inputs: list[Input]) -> Cost: def run(self, task: Task, depot) -> Output: raise NotImplementedError + def batch(self, task: Task) -> str | None: + """The batch of a task to run: the tasks of a stage with the same batch run one after another in one process + (a job of the orchestrator), so what they share is loaded once. None: the task runs alone.""" + return None + + def end_batch(self, batch: str) -> None: + """The process ran the last task of `batch`: drop what the stage kept for it.""" + def registry(*stages: Stage) -> dict[str, Stage]: """{name: stage}; names must be unique.""" diff --git a/tests/pipeline/test_orchestrate.py b/tests/pipeline/test_orchestrate.py index a4cee60..47c2914 100644 --- a/tests/pipeline/test_orchestrate.py +++ b/tests/pipeline/test_orchestrate.py @@ -1,6 +1,7 @@ """The local orchestrator (toy stages of test_stages; spawned worker processes where the pool matters).""" import hashlib +import os from pathlib import Path import pytest @@ -8,10 +9,10 @@ from nnnotes.pipeline import contract from nnnotes.pipeline.contract import Cost, Task from nnnotes.pipeline.depot import Depot -from nnnotes.pipeline.orchestrate import Orchestrator, Scheduler, coverage +from nnnotes.pipeline.orchestrate import Job, Orchestrator, Scheduler, coverage, jobs from nnnotes.pipeline.plan import plan -from nnnotes.pipeline.stages import execute -from tests.pipeline.test_stages import PIPELINE, toy_files, toy_stages +from nnnotes.pipeline.stages import execute, registry +from tests.pipeline.test_stages import PIPELINE, Count, Src, Total, toy_files, toy_stages TEXTS = { "a": "alpha\nbeta\n\nFAIL here\n?what", @@ -144,30 +145,126 @@ def test_a_bumped_stage_reruns_itself_and_what_reads_it(tmp_path): assert {t for t, v in st.items() if v == "ran"} == {f"toy.count:{s}" for s in TEXTS} | {"toy.total:all"} +# ---------------------------------------------------------------- batches +class Batched(Src): + """toy.src whose tasks are batched by the first letter of their subject; every process logs to /.log + what it loads, runs and drops.""" + + def __init__(self, log: Path): + self.log, self.loaded = str(log), None + + def batch(self, task): + return task.subject[0] + + def run(self, task, depot): + if self.loaded != task.subject[0]: + self.loaded = task.subject[0] + self._note(f"load {self.loaded}") + self._note(f"run {task.subject}") + return super().run(task, depot) + + def end_batch(self, batch): + self._note(f"drop {batch}") + self.loaded = None + + def _note(self, line: str) -> None: + with open(Path(self.log) / f"{os.getpid()}.log", "a", encoding="utf-8") as f: + f.write(line + "\n") + + +def batched(tmp_path, name, texts, workers): + log = tmp_path / f"{name}-log" + log.mkdir() + stages = registry(Batched(log), Count(), Total()) + depot, run, orch = orchestrate(tmp_path, name, texts, workers=workers, stages=stages) + return depot, run, orch, [" ".join(f.read_text(encoding="utf-8").split()) for f in sorted(log.iterdir())] + + +BATCHED = {"a1": "x", "a2": "y\nz", "a3": "w", "b1": "v", "b2": "FAIL"} + + +def test_the_tasks_of_a_batch_run_in_one_process_which_drops_it_after_the_last(tmp_path): + plain, _, _ = orchestrate(tmp_path, "plain", BATCHED) + for workers in (0, 2): + depot, run, orch, logs = batched(tmp_path, f"w{workers}", BATCHED, workers) + assert tree(depot.root, "cas", "ac") == tree(plain.root, "cas", "ac") # batches change no result + assert all(v["status"] == "ran" for v in run.tasks.values()) + a, b = "load a run a1 run a2 run a3 drop a", "load b run b1 run b2 drop b" + if workers: + assert sorted(logs) == [a, b] and orch._started == 2 + else: + assert logs == [f"{a} {b}"] # equal estimates: by id + jobs_of = {e["task"]: e["job"] for e in run.log if e["event"] == "start" and e["task"].startswith("toy.src")} + assert jobs_of == {f"toy.src:{s}": f"toy.src:{s[0]}1" for s in BATCHED} + + +def test_a_worker_dying_in_a_job_fails_its_task_and_the_rest_runs_in_another(tmp_path): + texts = {"a1": "x", "a2": "DIE", "a3": "y", "a4": "z", "b1": "v"} + _, run, orch, logs = batched(tmp_path, "s", texts, workers=1) + st = {t: v["status"] for t, v in run.tasks.items() if t.startswith("toy.src")} + assert st == { + "toy.src:a1": "ran", + "toy.src:a2": "failed", + "toy.src:a3": "ran", + "toy.src:a4": "ran", + "toy.src:b1": "ran", + } + assert [f["code"] for f in run.failures() if f["task"] == "toy.src:a2"] == ["task.worker_died"] + assert sorted(logs) == ["load a run a1 run a2", "load a run a3 run a4 drop a load b run b1 drop b"] + assert orch._started >= 2 + + +def test_jobs_group_a_stages_batches(): + def task(stage, name, cpu, peak): + return Task(stage, 1, name, cost=Cost(cpu, peak)) + + class Firsts(Src): + def batch(self, task): + return task.subject[0] if task.subject != "a9" else None + + stages = registry(Firsts(), Count()) + ts = [task("toy.src", s, 1.0, p) for s, p in (("a1", 10), ("a2", 30), ("b1", 5), ("a9", 7))] + ts.append(task("toy.count", "a1", 2.0, 1)) + js = jobs(reversed(ts), stages) + assert [([t.subject for t in j.tasks], j.batch) for j in js] == [ + (["a1"], None), + (["a9"], None), + (["a1", "a2"], ("toy.src", "a")), + (["b1"], ("toy.src", "b")), + ] + assert (js[2].id, js[2].cpu, js[2].peak) == ("toy.src:a1", 2.0, 30) + + # ---------------------------------------------------------------- admission -def toy_task(name, cpu, peak): - return Task("toy.src", 1, name, cost=Cost(cpu, peak)) +def toy_job(name, cpu, peak): + return Job([Task("toy.src", 1, name, cost=Cost(cpu, peak))]) def test_scheduler_orders_longest_first_and_admits_within_the_budget(): - s = Scheduler([toy_task("a", 1, 40), toy_task("b", 5, 70), toy_task("c", 3, 30), toy_task("d", 3, 20)], 100) - assert s.next().subject == "b" # longest first - assert s.next().subject == "c" # 70 + 30 fits + s = Scheduler([toy_job("a", 1, 40), toy_job("b", 5, 70), toy_job("c", 3, 30), toy_job("d", 3, 20)], 100) + assert s.next().id == "toy.src:b" # longest first + assert s.next().id == "toy.src:c" # 70 + 30 fits assert s.next() is None # d (20) would pass 100 s.finished("toy.src:b") - assert [s.next().subject, s.next().subject] == ["d", "a"] + assert [s.next().id, s.next().id] == ["toy.src:d", "toy.src:a"] assert s.next() is None and s.used() == 90 and s.oversize == [] + s.put(toy_job("e", 9, 1)) + assert s.next().id == "toy.src:e" # queued again in its place + s.stop() + s.put(toy_job("f", 1, 1)) + assert s.next() is None and s.queue == [] -def test_scheduler_runs_an_oversize_task_alone(): - s = Scheduler([toy_task("big", 1, 500), toy_task("x", 9, 10), toy_task("y", 0.5, 10)], 100) - assert s.next().subject == "x" +def test_scheduler_runs_an_oversize_job_alone(): + big = Job([Task("toy.src", 1, n, cost=Cost(0.5, 500)) for n in ("big1", "big2")]) + s = Scheduler([big, toy_job("x", 9, 10), toy_job("y", 0.5, 10)], 100) + assert s.next().id == "toy.src:x" assert s.next() is None # big waits for the pool to drain; y must not pass it s.finished("toy.src:x") - assert s.next().subject == "big" and s.oversize == ["toy.src:big"] + assert s.next() is big and s.oversize == ["toy.src:big1", "toy.src:big2"] assert s.next() is None # nothing joins it - s.finished("toy.src:big") - assert s.next().subject == "y" + s.finished("toy.src:big1") + assert s.next().id == "toy.src:y" def test_pool_admission_respects_the_budget(tmp_path): diff --git a/tests/pipeline/test_plan.py b/tests/pipeline/test_plan.py index df9bbe5..76cb105 100644 --- a/tests/pipeline/test_plan.py +++ b/tests/pipeline/test_plan.py @@ -7,7 +7,8 @@ from nnnotes.pipeline.depot import Depot from nnnotes.pipeline.orchestrate import Orchestrator from nnnotes.pipeline.plan import plan, why -from tests.pipeline.test_stages import PIPELINE, toy_files, toy_stages +from nnnotes.pipeline.stages import registry +from tests.pipeline.test_stages import PIPELINE, Count, Src, Total, toy_files, toy_stages TEXTS = {"a": "alpha\nbeta", "b": "gamma", "c": "delta\n?x"} @@ -165,6 +166,13 @@ def test_peak_estimate_oversize_and_layout_counts(tmp_path): assert "layout original: write 0 / remove 0 / keep 4" in p.text() p = plan(toy_stages(), PIPELINE, depot, facts=facts, workers=2, memory=None, worker_bytes=10) assert p.doc["summary"]["peakBytes"] == 1000 + 100 + 2 * 10 + + class OneBatch(Src): + def batch(self, task): + return "all" + + p = plan(registry(OneBatch(), Count(), Total()), PIPELINE, depot, facts=facts, workers=2, worker_bytes=10) + assert p.doc["summary"]["peakBytes"] == 1000 + 2 * 10 # one job: its tasks run one after another run(depot, facts) p = plan(toy_stages(), PIPELINE, depot, facts=facts, layouts={"cas": {"write": 1, "remove": 0, "keep": 3}}) assert p.check() == 1 # nothing runs, but a file would be written