File size: 1,344 Bytes
5953ef9 |
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 |
import argparse
import random
import string
import sys
from watchmen import WatchClient
def parse_args(in_args=None):
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument("--task_name", type=str, required=True, help="Take Name")
arg_parser.add_argument("--cuda", type=str, required=True, help="cuda to be waited")
arg_parser.add_argument(
"--req_gpu_num",
type=int,
required=False,
default=1,
help="request number of gpus",
)
arg_parser.add_argument(
"--wait",
choices=["schedule", "queue", "none"],
default="none",
help="scheduling/queue wait",
)
arg_info = arg_parser.parse_args(args=in_args)
return arg_info
if __name__ == "__main__":
in_argv = parse_args()
if in_argv.wait == "none":
sys.exit(0)
random_id = "-" + "".join(random.sample(string.ascii_letters + string.digits, 8))
exp_id = in_argv.task_name + random_id
watch_client = WatchClient(
id=exp_id,
gpus=eval(f"[{in_argv.cuda}]"),
server_host="localhost",
server_port=62333,
req_gpu_num=in_argv.req_gpu_num,
mode=in_argv.wait,
timeout=60,
)
available_gpus = watch_client.wait()
available_gpus = [str(x) for x in available_gpus]
print(",".join(available_gpus))
|