Spaces:
Sleeping
Sleeping
| """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"), | |
| ) | |
| def index(): | |
| return RedirectResponse("/docs") | |