Spaces:
Sleeping
Sleeping
| """Command-line access to the file API and HTTP service.""" | |
| import argparse | |
| import json | |
| import os | |
| from pathlib import Path | |
| from .files import prune_manifest | |
| def main(argv=None): | |
| parser = argparse.ArgumentParser(prog="tensor-prune", description=__doc__) | |
| commands = parser.add_subparsers(dest="command", required=True) | |
| prune = commands.add_parser("prune", help="prune a memory-mapped manifest") | |
| prune.add_argument("--input-root", required=True, type=Path) | |
| prune.add_argument("--manifest", required=True) | |
| prune.add_argument("--output", required=True, type=Path) | |
| prune.add_argument("--scratch-dir", type=Path) | |
| prune.add_argument("--threads", type=int, default=1) | |
| prune.add_argument("--seed", type=int, default=0) | |
| serve = commands.add_parser("serve", help="run the bounded HTTP job service") | |
| serve.add_argument("--input-root", required=True, type=Path) | |
| serve.add_argument("--job-root", required=True, type=Path) | |
| serve.add_argument("--host", default="127.0.0.1") | |
| serve.add_argument("--port", type=int, default=8080) | |
| serve.add_argument("--workers", type=int, default=1, help="simultaneous jobs in this process") | |
| serve.add_argument("--queue-size", type=int, default=2) | |
| serve.add_argument("--threads-per-job", type=int, default=1) | |
| serve.add_argument("--max-weights", type=int, default=100_000_000) | |
| args = parser.parse_args(argv) | |
| if args.command == "prune": | |
| result = prune_manifest(args.input_root, args.manifest, args.output, | |
| scratch_dir=args.scratch_dir, threads=args.threads, seed=args.seed) | |
| print(json.dumps(result)) | |
| return | |
| try: | |
| import uvicorn | |
| from .service import create_app | |
| except ImportError: | |
| parser.error("install the HTTP extras with pip install '.[http]'") | |
| token = os.environ.get("TENSOR_PRUNING_API_TOKEN") | |
| if args.host not in {"127.0.0.1", "::1", "localhost"} and not token: | |
| parser.error("set TENSOR_PRUNING_API_TOKEN before binding to a non-loopback address") | |
| app = create_app(args.input_root, args.job_root, workers=args.workers, | |
| queue_size=args.queue_size, threads_per_job=args.threads_per_job, | |
| max_weights=args.max_weights, token=token) | |
| uvicorn.run(app, host=args.host, port=args.port) | |
| if __name__ == "__main__": | |
| main() | |