graduate_code/solutions/scoring.py

"""GP2: reproducible synthetic experiment and retry-safe local recorded effects.
The crashworker CLI deliberately exits only its own process with status 23.
"""
from contextlib import contextmanager
from pathlib import Path
import argparse
import hashlib
import json
import math
import os
import re
import sqlite3
from . import data, ml, ethics


FEATURES = ["x0", "x1"]
ARTIFACT_KEYS = {"schema","features","mean","scale","weights","bias","threshold"}


def canonical(value):
    return json.dumps(value,sort_keys=True,separators=(",",":"),allow_nan=False)


def fingerprint(value):
    return hashlib.sha256(canonical(value).encode("utf-8")).hexdigest()


def finite_number(value, limit=1e6):
    return type(value) in (int,float) and math.isfinite(value) and abs(value) <= limit


def validate_artifact(artifact):
    if type(artifact) is not dict or set(artifact) != ARTIFACT_KEYS:
        raise ValueError("unexpected model schema keys")
    if type(artifact["schema"]) is not int or artifact["schema"] != 1 or artifact["features"] != FEATURES:
        raise ValueError("wrong schema version or feature order")
    for key in ("mean","scale","weights"):
        if type(artifact[key]) is not list or len(artifact[key]) != 2:
            raise ValueError("two model dimensions required")
        if not all(finite_number(x) for x in artifact[key]):
            raise ValueError("invalid numeric vector")
    if any(x < 1e-12 for x in artifact["scale"]):
        raise ValueError("scale too small")
    if not finite_number(artifact["bias"]) or not finite_number(artifact["threshold"]) or not 0 <= artifact["threshold"] <= 1:
        raise ValueError("invalid bias or threshold")
    return artifact


def _unique_object(pairs):
    result = {}
    for key,value in pairs:
        if key in result:
            raise ValueError("duplicate JSON key")
        result[key] = value
    return result


def load_artifact(path):
    with open(path,"rb") as handle:
        raw = handle.read(16385)
    if len(raw) > 16384:
        raise ValueError("artifact exceeds 16 KiB")
    try:
        value = json.loads(raw,object_pairs_hook=_unique_object)
    except (UnicodeError,json.JSONDecodeError) as error:
        raise ValueError("invalid artifact JSON") from error
    return validate_artifact(value)


def experiment(seed=21):
    rows = data.fixture(seed=seed)
    groups = [row["group"] for row in rows]
    parts = data.split_groups(groups,seed=seed)
    X = {name:[rows[i]["x"] for i in indices] for name,indices in parts.items()}
    y = {name:[rows[i]["y"] for i in indices] for name,indices in parts.items()}
    scaler = data.fit_scaler(X["train"])
    scaled = {name:data.transform(values,scaler) for name,values in X.items()}
    candidates = []
    for lam in (0.,.1,1.):
        model = ml.fit(scaled["train"],y["train"],lam=lam,rate=.1,steps=400)
        loss,_,_ = ml.loss_gradient(model["weights"],model["bias"],
                                  scaled["validation"],y["validation"],lam=0.)
        candidates.append((loss,lam,model))
    validation_loss,selected,model = min(candidates,key=lambda item:(item[0],item[1]))
    artifact = {"schema":1,"features":list(FEATURES),"mean":scaler["mean"],
                "scale":scaler["scale"],"weights":model["weights"],
                "bias":model["bias"],"threshold":.5}
    validate_artifact(artifact)
    identity = fingerprint(artifact)  # freeze before evaluating the test set
    probabilities = ml.predict(model,scaled["test"])
    metrics = ethics.binary_report(y["test"],probabilities,threshold=.5)
    report = {"seed":seed,"partitions":parts,
              "group_partitions":{name:sorted({groups[i] for i in indices})
                                  for name,indices in parts.items()},
              "selected_regularization":selected,"validation_log_loss":validation_loss,
              "test_metrics":metrics,"model_digest":identity,
              "limitation":"synthetic grouped data, not evidence of real-world utility"}
    return artifact,report


def connect_store(path, role):
    if role not in ("producer","consumer"):
        raise ValueError("unknown store role")
    connection = sqlite3.connect(path,timeout=.1,isolation_level=None,
                                 autocommit=sqlite3.LEGACY_TRANSACTION_CONTROL)
    connection.execute("PRAGMA foreign_keys=ON")
    connection.execute("PRAGMA synchronous=FULL")
    if role == "producer":
        connection.executescript("""
            CREATE TABLE IF NOT EXISTS jobs(id TEXT PRIMARY KEY,payload TEXT NOT NULL);
            CREATE TABLE IF NOT EXISTS outbox(
              id TEXT PRIMARY KEY REFERENCES jobs(id),digest TEXT NOT NULL,
              ack INTEGER NOT NULL CHECK(ack IN (0,1)));
        """)
    else:
        connection.executescript("""
            CREATE TABLE IF NOT EXISTS receipts(id TEXT PRIMARY KEY,payload TEXT NOT NULL);
            CREATE TABLE IF NOT EXISTS results(
              id TEXT PRIMARY KEY REFERENCES receipts(id),
              score REAL NOT NULL CHECK(score>=0 AND score<=1),model TEXT NOT NULL);
        """)
    return connection


