Artifex/control_plane/model_studio/profiles.py
2026-08-17 02:04:43 +07:00

165 lines
10 KiB
Python

from __future__ import annotations
import hashlib
import json
import shlex
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol
@dataclass
class ArchaeologyFinding:
subject: str
confidence: str
evidence: list[str]
details: dict[str, Any]
class ModelProjectProfile(Protocol):
name: str
def archaeology(self, repository_path: str) -> dict[str, Any]: ...
def validate_evaluation_suite(self, repository_path: str, suite_reference: str) -> dict[str, Any]: ...
def training_command(self, recipe: dict[str, Any], output_directory: str) -> list[str]: ...
class GuardModelProfile:
name = "guard"
base_model = "Qwen/Qwen2.5-Coder-3B-Instruct"
def archaeology(self, repository_path: str) -> dict[str, Any]:
root = Path(repository_path)
findings: list[dict[str, Any]] = []
def add(subject: str, confidence: str, evidence: list[Path], **details: Any) -> None:
findings.append({"subject": subject, "confidence": confidence, "evidence": [str(item) for item in evidence], "details": details})
trainer = root / "scripts" / "05_finetune.py"
evaluator = root / "scripts" / "run_evmbench.py"
strict_scorer = root / "scripts" / "score_evmbench_strict.py"
add("base_model", "CONFIRMED" if trainer.exists() else "UNKNOWN", [trainer] if trainer.exists() else [], value=self.base_model)
add("training_entrypoint", "CONFIRMED" if trainer.exists() else "UNKNOWN", [trainer] if trainer.exists() else [], command="python scripts/05_finetune.py --model ... --training-data ... --output ...")
add("evaluation_entrypoint", "CONFIRMED" if evaluator.exists() else "UNKNOWN", [item for item in [evaluator, strict_scorer] if item.exists()], command="python scripts/run_evmbench.py ...", strict_scoring=bool(strict_scorer.exists()))
checkpoints = []
for adapter in sorted((root / "models").glob("**/adapter_model.*")) if (root / "models").exists() else []:
checkpoints.append({"name": adapter.parent.name, "reference": str(adapter.parent), "hash": self._hash_file(adapter), "adapter_type": "LORA", "confidence": "CONFIRMED"})
datasets = []
for manifest in sorted((root / "data" / "production_authorized").glob("*.json")) if (root / "data" / "production_authorized").exists() else []:
try:
payload = json.loads(manifest.read_text(encoding="utf-8"))
count = len(payload) if isinstance(payload, list) else payload.get("record_count")
except (OSError, json.JSONDecodeError):
count = None
datasets.append({"name": manifest.stem, "reference": str(manifest), "hash": self._hash_file(manifest), "record_count": count, "tags": ["imported", "training"]})
reports = [str(path) for path in sorted((root / "results").glob("qwen25_coder_3b_*") if (root / "results").exists() else [])]
benchmark = root / "evmbench"
add("benchmark_checkout", "CONFIRMED" if benchmark.exists() else "UNKNOWN", [benchmark] if benchmark.exists() else [], strict_scorer_exists=strict_scorer.exists())
return {"root_exists": root.exists(), "findings": findings, "checkpoints": checkpoints, "datasets": datasets, "historical_reports": reports, "benchmark_reference": str(benchmark), "trainer_reference": str(trainer), "evaluator_reference": str(evaluator), "strict_scorer_reference": str(strict_scorer)}
def validate_evaluation_suite(self, repository_path: str, suite_reference: str) -> dict[str, Any]:
root = Path(repository_path)
evaluator = root / "scripts" / "run_evmbench.py"
scorer = root / "scripts" / "score_evmbench_strict.py"
benchmark = Path(suite_reference)
available = evaluator.exists() and scorer.exists() and benchmark.exists()
return {"status": "WARNING" if available else "BLOCKED", "evidence": {"evaluator": str(evaluator), "strict_scorer": str(scorer), "benchmark": str(benchmark), "reason": "Scripts and checkout are present, but no fresh Spark baseline has yet proven model loading, output parsing, reproducibility, or contamination checks." if available else "Guard evaluation harness or strict scorer missing."}}
def training_command(self, recipe: dict[str, Any], output_directory: str) -> list[str]:
dataset = str(recipe["training_data"])
return ["python", "scripts/05_finetune.py", "--model", str(recipe.get("base_model", self.base_model)), "--training-data", dataset, "--output", output_directory]
@staticmethod
def _hash_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for block in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
class SparkGuardDatasetInventory:
"""Read-only manifest metadata probe for a registered Spark Guard workspace."""
def __init__(self, ssh_alias: str = "spark") -> None:
self.ssh_alias = ssh_alias
def inspect(self, references: list[str]) -> list[dict[str, Any]]:
script = """import hashlib,json,os,sys
keys=['dataset_status','final_model_eligible','production_training_authorized','benchmark_source_included','synthetic_code_included','split','c4_invalid','teacher','vulnerability_type']
for path in sys.argv[1:]:
try:
raw=open(path,'rb').read(); value=json.loads(raw); first=value[0] if isinstance(value,list) and value else {}
print(json.dumps({'reference':path,'content_hash':hashlib.sha256(raw).hexdigest(),'record_count':len(value) if isinstance(value,list) else None,'first':{key:first.get(key) for key in keys if key in first}},ensure_ascii=True))
except Exception as exc: print(json.dumps({'reference':path,'error':str(exc)},ensure_ascii=True))
"""
if not references:
return []
output = self._remote("python3 -c " + shlex.quote(script) + " " + " ".join(self._quote(item) for item in references), timeout=120)
records = []
for line in output.splitlines():
try:
item = json.loads(line)
except json.JSONDecodeError:
continue
if "content_hash" in item and item.get("record_count") is not None:
records.append(item)
return records
def _remote(self, command: str, *, timeout: int = 120) -> str:
completed = subprocess.run(["ssh", self.ssh_alias, command], capture_output=True, text=True, encoding="utf-8", errors="replace", timeout=timeout, check=False)
if completed.returncode:
raise RuntimeError(completed.stderr.strip() or f"Spark inventory command failed: {command}")
return completed.stdout.strip()
@staticmethod
def _quote(value: str) -> str:
return "'" + value.replace("'", "'\\''") + "'"
class SparkGuardDatasetMaterializer(SparkGuardDatasetInventory):
"""Creates new versioned JSON manifests from authorized Guard source records on Spark."""
def materialize(self, source_references: list[str], output_directory: str, *, strict: bool = False) -> dict[str, Any]:
script = """import hashlib,json,os,sys
out=sys.argv[1]; strict=sys.argv[2]=='1'; sources=sys.argv[3:]; os.makedirs(out,exist_ok=True)
train=[]; validation=[]; regression=[]; seen=set(); stats={'sources':len(sources),'read':0,'accepted':0,'malformed':0,'duplicates':0}
for source in sources:
value=json.load(open(source,encoding='utf-8'))
if not isinstance(value,list): continue
for record in value:
stats['read']+=1
if not isinstance(record,dict) or not isinstance(record.get('input'),str) or not record['input'].strip() or not isinstance(record.get('output'),str) or not record['output'].strip(): stats['malformed']+=1; continue
if strict:
try: target=json.loads(record['output'])
except Exception: stats['malformed']+=1; continue
if not isinstance(target,dict) or not isinstance(target.get('findings'),list) or record.get('c4_invalid') is True or 'pragma solidity' not in record['input'].lower(): stats['malformed']+=1; continue
normalized=[]
for finding in target['findings']:
if not isinstance(finding,dict): continue
kind=finding.get('canonical_type') or finding.get('type')
if kind: normalized.append({key:finding[key] for key in ['canonical_type','type','severity','location','description','evidence','vulnerable_line_spans'] if key in finding})
record={**record,'output':json.dumps({'findings':normalized},ensure_ascii=False),'artifex_schema':'guard_finding_v1'}
source_key=str(record.get('source_sha256') or record.get('parent_source_sha256') or hashlib.sha256(record['input'].encode()).hexdigest())
key=hashlib.sha256((record['input']+'\\0'+record['output']).encode()).hexdigest()
if key in seen: stats['duplicates']+=1; continue
seen.add(key); record={**record,'artifex_source_key':source_key,'artifex_record_hash':key,'artifex_source_manifest':source}; bucket=int(hashlib.sha256(source_key.encode()).hexdigest(),16)%10
(train if bucket<8 else validation if bucket<9 else regression).append(record); stats['accepted']+=1
for name,rows in [('train',train),('validation',validation),('regression',regression)]:
path=os.path.join(out,name+'.json'); json.dump(rows,open(path,'w',encoding='utf-8'),ensure_ascii=False); stats[name]=len(rows); stats[name+'_sha256']=hashlib.sha256(open(path,'rb').read()).hexdigest()
json.dump(stats,open(os.path.join(out,'manifest.json'),'w',encoding='utf-8'),indent=2); print(json.dumps({'output_directory':out,**stats}))
"""
if not source_references:
raise ValueError("Dataset materialization requires source manifests.")
output = self._remote("python3 -c " + shlex.quote(script) + " " + self._quote(output_directory) + " " + ("1" if strict else "0") + " " + " ".join(self._quote(item) for item in source_references), timeout=900)
return json.loads(output)
def sample(self, reference: str, *, count: int = 8, maximum_field_chars: int = 3000) -> list[dict[str, Any]]:
script = """import json,sys
rows=json.load(open(sys.argv[1],encoding='utf-8')); count=int(sys.argv[2]); limit=int(sys.argv[3]); step=max(1,len(rows)//max(1,count)); out=[]
for index in range(0,len(rows),step):
record=rows[index]; out.append({key:(value[:limit] if isinstance(value,str) else value) for key,value in record.items() if key in ['input','output','record_id','source_sha256','parent_source_sha256','vulnerability_type','artifex_source_manifest']})
if len(out)>=count: break
print(json.dumps(out,ensure_ascii=True))
"""
output = self._remote("python3 -c " + shlex.quote(script) + " " + self._quote(reference) + f" {count} {maximum_field_chars}", timeout=120)
return json.loads(output)