add lifetime option for server and dht

scripts/run_dht.py

@@ -1,9 +1,4 @@
 import argparse
-import resource
-import os
-import sys
-
-import torch
 import tesseract
 from tesseract.utils import find_open_port
 
@@ -12,6 +7,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--port', type=int, default=None, required=False)
     parser.add_argument('--initial_peers', type=str, default="[]", required=False)
+    parser.add_argument('--lifetime_seconds', type=int, default=None, required=False)
 
     args = parser.parse_args()
     initial_peers = eval(args.initial_peers)
@@ -21,6 +17,7 @@ if __name__ == "__main__":
     print(f"Running network node on port {network.port}")
 
     try:
-        network.run()
+        network.start()
+        network.join(timeout=args.lifetime_seconds)
     finally:
         network.shutdown()

scripts/run_server.py

@@ -25,6 +25,7 @@ if __name__ == "__main__":
     parser.add_argument('--no_network', action='store_true')
     parser.add_argument('--initial_peers', type=str, default="[]", required=False)
     parser.add_argument('--network_port', type=int, default=None, required=False)
+    parser.add_argument('--lifetime_seconds', type=int, default=None, required=False)
     parser.add_argument('--increase_file_limit', action='store_true')
 
     args = parser.parse_args()
@@ -65,9 +66,10 @@ if __name__ == "__main__":
     server = tesseract.TesseractServer(
         network, experts, addr=args.host, port=args.port or find_open_port(),
         conn_handler_processes=args.num_handlers, device=device)
-    print(f"Running server at {server.addr}:{server.port}")
-    print(f"Active experts of type {args.expert_cls}: {list(experts.keys())}")
     try:
-        server.run()
+        server.run_in_background(await_ready=True)
+        print(f"Running server at {server.addr}:{server.port}")
+        print(f"Active experts of type {args.expert_cls}: {list(experts.keys())}")
+        server.join(timeout=args.lifetime_seconds)
     finally:
         server.shutdown()