SniperGold_ML/engine/dispatcher.py

"""Dispatcher / ingestion orchestration (Layer 4, spec 10 + 13).

Responsibilities:
  - G-2 worker creation bounded by CPU and memory estimates (recorded; never
    claimed as CPU saturation)
  - deterministic ordered dispatch; exactly one predefined recovery per
    failure class; bounded redispatch (retry_limit=2)
  - hash-chained commit journal; checkpoint per workload; lock lifecycle
  - sequential carry-aware bar aggregation over ordered canonical chunks
  - fail-closed state machine; RESUME_BLOCKED guards; telemetry
"""

import json
import os
import threading
import time
from multiprocessing import get_context

from . import versions, storage
from .aggregate import Aggregator
from .canonical import tick_chunk_content_id
from .checkpoint import (
    STATUS_INITIALIZED, STATUS_RUNNING, STATUS_CHUNK_COMMITTED, STATUS_PAUSED,
    STATUS_FAILED, STATUS_COMPLETED, STATUS_RESUME_BLOCKED,
    empty_malformed_counts, validate_transition,
    write_checkpoint, load_checkpoint,
)
from .config import config_sha256, config_snapshot, validate_config
from .journal import append_commit, read_journal, seq_after
from .lock import LockHandle
from .parse import ALL_MALFORMED_CLASSES
from .storage import (
    checkpoint_path, control_path, journal_path, lock_path, progress_path,
    state_dir, ticks_chunk_path, bar_part_path, bar_dir, read_ticks_chunk,
    read_bar_part, write_bar_part, cert_path, chunkmap_path,
)
from .util import atomic_write_json, sha256_file
from .worker import worker_main

try:
    import psutil as _psutil
except ImportError:
    _psutil = None

RUN_ID_PREFIX = "RUN-"
NON_EMPTY_MALFORMED = [c for c in ALL_MALFORMED_CLASSES if c != "MALFORMED_EMPTY_LINE"]


class IngestError(Exception):
    """Material condition; run must STOP (fail-closed)."""


def new_run_id():
    import datetime
    import uuid
    stamp = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%d-%H%M%S")
    return "%s%s-%s" % (RUN_ID_PREFIX, stamp, uuid.uuid4().hex[:8])


def _utc_now_iso():
    import datetime
    return datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z")


def _machine_limits():
    cpu = os.cpu_count() or 8
    mem_mb = None
    if _psutil is not None:
        mem_mb = _psutil.virtual_memory().total // (1024 * 1024)
    return {"cpu": cpu, "memory_total_mb": mem_mb}


