Spaces:
Sleeping
Sleeping
File size: 7,987 Bytes
a42d9ad | 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 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """Bounded HTTP job service for local checkpoint manifests."""
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
import fcntl
import json
import logging
from pathlib import Path
import secrets
from threading import BoundedSemaphore
from uuid import UUID, uuid4
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import FileResponse
from fastapi.security import APIKeyHeader, HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel, ConfigDict, Field
from .core import _kernel
from .files import prune_manifest
class JobRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
manifest: str = Field(min_length=1, max_length=1024)
seed: int = Field(default=0, ge=0, lt=1 << 64)
class CheckpointJobRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
checkpoint: str = Field(default=".", min_length=1, max_length=1024)
keep: int = Field(strict=True, ge=0, lt=1 << 63)
moments_manifest: str | None = Field(default=None, min_length=1, max_length=1024)
seed: int = Field(default=0, ge=0, lt=1 << 64)
def _write_state(directory, run_id, state, **fields):
path = directory / "state.json"
tmp = directory / ".state.json.tmp"
tmp.write_text(json.dumps(dict(state=state, run_id=run_id, **fields)))
tmp.replace(path)
def create_app(input_root, job_root, *, workers=1, queue_size=2, threads_per_job=1,
max_weights=100_000_000, token=None, require_token=False,
checkpoint_root=None):
"""Create a single-process HTTP service with configurable CPU and queue limits."""
if min(workers, threads_per_job, max_weights) < 1 or queue_size < 0:
raise ValueError("invalid worker, queue, thread, or weight limit")
inputs = Path(input_root).resolve(strict=True)
jobs = Path(job_root).resolve()
checkpoints = Path(checkpoint_root).resolve(strict=True) if checkpoint_root is not None else None
if jobs.is_relative_to(inputs) or inputs.is_relative_to(jobs):
raise ValueError("input and job roots must be separate directories")
@asynccontextmanager
async def lifespan(app):
jobs.mkdir(parents=True, exist_ok=True)
with (jobs / ".server.lock").open("a+b") as lock:
try:
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError as exc:
raise RuntimeError("another service already owns this job directory") from exc
_kernel()
app.state.run_id = uuid4().hex
app.state.slots = BoundedSemaphore(workers + queue_size)
with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="prune") as pool:
app.state.pool = pool
yield
app = FastAPI(title="Tensor Pruning", version="1.0.0", lifespan=lifespan)
bearer = HTTPBearer(auto_error=False)
api_key = APIKeyHeader(name="X-Tensor-Pruning-Token", auto_error=False)
def authenticate(credentials: HTTPAuthorizationCredentials | None = Depends(bearer),
space_token: str | None = Depends(api_key)):
if not token:
raise HTTPException(503, "The operator must configure the API token")
supplied = space_token if space_token is not None else (credentials.credentials if credentials else "")
if not secrets.compare_digest(supplied.encode(), token.encode()):
raise HTTPException(401, "A valid bearer token is required",
headers={"WWW-Authenticate": "Bearer"})
protected = [Depends(authenticate)] if token or require_token else []
@app.get("/health")
def health():
return {"status": "ok"}
def read_job(job_id, request):
directory = jobs / str(job_id)
try:
state = json.loads((directory / "state.json").read_text())
except FileNotFoundError as exc:
raise HTTPException(404, "Unknown job") from exc
run_id = state.pop("run_id")
if run_id != request.app.state.run_id and state["state"] in {"queued", "running"}:
state["state"] = "interrupted"
return directory, state
def enqueue(request, operation):
state = request.app.state
if not state.slots.acquire(blocking=False):
raise HTTPException(429, "Pruning queue is full", headers={"Retry-After": "1"})
job_id = uuid4()
directory = jobs / str(job_id)
def run():
try:
_write_state(directory, state.run_id, "running")
result = operation(directory)
_write_state(directory, state.run_id, "completed", result=result)
except (ValueError, TypeError, OSError, MemoryError) as exc:
_write_state(directory, state.run_id, "failed", error=str(exc))
except Exception:
logging.getLogger(__name__).exception("Pruning job failed")
_write_state(directory, state.run_id, "failed", error="Internal pruning failure")
finally:
state.slots.release()
try:
directory.mkdir()
_write_state(directory, state.run_id, "queued")
state.pool.submit(run)
except Exception:
state.slots.release()
raise
return {"id": str(job_id), "state": "queued"}
@app.post("/v1/jobs", status_code=202, dependencies=protected)
def submit(body: JobRequest, request: Request):
manifest = Path(body.manifest)
if manifest.is_absolute() or not (inputs / manifest).resolve().is_relative_to(inputs):
raise HTTPException(400, "Manifest must be inside the input root")
return enqueue(request, lambda directory: prune_manifest(
inputs, body.manifest, directory / "output", threads=threads_per_job,
seed=body.seed, max_weights=max_weights))
if checkpoints is not None:
from .checkpoint import prune_checkpoint
@app.post("/v1/checkpoint-jobs", status_code=202, dependencies=protected)
def submit_checkpoint(body: CheckpointJobRequest, request: Request):
relative = Path(body.checkpoint)
checkpoint = (checkpoints / relative).resolve()
if relative.is_absolute() or not checkpoint.is_relative_to(checkpoints):
raise HTTPException(400, "Checkpoint must be inside the model mount")
if not checkpoint.is_dir():
raise HTTPException(404, "Mounted checkpoint directory not found")
return enqueue(request, lambda directory: prune_checkpoint(
checkpoint, directory / "output", keep=body.keep, workspace=directory,
max_weights=max_weights, input_root=inputs,
moments_manifest=body.moments_manifest, threads=threads_per_job, seed=body.seed))
@app.get("/v1/jobs/{job_id}", dependencies=protected)
def status(job_id: UUID, request: Request):
_, state = read_job(job_id, request)
return {"id": str(job_id), **state}
@app.get("/v1/jobs/{job_id}/masks/{index}", dependencies=protected,
response_class=FileResponse, responses={
200: {"content": {"application/octet-stream": {
"schema": {"type": "string", "format": "binary"}
}}}
})
def mask(job_id: UUID, index: int, request: Request):
directory, state = read_job(job_id, request)
if state["state"] != "completed":
raise HTTPException(409, "The job has no completed masks")
projections = state["result"]["projections"]
if not 0 <= index < len(projections):
raise HTTPException(404, "Unknown projection index")
filename = projections[index]["mask"]
return FileResponse(directory / "output" / filename,
filename=filename, media_type="application/octet-stream")
return app
|