Spaces:
Sleeping
Sleeping
| """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") | |
| 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 [] | |
| 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"} | |
| 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 | |
| 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)) | |
| def status(job_id: UUID, request: Request): | |
| _, state = read_job(job_id, request) | |
| return {"id": str(job_id), **state} | |
| 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 | |