diff --git a/python/loom/client/client.py b/python/loom/client/client.py index 6aad0b225fc1e58c15cfe11233ed1052d5dff782..82b99a50e6c5b0cb0a77f4768294bd5b1a8edf97 100644 --- a/python/loom/client/client.py +++ b/python/loom/client/client.py @@ -159,8 +159,8 @@ class Client(object): self.symbols = {} for i, s in enumerate(cmsg.symbols): self.symbols[s] = i - self.array_id = self.symbols["loom/array"] - self.rawdata_id = self.symbols["loom/data"] + self.array_id = self.symbols.get("loom/array") + self.rawdata_id = self.symbols.get("loom/data") def process_error(self, cmsg): assert cmsg.HasField("error") diff --git a/python/loom/client/wait/__init__.py b/python/loom/client/wait/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/python/loom/client/wait/__main__.py b/python/loom/client/wait/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..0242906e1720b4735a36f01fc3df14f39ab5cb04 --- /dev/null +++ b/python/loom/client/wait/__main__.py @@ -0,0 +1,63 @@ + +# +# This program waits for X workers in server +# + +import argparse +import sys +import time +from ..client import Client + +def parse_args(): + parser = argparse.ArgumentParser( + description="Wait " + "-- waits for workers in server") + + parser.add_argument("server", + metavar="SERVER", + type=str, + help="Server hostname") + + parser.add_argument("port", + metavar="PORT", + type=int, + help="Server TCP port") + + parser.add_argument("n_workers", + metavar="N_WORKERS", + type=int, + help="Number of workers") + + parser.add_argument("timeout", + metavar="TIMEOUT", + type=int, + help="timeout (in seconds)") + + args = parser.parse_args() + return args + + + +def main(): + args = parse_args() + client = Client(args.server, args.port) + + t = 0 + workers = -1 + + while True: + if t > args.timeout: + print("Timeouted") + sys.exit(1) + n_workers = client.get_stats()["n_workers"] + if workers != n_workers: + workers = n_workers + print("Workers in server: {}/{}".format(workers, args.n_workers)) + if workers == args.n_workers: + return + time.sleep(1) + t += 1 + + +if __name__ == "__main__": + main()