Leave --no_auto_relay argument only

bigscience-workshop · Jan 9, 2023 · a39e544 · a39e544
1 parent a8f20de
commit a39e544
Showing 1 changed file with 4 additions and 12 deletions.
diff --git a/src/petals/cli/run_server.py b/src/petals/cli/run_server.py
@@ -12,9 +12,6 @@
 logger = get_logger(__file__)
 
 
-TRUE_CONSTANTS = ["TRUE", "1"]
-
-
 def main():
     # fmt:off
     parser = configargparse.ArgParser(default_config_files=["config.yml"],
@@ -41,6 +38,9 @@ def main():
                              'This is a simplified way to set the --announce_maddrs option (see below).'
                              'Default: server announces IPv4/IPv6 addresses of your network interfaces')
 
+    parser.add_argument("--no_auto_relay", action="store_false", dest="use_auto_relay",
+                        help="Do not look for libp2p relays to reach peers behind NATs/firewalls")
+
     parser.add_argument('--host_maddrs', nargs='+', required=False,
                         help='Multiaddrs to listen for external connections from other peers')
     parser.add_argument('--announce_maddrs', nargs='+', required=False,
@@ -130,11 +130,7 @@ def main():
     parser.add_argument("--mean_balance_check_period", type=float, default=60,
                         help="Check the swarm's balance every N seconds (and rebalance it if necessary)")
 
-    parser.add_argument("--use_auto_relay", type=str, default="True",
-                        help="Look for libp2p relays for NAT traversal. "
-                             "Use `--use_auto_relay False/True` to disable/enable this")
     parser.add_argument("--use_auth_token", action='store_true', help="auth token for from_pretrained")
-
     parser.add_argument('--load_in_8bit', type=str, default=None,
                         help="Convert the loaded transformer blocks into mixed-8bit quantized model. "
                              "Default: True if GPU is available. Use `--load_in_8bit False` to disable this")
@@ -193,11 +189,9 @@ def main():
     if args.pop("new_swarm"):
         args["initial_peers"] = []
 
-    use_auto_relay = args.pop("use_auto_relay").upper() in TRUE_CONSTANTS
-
     load_in_8bit = args.pop("load_in_8bit")
     if load_in_8bit is not None:
-        load_in_8bit = load_in_8bit.upper() in TRUE_CONSTANTS
+        args["load_in_8bit"] = load_in_8bit.lower() in ["true", "1"]
 
     server = Server(
         **args,
@@ -206,8 +200,6 @@ def main():
         compression=compression,
         max_disk_space=max_disk_space,
         attn_cache_size=attn_cache_size,
-        use_auto_relay=use_auto_relay,
-        load_in_8bit=load_in_8bit,
     )
     try:
         server.run()