def workers_to_create(cfg):
    """G-2: workers_created = min(requested, cpu_count, memory bound)."""
    lim = _machine_limits()
    n = cfg["workers_requested"]
    n = min(n, lim["cpu"])
    worker_mb = cfg.get("memory_limits_mb", {}).get("worker", 1024)
    if lim["memory_total_mb"]:
        n = min(n, max(1, (lim["memory_total_mb"] // 2) // worker_mb))
    return max(1, n)


def cpu_utilization_sample(interval=1.0):
    """Parent-side measured CPU utilization; None when psutil is absent."""
    if _psutil is None:
        return None
    return _psutil.cpu_percent(interval=interval)


def _line_start_offsets(chunkmap):
    """Global 1-based line number of the first line of each chunk.
    chunk[0] starts at global line 1 (0-based offset 0 in worker code:
    parse_chunk adds +1 per line). We pass the 0-based line offset."""
    offsets = []
    acc = 0
    for c in chunkmap["chunks"]:
        offsets.append(acc)
        acc += c["line_count"]
    return offsets


def _next_chunk_from_journal(records, total_chunks):
    committed = {r["chunk"] for r in records}
    for i in range(total_chunks):
        if i not in committed:
            return i
    return total_chunks


# ---------------------------------------------------------------------------
# Run initialization
# ---------------------------------------------------------------------------

def init_run(cfg, source_path, source_id, cert, run_id, out_root, chunkmap):
    """Create a new run: directories, persisted config, INITIALIZED checkpoint.
    Refuses a live lock; VERIFIED_WITH_LIMITATION is never silently accepted."""
    from .certify import cert_matches_source
    guard = cert_matches_source(cert, source_path)
    if guard is not None:
        raise IngestError("certificate does not match source: %s" % guard)
    status = cert.get("certification_status")
    if status != "FULLY_VERIFIED":
        raise IngestError("source not FULLY_VERIFIED (cert status %r); "
                          "VERIFIED_WITH_LIMITATION requires human gate G-3" % status)
    if cert.get("source_format") != "csv_tickstory_mt5":
        raise IngestError("certified source_format %r is not csv_tickstory_mt5; "
                          "init blocked (spec 8.1 certificate/config mismatch)"
                          % cert.get("source_format"))
    cfg_grammar = cfg.get("source_grammar")
    if cfg_grammar not in versions.SUPPORTED_GRAMMARS:
        raise IngestError("configured source_grammar %r is unsupported; "
                          "init blocked (spec 8.1)" % cfg_grammar)
    expect_grammar_id = (versions.GRAMMAR_ID_TICKSTORY
                         if cfg_grammar == versions.GRAMMAR_TICKSTORY_MT5
                         else versions.GRAMMAR_ID_DOTTED)
    cert_grammar_id = cert.get("grammar_id")
    if cert_grammar_id != expect_grammar_id:
        raise IngestError(
            "certified grammar_id %r does not match configured source_grammar %r; "
            "init blocked (spec 8.1 certificate/config mismatch)"
            % (cert_grammar_id, cfg_grammar))

    dirs = [
        os.path.join(out_root, "certification"),
        os.path.join(out_root, "chunkmap"),
        os.path.join(out_root, "ticks"),
        os.path.join(out_root, "staging"),
        os.path.join(out_root, "malformed"),
        os.path.join(out_root, "source_preservation"),
        os.path.join(out_root, "bars"),
        os.path.join(out_root, "logs"),
        os.path.join(out_root, "verification"),
        os.path.join(out_root, "evidence"),
        storage.datasets_dir(out_root),
        state_dir(out_root),
    ]
    for d in dirs:
        os.makedirs(d, exist_ok=True)
    for tf in cfg["timeframes"]:
        os.makedirs(bar_dir(out_root, tf), exist_ok=True)

    cm_sha = chunkmap["chunkmap_sha256"]
    payload = {
        "schema_version": versions.CHECKPOINT_SCHEMA_VERSION,
        "run_id": run_id,
        "status": STATUS_INITIALIZED,
        "source_identity": {
            "path": os.path.abspath(source_path),
            "size": cert["source_size_bytes"],
            "mtime_ns": cert["source_mtime_ns"],
            "sha256_full": cert.get("sha256_full"),
            "cert_status": cert["certification_status"],
            "cert_path": cert_path(out_root),
        },
        "versions": {
            "engine": versions.ENGINE_VERSION,
            "parser": versions.PARSER_VERSION,
            "algorithm": versions.ALGORITHM_VERSION,
            "dataset": versions.DATASET_VERSION,
            "manifest_schema": versions.MANIFEST_SCHEMA_VERSION,
        },
        "chunkmap_sha256": cm_sha,
        "journal_tail": {"seq": -1, "hash": _empty_hash()},
        "last_completed_chunk": -1,
        "next_chunk": 0,
        "next_byte_start": 0,
        "cumulative": {
            "rows_parsed": 0, "rows_canonical": 0, "bytes_read": 0,
            "malformed": empty_malformed_counts(),
            "chunks_committed": 0,
            "bars_per_timeframe": {tf: 0 for tf in cfg["timeframes"]},
        },
        "carries": {tf: None for tf in cfg["timeframes"]},
        "config_sha256": config_sha256(cfg),
        "update_ts_utc": _utc_now_iso(),
        "status_reason": None,
    }
    write_checkpoint(checkpoint_path(out_root), payload)
    atomic_write_json(os.path.join(state_dir(out_root), "config.json"),
                      config_snapshot(cfg))
    return payload


def _empty_hash():
    import hashlib
    return hashlib.sha256(b"").hexdigest()


# ---------------------------------------------------------------------------
# Ingestion runner
# ---------------------------------------------------------------------------

class IngestRunner:
    def __init__(self, cfg, out_root, source_path, source_id, chunkmap, cert,
                 run_id, use_spawn=False, limit_chunks=None, pause_after=None,
                 worker_fn=None):
        self.cfg = cfg
        self.out_root = out_root
        self.source_path = source_path
        self.source_id = source_id
        self.chunkmap = chunkmap
        self.cert = cert
        self.run_id = run_id
        self.use_spawn = use_spawn
        self.limit_chunks = limit_chunks
        self.pause_after = pause_after
        self.worker_fn = worker_fn
        self.line_offsets = _line_start_offsets(chunkmap)
        self.lock = LockHandle(lock_path(out_root), run_id, os.getpid())
        self.heartbeat_stop = threading.Event()
        self.telemetry = {
            "workers_requested": cfg["workers_requested"],
            "workers_created": 1,
            "workers_with_data": 0,
            "cpu_utilization_pct": None,
        }
        self._cpu_samples = []

    # -- chunk execution ----------------------------------------------------
    def _execute_chunk(self, ci):
        c = self.chunkmap["chunks"][ci]
        cfg_path = os.path.join(state_dir(self.out_root), "config.json")
        args = (cfg_path, self.source_path, self.out_root, self.source_id,
                ci, c["byte_start"], c["byte_end"], self.line_offsets[ci],
                (ci == 0), _utc_now_ms())
        if self.worker_fn is not None:
            return self.worker_fn(*args)
        if self.use_spawn:
            ctx = get_context("spawn")
            with ctx.Pool(1) as pool:
                return pool.apply(worker_main, args)
        return worker_main(*args)

    # -- heartbeat ----------------------------------------------------------
    def _heartbeat_loop(self):
        while not self.heartbeat_stop.wait(30):
            self.lock.touch()
            sample = cpu_utilization_sample(interval=0.2)
            if sample is not None:
                self._cpu_samples.append(sample)

    # -- aggregation --------------------------------------------------------
    def _aggregate_workload(self, aggregator, wi, is_last, cum_malformed, prev_ts):
        wl = self.chunkmap["workloads"][wi]
        for ci in wl["chunks"]:
            rows = read_ticks_chunk(ticks_chunk_path(self.out_root, ci))
            for rec in rows:
                if prev_ts[0] is not None and rec[0] < prev_ts[0]:
                    cum_malformed["MALFORMED_NON_MONOTONIC_BOUNDARY"] += 1
                    continue
                prev_ts[0] = rec[0]
                aggregator.consume(rec)
        for tf in self.cfg["timeframes"]:
            rows = aggregator.finish_eof(tf) if is_last else aggregator.finish_workload(tf)
            # A part file is ALWAYS written per processed workload (empty when
            # no bar was finalized). An empty part is an explicit
            # aggregation-completeness marker for deterministic resume; it
            # contains zero bar rows, so no empty-bar inflation ever occurs.
            pp = bar_part_path(self.out_root, tf, wi)
            os.makedirs(os.path.dirname(pp), exist_ok=True)
            write_bar_part(pp, rows, tf, wi, self.cfg)

    # -- checkpoint ---------------------------------------------------------
    def _write_checkpoint(self, status, reason, ckpt, journal_records,
                          journal_tail_hash, next_chunk, carries, bar_counts,
                          cumsum):
        payload = dict(ckpt)
        payload["status"] = status
        payload["journal_tail"] = {
            "seq": (journal_records[-1]["seq"] if journal_records else -1),
            "hash": journal_tail_hash,
        }
        payload["last_completed_chunk"] = next_chunk - 1
        payload["next_chunk"] = next_chunk
        ch = (self.chunkmap["chunks"][next_chunk]
              if next_chunk < len(self.chunkmap["chunks"]) else None)
        payload["next_byte_start"] = (ch["byte_start"]
                                      if ch else self.chunkmap["total_bytes"])
        payload["carries"] = carries
        payload["cumulative"] = dict(cumsum)
        payload["cumulative"]["bars_per_timeframe"] = bar_counts
        payload["cumulative"]["malformed"] = dict(cumsum["malformed"])
        payload["update_ts_utc"] = _utc_now_iso()
        payload["status_reason"] = reason
        write_checkpoint(checkpoint_path(self.out_root), payload)
        return payload

    # -- main loop ----------------------------------------------------------
    def run(self):
        ckpt = load_checkpoint(checkpoint_path(self.out_root))
        if ckpt.get("run_id") != self.run_id:
            raise IngestError("checkpoint run_id mismatch")
        if self.chunkmap["chunkmap_sha256"] != ckpt.get("chunkmap_sha256"):
            raise IngestError("chunkmap_sha256 mismatch vs checkpoint")
        validate_transition(ckpt["status"], STATUS_RUNNING)

        self.lock.acquire()
        self.telemetry["workers_created"] = workers_to_create(self.cfg)
        self.telemetry["cpu_utilization_pct"] = None
        hb = threading.Thread(target=self._heartbeat_loop, daemon=True)
        hb.start()
        try:
            return self._dispatch_loop(ckpt)
        finally:
            self.heartbeat_stop.set()
            hb.join(timeout=2)
            if self._cpu_samples:
                self.telemetry["cpu_utilization_pct"] = round(
                    sum(self._cpu_samples) / len(self._cpu_samples), 1)
            self.lock.release()

    def _dispatch_loop(self, ckpt):
        cfg = self.cfg
        journal_records, journal_tail_hash = read_journal(journal_path(self.out_root))
        committed = {r["chunk"] for r in journal_records}
        next_chunk = _next_chunk_from_journal(journal_records, len(self.chunkmap["chunks"]))
        if next_chunk == len(self.chunkmap["chunks"]) and not self.chunkmap["workloads"]:
            return STATUS_CHUNK_COMMITTED

        carries = {tf: ckpt["carries"].get(tf) for tf in cfg["timeframes"]}
        bar_counts = dict(ckpt["cumulative"]["bars_per_timeframe"])
        last_final = _seed_last_final(self.out_root, cfg, carries, bar_counts,
                                      self.chunkmap)
        aggregator = Aggregator(cfg, carries=carries, bar_counts=bar_counts,
                                last_final_period=last_final)

        seq = seq_after(journal_records)
        prev_hash = journal_tail_hash
        cum = dict(ckpt["cumulative"])
        cum["malformed"] = dict(cum["malformed"])
        attempts = {}
        prev_ts = [None]
        started = time.time()
        limit_rate = cfg.get("malformed_rate_limit", 0.001)
        limit_abs = cfg.get("malformed_abs_cap", 10_000)

        for wi, wl in enumerate(self.chunkmap["workloads"]):
            pending = [ci for ci in wl["chunks"] if ci not in committed]
            if self.limit_chunks is not None:
                pending = [ci for ci in pending if ci < self.limit_chunks]
            wl_mal_start = {k: cum["malformed"][k] for k in ALL_MALFORMED_CLASSES}
            wl_rows_start = cum["rows_parsed"]
            while pending:
                ci = pending.pop(0)
                if attempts.get(ci, 0) > cfg["retry_limit"]:
                    return self._fail("worker_retry_limit", ckpt, journal_records,
                                      prev_hash, aggregator, cum)
                try:
                    stat = self._execute_chunk(ci)
                except Exception as e:
                    attempts[ci] = attempts.get(ci, 0) + 1
                    if attempts[ci] > cfg["retry_limit"]:
                        return self._fail("worker_exception:%s" % e, ckpt,
                                          journal_records, prev_hash,
                                          aggregator, cum)
                    pending.append(ci)
                    continue
                prev_hash = append_commit(journal_path(self.out_root), seq,
                                          stat["chunk_index"], stat["content_id"],
                                          stat["file_sha256"], stat["ts_utc"],
                                          prev_hash)
                journal_records.append({"seq": seq, "chunk": ci,
                                        "content_id": stat["content_id"],
                                        "file_sha256": stat["file_sha256"],
                                        "prev": None, "ts_utc": stat["ts_utc"]})
                seq += 1
                committed.add(ci)
                cum["rows_parsed"] += stat["rows_parsed"]
                cum["rows_canonical"] += stat["canonical_rows"]
                c = self.chunkmap["chunks"][ci]
                cum["bytes_read"] += c["byte_end"] - c["byte_start"]
                cum["chunks_committed"] += 1
                for k in ALL_MALFORMED_CLASSES:
                    cum["malformed"][k] += stat["malformed"].get(k, 0)

            # Per-workload malformed policy (F17: carry counters, check at end)
            wl_mal = {k: cum["malformed"][k] - wl_mal_start[k]
                      for k in ALL_MALFORMED_CLASSES}
            wl_rows = cum["rows_parsed"] - wl_rows_start
            non_empty = sum(wl_mal[k] for k in NON_EMPTY_MALFORMED)
            if wl_rows > 0 and (non_empty > limit_abs
                                or non_empty / wl_rows > limit_rate):
                return self._fail("malformed_rate_exceeded", ckpt,
                                  journal_records, prev_hash, aggregator, cum)

            # Aggregation of a workload happens at most once (append-only);
            # on resume a workload whose bar parts already exist is skipped and
            # the aggregator state is restored from the checkpoint.
            part_exists = any(
                os.path.exists(bar_part_path(self.out_root, tf, wi))
                for tf in cfg["timeframes"])
            if not part_exists:
                try:
                    self._aggregate_workload(aggregator, wi, wi == len(self.chunkmap["workloads"]) - 1,
                                             cum["malformed"], prev_ts)
                except ValueError as e:
                    return self._fail("carry_continuity:%s" % e, ckpt,
                                      journal_records, prev_hash, aggregator, cum)
            bar_counts = {tf: aggregator.bar_counts[tf] for tf in cfg["timeframes"]}
            next_chunk = _next_chunk_from_journal(journal_records,
                                                  len(self.chunkmap["chunks"]))
            ckpt = self._write_checkpoint(STATUS_CHUNK_COMMITTED, None, ckpt,
                                          journal_records, prev_hash, next_chunk,
                                          aggregator.carr_state(), bar_counts, cum)
            self._write_progress(wi, started)
            if self.pause_after is not None and wi + 1 >= self.pause_after:
                ckpt = self._write_checkpoint(STATUS_PAUSED, "pause_after_workload",
                                              ckpt, journal_records, prev_hash,
                                              next_chunk, aggregator.carr_state(),
                                              bar_counts, cum)
                return STATUS_PAUSED
            if self._control_requested():
                ckpt = self._write_checkpoint(STATUS_PAUSED, "control_requested",
                                              ckpt, journal_records, prev_hash,
                                              next_chunk, aggregator.carr_state(),
                                              bar_counts, cum)
                return STATUS_PAUSED

        return STATUS_CHUNK_COMMITTED

    def _control_requested(self):
        cp = control_path(self.out_root)
        if os.path.exists(cp):
            try:
                with open(cp, "r", encoding="utf-8") as fh:
                    ctl = json.load(fh)
                return ctl.get("action") in ("pause", "stop")
            except Exception:
                return False
        return False

    def _fail(self, reason, ckpt, journal_records, journal_tail_hash,
              aggregator, cumsum):
        next_chunk = _next_chunk_from_journal(journal_records, len(self.chunkmap["chunks"]))
        self._write_checkpoint(STATUS_FAILED, reason, ckpt, journal_records,
                               journal_tail_hash, next_chunk,
                               aggregator.carr_state(),
                               {tf: aggregator.bar_counts[tf]
                                for tf in self.cfg["timeframes"]},
                               cumsum)
        return STATUS_FAILED

    def _write_progress(self, wi, started):
        total_wl = len(self.chunkmap["workloads"])
        progress = {
            "run_id": self.run_id,
            "status": "RUNNING",
            "chunks_done": self.chunkmap["chunks"][wi]["index"] + 1,
            "chunks_total": len(self.chunkmap["chunks"]),
            "bytes": None,
            "workloads": {"done": wi + 1, "total": total_wl},
            "workers": self.telemetry,
            "started_utc": _utc_now_iso(),
            "updated_utc": _utc_now_iso(),
            "eta_sec": None,
        }
        atomic_write_json(progress_path(self.out_root), progress)


def _utc_now_ms():
    import datetime
    return (datetime.datetime.now(datetime.timezone.utc).timestamp() * 1000)


def _seed_last_final(out_root, cfg, carries, bar_counts, chunkmap):
    """Seed the last finalized period per TF from existing bar parts."""
    last_final = {}
    for tf in cfg["timeframes"]:
        carry = carries.get(tf)
        if carry is not None:
            last_final[tf] = max(0, carry["period_id"] - 1)
            continue
        found = -1
        for wi in range(len(chunkmap["workloads"])):
            pp = bar_part_path(out_root, tf, wi)
            if os.path.exists(pp):
                rows = read_bar_part(pp)
                if rows:
                    found = max(found, rows[-1][1])
        last_final[tf] = found
    return last_final


# ---------------------------------------------------------------------------
# Resume guards (spec 13.5)
# ---------------------------------------------------------------------------

def resume_checks(cfg, out_root, source_path, chunkmap, cert):
    """Run all §13.5 guards. Returns (checkpoint, None) when RESUME is
    allowed, or (checkpoint_or_None, reason) when RESUME_BLOCKED / material."""
    p = checkpoint_path(out_root)
    try:
        ckpt = load_checkpoint(p)
    except Exception as e:
        return None, "checkpoint_corrupt:%s" % e

    from .certify import cert_matches_source
    guard = cert_matches_source(cert, source_path)
    if guard is not None:
        return ckpt, guard
    if ckpt.get("source_identity", {}).get("sha256_full") != cert.get("sha256_full"):
        return ckpt, "sha256_mismatch"
    v = ckpt.get("versions", {})
    if v.get("engine") != versions.ENGINE_VERSION:
        return ckpt, "engine_version_mismatch"
    if v.get("parser") != versions.PARSER_VERSION:
        return ckpt, "parser_version_mismatch"
    if v.get("algorithm") != versions.ALGORITHM_VERSION:
        return ckpt, "algorithm_version_mismatch"
    if ckpt.get("chunkmap_sha256") != chunkmap["chunkmap_sha256"]:
        return ckpt, "chunkmap_mismatch"

    try:
        records, tail = read_journal(journal_path(out_root))
    except Exception as e:
        return ckpt, "journal_corrupt:%s" % e
    if ckpt["journal_tail"]["seq"] != (records[-1]["seq"] if records else -1):
        return ckpt, "journal_tail_mismatch"
    if ckpt["journal_tail"]["hash"] != tail:
        return ckpt, "journal_tail_mismatch"

    for rec in records:
        pp = ticks_chunk_path(out_root, rec["chunk"])
        if not os.path.exists(pp):
            return ckpt, "committed_chunk_missing:%d" % rec["chunk"]
        rows = read_ticks_chunk(pp)
        actual = tick_chunk_content_id([(r[0], r[1], r[2], r[4]) for r in rows])
        if actual != rec["content_id"]:
            return ckpt, "committed_content_mismatch:%d" % rec["chunk"]
    if ckpt.get("status") in (STATUS_FAILED,):
        return ckpt, "failed_state_requires_human_resolution"
    if ckpt.get("status") in (STATUS_COMPLETED,):
        return ckpt, "already_completed"
    return ckpt, None


# ---------------------------------------------------------------------------
# Completion
# ---------------------------------------------------------------------------

def finalize_completed(cfg, out_root, run_id):
    """CHUNK_COMMITTED -> COMPLETED. Only valid after the independent verifier
    accepted the run; the caller enforces that gate."""
    p = checkpoint_path(out_root)
    ckpt = load_checkpoint(p)
    validate_transition(ckpt["status"], STATUS_COMPLETED)
    ckpt["status"] = STATUS_COMPLETED
    ckpt["update_ts_utc"] = _utc_now_iso()
    ckpt["status_reason"] = None
    write_checkpoint(p, ckpt)
    return ckpt