@contextmanager
def transaction(connection):
    connection.execute("BEGIN IMMEDIATE")
    try:
        yield
        connection.execute("COMMIT")
    except BaseException:
        if connection.in_transaction:
            connection.execute("ROLLBACK")
        raise


def payload_for(job_id, features, model_digest):
    if type(job_id) is not str or re.fullmatch(r"[A-Za-z0-9_-]{1,64}",job_id) is None:
        raise ValueError("invalid job ID")
    if type(features) not in (list,tuple) or len(features) != 2 or not all(finite_number(x) for x in features):
        raise ValueError("two finite bounded features required")
    if type(model_digest) is not str or re.fullmatch(r"[0-9a-f]{64}",model_digest) is None:
        raise ValueError("invalid model digest")
    return {"id":job_id,"features":[float(x) for x in features],"model":model_digest}


def submit(connection, job_id, features, model_digest):
    payload = canonical(payload_for(job_id,features,model_digest))
    digest = hashlib.sha256(payload.encode()).hexdigest()
    with transaction(connection):
        old = connection.execute("SELECT payload FROM jobs WHERE id=?",(job_id,)).fetchone()
        if old is not None:
            if old[0] != payload:
                raise ValueError("conflicting job ID")
            return False
        connection.execute("INSERT INTO jobs VALUES(?,?)",(job_id,payload))
        connection.execute("INSERT INTO outbox VALUES(?,?,0)",(job_id,digest))
    return True


def consume(connection, payload, artifact):
    validate_artifact(artifact)
    raw = json.loads(payload,object_pairs_hook=_unique_object)
    if type(raw) is not dict or set(raw) != {"id","features","model"}:
        raise ValueError("invalid payload keys")
    normalized = payload_for(raw["id"],raw["features"],raw["model"])
    if canonical(normalized) != payload or raw["model"] != fingerprint(artifact):
        raise ValueError("noncanonical payload or wrong model")
    scaled = data.transform([raw["features"]],{"mean":artifact["mean"],"scale":artifact["scale"]})
    probability = ml.predict(artifact,scaled)[0]  # pure computation outside write lock
    with transaction(connection):
        old = connection.execute("SELECT payload FROM receipts WHERE id=?",(raw["id"],)).fetchone()
        if old is not None:
            if old[0] != payload:
                raise ValueError("consumer identity conflict")
            result = connection.execute("SELECT score,model FROM results WHERE id=?",(raw["id"],)).fetchone()
            if result is None:
                raise RuntimeError("receipt/result invariant violated")
            return {"id":raw["id"],"score":result[0],"model":result[1]}
        connection.execute("INSERT INTO receipts VALUES(?,?)",(raw["id"],payload))
        connection.execute("INSERT INTO results VALUES(?,?,?)",(raw["id"],probability,raw["model"]))
    return {"id":raw["id"],"score":probability,"model":raw["model"]}


def worker(producer_path, consumer_path, artifact, crash_after_consumer=False):
    if Path(producer_path).resolve() == Path(consumer_path).resolve():
        raise ValueError("use separate producer and consumer files")
    producer = connect_store(producer_path,"producer")
    consumer = None
    try:
        consumer = connect_store(consumer_path,"consumer")
        row = producer.execute("""SELECT j.id,j.payload,o.digest FROM jobs j
          JOIN outbox o ON j.id=o.id WHERE o.ack=0 ORDER BY j.id LIMIT 1""").fetchone()
        if row is None: return None
        job_id,payload,digest = row
        if hashlib.sha256(payload.encode()).hexdigest() != digest:
            raise ValueError("outbox integrity mismatch")
        result = consume(consumer,payload,artifact)
        if crash_after_consumer:
            os._exit(23)  # explicit disposable child-process test only
        with transaction(producer):
            changed = producer.execute("UPDATE outbox SET ack=1 WHERE id=? AND digest=?",
                                       (job_id,digest)).rowcount
            if changed != 1: raise RuntimeError("ack identity mismatch")
        return result
    finally:
        if consumer is not None: consumer.close()
        producer.close()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command",required=True)
    demo = sub.add_parser("demo")
    demo.add_argument("directory")
    crash = sub.add_parser("crashworker")
    crash.add_argument("producer");crash.add_argument("consumer");crash.add_argument("artifact")
    args = parser.parse_args()
    if args.command == "crashworker":
        worker(args.producer,args.consumer,load_artifact(args.artifact),crash_after_consumer=True)
        return
    root = Path(args.directory)
    root.mkdir(parents=True,exist_ok=False)
    artifact,report = experiment()
    (root/"model.json").write_text(canonical(artifact),encoding="utf-8")
    (root/"experiment.json").write_text(json.dumps(report,indent=2),encoding="utf-8")
    producer = connect_store(root/"producer.db","producer")
    try: submit(producer,"example-1",[.3,-.2],fingerprint(artifact))
    finally: producer.close()
    print(worker(root/"producer.db",root/"consumer.db",artifact))
    print("Saved local experiment to",root)


if __name__ == "__main__":
    main()