Merge pull request #3306 from MervinPraison/feat/praisonai-train-gen-guard

feat(data): in-loop cost/sponsorship guard + per-row provenance stamping

src/praisonai-train/praisonai_train/data/__init__.py

@@ -8,7 +8,7 @@ from praisonai_train.data import checks as _checks  # noqa: F401  registers chec
 from praisonai_train.data import recipes as _recipes  # noqa: F401  registers recipes
 from praisonai_train.data.benchmark import BenchResult, benchmark_deployments
 from praisonai_train.data.dedup import MinHashLSH, global_dedup, near_dedup
-from praisonai_train.data.generate import generate_dataset
+from praisonai_train.data.generate import azure_sponsorship_guard, generate_dataset
 from praisonai_train.data.intent import translation_intent
 from praisonai_train.data.qc import filter_rows, score
 from praisonai_train.data.registry import checks, recipes
@@ -16,6 +16,7 @@ from praisonai_train.data.registry import checks, recipes
 __all__ = [
     "BenchResult",
     "MinHashLSH",
+    "azure_sponsorship_guard",
     "benchmark_deployments",
     "checks",
     "filter_rows",

src/praisonai-train/praisonai_train/data/generate.py

@@ -99,8 +99,58 @@ def _call(cfg: dict, spec: dict) -> dict | None:
                 return None
     if not isinstance(row, dict) or not (row.get("instruction") and row.get("output")):
         return None
-    return {"instruction": row["instruction"], "input": row.get("input", ""),
-            "output": row["output"]}
+    out = {"instruction": row["instruction"], "input": row.get("input", ""),
+           "output": row["output"]}
+    # Per-row provenance: which model + category produced this row, so the dataset
+    # is self-describing (downstream splits/QC key off metadata, not text). Additive
+    # and backward-compatible — a key is added only when its value is available.
+    if cfg.get("deployment"):
+        out["model"] = cfg["deployment"]
+    if spec.get("task_type"):
+        out["task_type"] = spec["task_type"]
+    if spec.get("topic"):
+        out["topic"] = spec["topic"]
+    return out
+
+
+def azure_sponsorship_guard(
+    subscription_id: str, quota_id: str = "Sponsored_2016-01-01"
+) -> Callable[[], bool]:
+    """Build a ``should_continue`` predicate that is True while an Azure subscription
+    still reads as sponsored (``subscriptionPolicies.quotaId == quota_id``).
+
+    Fail-OPEN: any transient error (CLI missing, network blip, unparseable output)
+    returns True. This is a belt-and-suspenders guard against runaway spend once a
+    sponsorship converts to pay-as-you-go — an external watchdog is the hard,
+    fail-closed guard; this one must never halt a healthy run on a flaky check.
+    """
+    def _check() -> bool:
+        import subprocess
+        try:
+            out = subprocess.run(
+                ["az", "rest", "--method", "GET", "--url",
+                 f"https://management.azure.com/subscriptions/{subscription_id}"
+                 "?api-version=2022-12-01"],
+                capture_output=True, text=True, timeout=10).stdout
+            return json.loads(out)["subscriptionPolicies"]["quotaId"] == quota_id
+        except Exception:
+            return True
+
+    return _check
+
+
+def _resolve_guard(cfg: dict,
+                   should_continue: Callable[[], bool] | None) -> Callable[[], bool] | None:
+    """Pick the continuation guard: an explicit ``should_continue`` callback wins;
+    otherwise fall back to the built-in Azure sponsorship check when a subscription
+    id is supplied (``cfg['subscription_id']`` or env ``AZ_SUBSCRIPTION``). Returns
+    ``None`` when neither is configured (no guard)."""
+    if should_continue is not None:
+        return should_continue
+    sub = cfg.get("subscription_id") or os.environ.get("AZ_SUBSCRIPTION")
+    if sub:
+        return azure_sponsorship_guard(sub)
+    return None
 
 
 def _norm_row(row: dict) -> str:
@@ -111,13 +161,24 @@ def generate_dataset(
     config: dict,
     on_row: Callable[[dict], None] | None = None,
     progress_callback: Callable[[int, int, int], None] | None = None,
+    should_continue: Callable[[], bool] | None = None,
 ) -> Iterator[dict]:
     """Yield unique synthetic rows.
 
     Required config: deployment, num_examples (+ endpoint/api_key or env).
     Optional: recipe (name|dict, default 'tamil'), concurrency, start_offset,
     dedup_from (rows or JSONL paths), stop_file, azure (bool), api_version,
-    max_completion_tokens.
+    max_completion_tokens, subscription_id, sponsor_check_rows.
+
+    ``should_continue`` is an optional guard ``fn() -> bool`` re-checked every
+    ``config['sponsor_check_rows']`` completed requests (default 5000). When it
+    returns False the run halts cleanly — the stop-file is touched, pending futures
+    are cancelled and the loop breaks — the same shutdown path as the stop-file
+    circuit-breaker. If omitted but ``config['subscription_id']`` (or env
+    ``AZ_SUBSCRIPTION``) is set, the built-in Azure sponsorship guard is used
+    (see ``azure_sponsorship_guard``); with neither, generation is unguarded. The
+    guard should fail-OPEN (return True on transient errors) — an external watchdog
+    is the hard fail-closed guard; this in-loop check is belt-and-suspenders.
 
     ``progress_callback`` is an optional ``fn(done, total, kept)`` invoked as work
     completes so a CLI/monitor/notebook can show completion + remaining. ``done``
@@ -141,6 +202,11 @@ def generate_dataset(
             seen.add(_norm_row(r))
 
     stop_file = cfg.get("stop_file") or os.path.expanduser("~/.praisonai_train_stop")
+    guard = _resolve_guard(cfg, should_continue)
+    # Guard against a misconfigured interval: 0 would divide-by-zero and a
+    # negative value would silently never poll (disabling the safety check).
+    _interval = cfg.get("sponsor_check_rows")
+    check_every = max(1, int(_interval)) if _interval is not None else 5000
     specs = recipe.prompts(cfg["num_examples"], start=cfg.get("start_offset", 0))
     total = len(specs)
     done = kept = 0
@@ -157,6 +223,8 @@ def generate_dataset(
             except Exception:
                 row = None
             done += 1
+            # Emit the just-completed row first so a paid, already-finished request
+            # is never dropped by the guard/shutdown that may follow this iteration.
             emit = False
             if row:
                 key = _norm_row(row)
@@ -170,3 +238,13 @@ def generate_dataset(
                 yield row
             if progress_callback:
                 progress_callback(done, total, kept)
+            # In-loop continuation guard (billing/sponsorship self-check): every
+            # ``check_every`` completions, if the predicate says stop, halt cleanly
+            # via the same shutdown path as the stop-file breaker.
+            if guard is not None and done % check_every == 0 and not guard():
+                try:
+                    open(stop_file, "w").close()
+                except Exception:
+                    pass
+                pool.shutdown(wait=False, cancel_futures=True)
+                break

src/praisonai-train/praisonai_train/data/recipes.py

@@ -34,6 +34,12 @@ class _AxisRecipe:
                     audience=auds[(i // max(len(styles), 1)) % len(auds)],
                     variant=i // len(grid),
                 ),
+                # Ground-truth labels carried through to the emitted row so QC and
+                # the translation-share metric key off metadata, not brittle text
+                # (see ``intent.translation_intent`` / ``qc.score``). The generator
+                # stamps these onto each row as provenance.
+                "task_type": task,
+                "topic": topic,
             })
         return specs
 

src/praisonai-train/tests/unit/data/test_generate_guard.py

@@ -0,0 +1,209 @@
+"""Tests for the in-loop continuation guard and per-row provenance stamping.
+
+Network is stubbed (``_call`` is monkeypatched, or a fake recipe is used) so these
+run offline and deterministically. They exercise two additive generator features:
+
+* ``should_continue`` guard: halts generation cleanly once the predicate returns
+  False after ``sponsor_check_rows`` completions, and never halts while True; and
+* provenance stamping: emitted rows carry ``model``/``task_type``/``topic`` when
+  those are available, and omit them when they are not.
+"""
+from praisonai_train.data import generate as gen_mod
+from praisonai_train.data.generate import generate_dataset
+
+
+def _cfg(n, **extra):
+    cfg = {"endpoint": "http://x", "api_key": "k", "deployment": "d", "num_examples": n,
+           "concurrency": 1, "stop_file": "/nonexistent/stop"}
+    cfg.update(extra)
+    return cfg
+
+
+def _fake_call(cfg, spec):
+    # Unique instruction per spec so dedup keeps every row.
+    return {"instruction": spec["user"], "input": "", "output": "ஒரு முழுமையான தமிழ் பதில் இங்கே."}
+
+
+# --- Feature 1: the continuation guard ---------------------------------------
+
+def test_guard_halts_generation_after_n_rows(monkeypatch):
+    monkeypatch.setattr(gen_mod, "_call", _fake_call)
+    calls = {"n": 0}
+
+    def should_continue():
+        # True on the first check, False on the second -> halt at 2*check_every.
+        calls["n"] += 1
+        return calls["n"] < 2
+
+    # check_every=3: guard runs at done==3 (True) and done==6 (False -> halt).
+    rows = list(generate_dataset(_cfg(20, sponsor_check_rows=3),
+                                 should_continue=should_continue))
+
+    assert calls["n"] == 2                 # guard was polled exactly twice
+    # Halted at the second poll (done==6); the run stopped well short of 20 rows.
+    assert len(rows) <= 6
+    assert len(rows) < 20
+
+
+def test_guard_does_not_halt_while_true(monkeypatch):
+    monkeypatch.setattr(gen_mod, "_call", _fake_call)
+    polls = {"n": 0}
+
+    def should_continue():
+        polls["n"] += 1
+        return True                        # always healthy
+
+    n = 12
+    rows = list(generate_dataset(_cfg(n, sponsor_check_rows=3),
+                                 should_continue=should_continue))
+
+    assert len(rows) == n                  # nothing halted; all rows emitted
+    assert polls["n"] >= 1                  # guard was actually exercised
+
+
+def test_no_guard_by_default(monkeypatch):
+    # Backward compatible: no should_continue and no subscription_id => unguarded.
+    monkeypatch.setattr(gen_mod, "_call", _fake_call)
+    rows = list(generate_dataset(_cfg(5)))
+    assert len(rows) == 5
+
+
+def test_guard_resolved_from_subscription_id(monkeypatch):
+    # When no callback is passed but a subscription id is present, the built-in
+    # Azure sponsorship guard is wired in. Stub the factory to observe wiring
+    # without touching the network / az CLI.
+    monkeypatch.setattr(gen_mod, "_call", _fake_call)
+    built = {"sub": None, "polls": 0}
+
+    def fake_factory(sub):
+        built["sub"] = sub
+
+        def _check():
+            built["polls"] += 1
+            return False                   # not sponsored -> halt
+        return _check
+
+    monkeypatch.setattr(gen_mod, "azure_sponsorship_guard", fake_factory)
+    rows = list(generate_dataset(_cfg(20, subscription_id="sub-123",
+                                      sponsor_check_rows=2)))
+
+    assert built["sub"] == "sub-123"        # factory built from cfg subscription_id
+    assert built["polls"] >= 1              # guard was polled
+    assert len(rows) < 20                   # halted before completing
+
+
+def test_guard_emits_trigger_row_before_halt(monkeypatch):
+    # The completed request that triggers the halt must still be emitted — it is a
+    # paid, finished row and must not be silently dropped by the shutdown path.
+    monkeypatch.setattr(gen_mod, "_call", _fake_call)
+
+    def should_continue():
+        return False                       # halt at the first check
+
+    rows = list(generate_dataset(_cfg(20, sponsor_check_rows=1),
+                                 should_continue=should_continue))
+
+    assert len(rows) == 1                  # the trigger row is emitted, not dropped
+
+
+def test_guard_invalid_interval_does_not_crash(monkeypatch):
+    # A misconfigured interval (0 or negative) must not divide-by-zero or silently
+    # disable the guard; it is coerced to a safe minimum of 1.
+    monkeypatch.setattr(gen_mod, "_call", _fake_call)
+
+    def should_continue():
+        return False                       # halt at the first check
+
+    for bad in (0, -5):
+        rows = list(generate_dataset(_cfg(20, sponsor_check_rows=bad),
+                                     should_continue=should_continue))
+        assert len(rows) == 1              # coerced to 1 -> halts after first row
+
+
+def test_azure_sponsorship_guard_fails_open(monkeypatch):
+    # The built-in guard must return True (fail-OPEN) when the az CLI is missing /
+    # errors, so a flaky check never halts a healthy run.
+    guard = gen_mod.azure_sponsorship_guard("sub-xyz")
+
+    def boom(*a, **k):
+        raise FileNotFoundError("az not installed")
+
+    monkeypatch.setattr("subprocess.run", boom)
+    assert guard() is True
+
+
+# --- Feature 2: per-row provenance stamping ----------------------------------
+
+class _FakeRecipe:
+    """A recipe whose specs carry task_type/topic, to verify pass-through."""
+    name = "fake"
+
+    def prompts(self, n, start=0):
+        return [{"system": "s", "user": f"u{i}",
+                 "task_type": "translation", "topic": f"topic-{i}"}
+                for i in range(start, start + n)]
+
+
+class _BareRecipe:
+    """A recipe whose specs have no task_type/topic metadata."""
+    name = "bare"
+
+    def prompts(self, n, start=0):
+        return [{"system": "s", "user": f"u{i}"} for i in range(start, start + n)]
+
+
+def _patch_http(monkeypatch, payload_output="ஒரு முழுமையான தமிழ் பதில் இங்கே."):
+    """Stub urllib so the real ``_call`` (and thus its stamping) runs offline."""
+    import json as _json
+
+    class _Resp:
+        def __init__(self, user):
+            self._user = user
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *a):
+            return False
+
+        def read(self):
+            content = _json.dumps({"instruction": self._user, "input": "",
+                                   "output": payload_output})
+            return _json.dumps(
+                {"choices": [{"message": {"content": content}}]}).encode()
+
+    # urllib.request.urlopen is called with a Request built from the spec's user
+    # message; decode it to echo a unique instruction back.
+    def fake_urlopen(req, timeout=0):
+        body = _json.loads(req.data.decode())
+        return _Resp(body["messages"][1]["content"])
+
+    monkeypatch.setattr(gen_mod.urllib.request, "urlopen", fake_urlopen)
+
+
+def test_rows_carry_provenance_when_available(monkeypatch):
+    monkeypatch.setattr(gen_mod, "resolve", lambda name: _FakeRecipe())
+    _patch_http(monkeypatch)
+    rows = list(generate_dataset(_cfg(3, deployment="gpt-teacher")))
+
+    assert len(rows) == 3
+    for r in rows:
+        assert r["model"] == "gpt-teacher"           # deployment stamped
+        assert r["task_type"] == "translation"        # passed through from spec
+        assert r["topic"].startswith("topic-")        # passed through from spec
+        # Core fields remain intact.
+        assert set(["instruction", "input", "output"]).issubset(r)
+
+
+def test_rows_omit_provenance_when_absent(monkeypatch):
+    monkeypatch.setattr(gen_mod, "resolve", lambda name: _BareRecipe())
+    _patch_http(monkeypatch)
+    # deployment is required to build the request, so model is always present;
+    # task_type/topic must be OMITTED (not stamped as empty) when the spec lacks them.
+    rows = list(generate_dataset(_cfg(3, deployment="gpt-teacher")))
+
+    assert len(rows) == 3
+    for r in rows:
+        assert r["model"] == "gpt-teacher"
+        assert "task_type" not in r
+        assert "topic" not in r