Blame for space_app.py

"""Configure the Tensor Pruning API for a single Docker Space replica."""

import os
from pathlib import Path

from fastapi.responses import RedirectResponse

from tensor_pruning.service import create_app


work = Path(os.getenv("PRUNING_WORK_ROOT", "/home/user/work"))
inputs = Path(os.getenv("PRUNING_INPUT_ROOT", str(work / "inputs")))
inputs.mkdir(parents=True, exist_ok=True)

app = create_app(
    inputs,
    work / "jobs",
    workers=1,
    queue_size=int(os.getenv("PRUNING_QUEUE_SIZE", "1")),
    threads_per_job=int(os.getenv("PRUNING_THREADS", str(min(2, len(os.sched_getaffinity(0)))))),
    max_weights=int(os.getenv("PRUNING_MAX_WEIGHTS", "6000000000")),
    token=os.getenv("TENSOR_PRUNING_API_TOKEN"),
    require_token=True,
    checkpoint_root=os.getenv("PRUNING_CHECKPOINT_ROOT", "/models"),
)


@app.get("/", include_in_schema=False)
def index():
    return RedirectResponse("/docs")