feat: split __main__.py from ps.py

Author: specture724

checkpoint_engine/__main__.py

@@ -0,0 +1,28 @@
+import argparse
+import os
+
+from loguru import logger
+
+from checkpoint_engine.api import _init_api
+from checkpoint_engine.ps import ParameterServer
+
+
+@logger.catch(reraise=True)
+def run_from_cli():
+    import uvicorn
+
+    parser = argparse.ArgumentParser(description="Parameter Server")
+    parser.add_argument("--uds", type=str)
+
+    args = parser.parse_args()
+    logger.info(
+        f"Parameter Server {args=}, master addr: {os.getenv('MASTER_ADDR')}, master port {os.getenv('MASTER_PORT')}"
+    )
+
+    assert args.uds and len(args.uds) > 0, args.uds
+    ps = ParameterServer(auto_pg=True)
+    uvicorn.run(_init_api(ps), uds=args.uds, timeout_keep_alive=60)
+
+
+if __name__ == "__main__":
+    run_from_cli()

checkpoint_engine/ps.py

@@ -1,4 +1,3 @@
-import argparse
 import os
 import threading
 import time
@@ -13,7 +12,6 @@
 from loguru import logger
 from torch.multiprocessing.reductions import reduce_tensor
 
-from checkpoint_engine.api import _init_api
 from checkpoint_engine.data_types import (
     BucketRange,
     DataToGather,
@@ -822,24 +820,3 @@ def _update_per_bucket(
                 self._p2p_store.unregister_named_tensors([h2d_buffer_name])
 
             self.device_manager.device_module.empty_cache()
-
-
-@logger.catch(reraise=True)
-def run_from_cli():
-    import uvicorn
-
-    parser = argparse.ArgumentParser(description="Parameter Server")
-    parser.add_argument("--uds", type=str)
-
-    args = parser.parse_args()
-    logger.info(
-        f"Parameter Server {args=}, master addr: {os.getenv('MASTER_ADDR')}, master port {os.getenv('MASTER_PORT')}"
-    )
-
-    assert args.uds and len(args.uds) > 0, args.uds
-    ps = ParameterServer(auto_pg=True)
-    uvicorn.run(_init_api(ps), uds=args.uds, timeout_keep_alive=60)
-
-
-if __name__ == "__main__":
-    run_from_cli()