Spaces:
Sleeping
Sleeping
File size: 2,373 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 | """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()
|