Balanced generation under the pipeline, per-kind brake; STATE: type mix fix

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

harness/pipeline.py

@@ -221,7 +221,7 @@ class Pipeline:
 
     def old_generators(self):
         out = subprocess.run(["pgrep", "-fl", r"harness\.trainset run --part"], capture_output=True, text=True).stdout
-        return [l for l in out.splitlines() if "--plan plan2" not in l]
+        return [l for l in out.splitlines() if "--plan balanced" not in l and "--plan plan2" not in l]
 
     def supervisor(self):
         launched = False
@@ -232,15 +232,14 @@ class Pipeline:
             if time.time() > self.gen_deadline or os.path.exists(trainset.STOP_FLAG):
                 continue
             if not launched and not self.old_generators():
-                trainset.ensure_plan("plan2")
                 dl = self.a.gen_deadline
                 for i in range(3):
                     f = open(os.path.join(ROOT, "runs", "gen_train", f"p2_w{i}.log"), "a")
                     self.children.append(subprocess.Popen(
-                        [sys.executable, "-m", "harness.trainset", "run", "--plan", "plan2", "--part", str(i),
+                        [sys.executable, "-m", "harness.trainset", "run", "--plan", "balanced", "--part", str(i),
                          "--parts", "3", "--deadline", dl], cwd=ROOT, stdout=f, stderr=f))
                 launched = True
-                log("plan 2 generators started (3 workers, deadline", dl + ")")
+                log("balanced generators started (3 workers, deadline", dl + ")")
             elif launched:
                 for i, c in enumerate(self.children):
                     if c.poll() not in (None, 0) and crashes < 3:
@@ -248,7 +247,7 @@ class Pipeline:
                         log("generator", i, "exited with", c.returncode, "- restarted")
                         f = open(os.path.join(ROOT, "runs", "gen_train", f"p2_w{i}.log"), "a")
                         self.children[i] = subprocess.Popen(
-                            [sys.executable, "-m", "harness.trainset", "run", "--plan", "plan2", "--part", str(i),
+                            [sys.executable, "-m", "harness.trainset", "run", "--plan", "balanced", "--part", str(i),
                              "--parts", "3", "--deadline", self.a.gen_deadline], cwd=ROOT, stdout=f, stderr=f)
             # K variants: free text, about 10 % of the other accepted training tasks
             if k_proc is None or k_proc.poll() is not None:

harness/trainset.py

@@ -133,6 +133,28 @@ def backlog():
     return len(acc - done)
 
 
+def backlog_by_kind():
+    """{kind: accepted tasks without a first trajectory run}."""
+    done = set()
+    sp = os.path.join(ROOT, "runs", "traj", "summary.jsonl")
+    if os.path.exists(sp):
+        done = {json.loads(l)["task"] for l in open(sp) if json.loads(l)["attempt"] == 0}
+    out = {}
+    for f in glob.glob(os.path.join(POOL, "G*", "generation.json")):
+        tid = os.path.basename(os.path.dirname(f))
+        try:
+            ok = json.load(open(f)).get("accepted")
+        except (OSError, ValueError):
+            ok = False
+        if ok and tid not in done:
+            k = mix.kind_of_task_dir(tid)
+            out[k] = out.get(k, 0) + 1
+    return out
+
+
+BAL_KIND_BACKLOG = 8      # a kind with more waiting tasks than this is not generated (its trajectories come first)
+
+
 def run(part, parts, target, stop_ledger, plan_name="plan", deadline=None):
     plan = ensure_plan(plan_name, stop_ledger, target)
     stop_at = plan["ledger_at_start"] + plan["stop_ledger"]
@@ -237,8 +259,6 @@ def run_balanced(part, parts, deadline):
     base_url = os.environ.get("LLM_BASE_URL", "http://127.0.0.1:11434/v1")
     evals = overlap.load_pool("eval")
     while True:
-        while backlog() > BACKLOG_LIMIT and not os.path.exists(STOP_FLAG) and not (deadline and time.time() > deadline):
-            time.sleep(120)
         if os.path.exists(STOP_FLAG):
             print("STOP flag", flush=True)
             return
@@ -258,7 +278,12 @@ def run_balanced(part, parts, deadline):
         blocked = {k for k in mix.TYPE_SHARE if att.get(k, 0) >= 6 and acc.get(k, 0) < 0.2 * att.get(k, 0)}
         if blocked:
             print("kinds skipped (low acceptance):", sorted(blocked), flush=True)
-        kind = mix.deficit_pick(counts, allowed=set(mix.TYPE_SHARE) - blocked)
+        waiting = {k for k, n in backlog_by_kind().items() if n > BAL_KIND_BACKLOG}
+        allowed = set(mix.TYPE_SHARE) - blocked - waiting
+        if not allowed:
+            time.sleep(120)     # every kind has a backlog: the trajectories are the slower side
+            continue
+        kind = mix.deficit_pick(counts, allowed=allowed)
         n, claim = _claim_slot()
         if n is None:
             print("no free slot", flush=True)