File size: 5,152 Bytes
3145545 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """Frugal PCD validation/benchmarking and optional authenticated serving on Modal.
modal run lfm25_pcd_modal.py --task prepare
modal run lfm25_pcd_modal.py --task validate
modal run lfm25_pcd_modal.py --task benchmark
Publication is a separate, explicit --task publish operation.
"""
import json
from pathlib import Path
import modal
app = modal.App("lfm25-26b-pcd")
root = Path(__file__).parent
cache = modal.Volume.from_name("huggingface-cache", create_if_missing=True)
results = modal.Volume.from_name("lfm25-26b-pcd-results", create_if_missing=True)
base = modal.Image.debian_slim(python_version="3.11").env(
{
"HF_HOME": "/root/.cache/huggingface",
"TOKENIZERS_PARALLELISM": "false",
"HF_HUB_DISABLE_TELEMETRY": "1",
"OMP_NUM_THREADS": "4",
}
)
cpu_image = base.uv_pip_install("huggingface-hub==1.31.0", "pyyaml==6.0.3").add_local_python_source(
"pcd"
)
gpu_image = (
base.pip_install_from_requirements(str(root / "requirements-pcd.txt"))
.add_local_python_source("pcd")
.add_local_dir(root / "tests" / "pcd", "/root/tests/pcd")
)
@app.function(
image=cpu_image,
secrets=[modal.Secret.from_name("huggingface")],
volumes={"/root/.cache/huggingface": cache},
timeout=900,
memory=4096,
max_containers=1,
)
def prepare():
from pcd.publishing import prepare as prepare_snapshot
info = prepare_snapshot()
cache.commit()
return info
@app.cls(
image=gpu_image,
gpu="L40S",
cpu=4,
memory=16384,
timeout=900,
min_containers=0,
max_containers=1,
scaledown_window=30,
volumes={"/root/.cache/huggingface": cache, "/results": results},
)
class PCDModel:
precision: str = modal.parameter(default="float16")
@modal.enter()
def load(self):
import torch
from pcd.engine import Engine
if self.precision not in {"float16", "bfloat16", "float32"}:
raise ValueError("unsupported precision")
torch.set_num_threads(4)
torch.set_float32_matmul_precision("highest")
self.engine = Engine(
device="cuda", dtype=self.precision, attention="sdpa", local_files_only=True
)
@modal.method()
def run(self, task: str, suite: str = "diagnostic", repeats: int = 3):
from pcd.benchmark import execute
report = execute(self.engine, task=task, suite=suite, repeats=repeats)
path = Path("/results") / (report["run_id"] + ".json")
path.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")
results.commit()
return report
@modal.method()
def infer(self, context: str, schema: dict, mode: str = "token"):
return self.engine.constrained(context, schema, mode=mode)
@modal.fastapi_endpoint(method="POST", requires_proxy_auth=True, docs=False)
def extract(self, payload: dict):
import jsonschema
from fastapi import HTTPException
try:
if set(payload) - {"context", "schema", "mode"}:
raise ValueError("unsupported request keys")
return self.engine.constrained(
payload["context"], payload["schema"], mode=payload.get("mode", "token")
)
except (ValueError, KeyError, TypeError, jsonschema.SchemaError) as exc:
raise HTTPException(status_code=422, detail=str(exc)) from None
release_dir = root / "release" / "pcd"
publish_image = (
cpu_image.add_local_dir(release_dir, "/release-src") if release_dir.is_dir() else cpu_image
)
@app.function(
image=publish_image,
secrets=[modal.Secret.from_name("huggingface")],
volumes={"/root/.cache/huggingface": cache, "/results": results},
timeout=1800,
cpu=4,
memory=8192,
max_containers=1,
)
def publish(public: bool = False):
from pcd.publishing import publish as publish_bundle
result = publish_bundle(public=public)
results.commit()
return result
@app.local_entrypoint()
def main(
task: str = "validate",
suite: str = "diagnostic",
repeats: int = 3,
precision: str = "float16",
public: bool = False,
):
if task not in {"prepare", "validate", "benchmark", "probe", "publish"}:
raise ValueError("task must be prepare, validate, benchmark, probe, or publish")
if repeats < 1 or repeats > 10:
raise ValueError("repeats must be between 1 and 10")
if public and task != "publish":
raise ValueError("--public only applies to --task publish")
if task == "publish":
if not release_dir.is_dir():
raise ValueError("stage the release with scripts/build_pcd_release.py first")
report = publish.remote(public)
elif task == "prepare":
report = prepare.remote()
else:
report = PCDModel(precision=precision).run.remote(task, suite, repeats)
output = root / "results" / "pcd"
output.mkdir(parents=True, exist_ok=True)
path = output / (report.get("run_id", task) + ".json")
path.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")
print(json.dumps({"report": str(path), "summary": report.get("summary", report)}, indent=2))
|