ladyrick 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

ladyrick/allgather.py

@@ -0,0 +1,106 @@
+import os
+import socket
+import struct
+import sys
+import time
+
+
+class AllGather:
+    def __init__(self):
+        if {"MASTER_ADDR", "MASTER_PORT", "WORLD_SIZE", "RANK"} - set(os.environ):
+            raise RuntimeError("MASTER_ADDR, MASTER_PORT, WORLD_SIZE, RANK env is required")
+        self.master_addr = os.environ["MASTER_ADDR"]
+        self.master_port = int(os.environ["MASTER_PORT"])
+        self.world_size = int(os.environ["WORLD_SIZE"])
+        self.rank = int(os.environ["RANK"])
+        assert 0 <= self.rank < self.world_size
+
+    def allgather(self, data: bytes):
+        if self.world_size == 1:
+            assert self.rank == 0
+            return [data]
+
+        if self.rank == 0:
+            gathered = self._run_master(data)
+        else:
+            gathered = self._run_worker(data)
+        return gathered
+
+    def _run_master(self, data: bytes):
+        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        try:
+            server_socket.bind(("0.0.0.0", self.master_port))
+            server_socket.listen(self.world_size - 1)
+        except OSError as e:
+            print(f"Master failed to bind port: {e}")
+            sys.exit(1)
+
+        try:
+            # collect
+            all_data = [data]
+            conns: list[socket.socket] = []
+            for worker_rank in range(1, self.world_size):
+                conn, _ = server_socket.accept()
+                conns.append(conn)
+                cur_data_len = struct.unpack("!I", conn.recv(4))[0]
+                cur_data = conn.recv(cur_data_len)
+                all_data.append(cur_data)
+
+            assert len(all_data) == self.world_size
+
+            # broadcast
+            data_len_pack = struct.pack("!" + "I" * len(all_data), *[len(d) for d in all_data])
+
+            for worker_rank in range(1, self.world_size):
+                conn = conns[worker_rank - 1]
+                conn.sendall(data_len_pack)
+                for d in all_data:
+                    conn.sendall(d)
+                conn.close()
+        finally:
+            server_socket.close()
+        return all_data
+
+    def _run_worker(self, data: bytes):
+        worker_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        while True:
+            try:
+                worker_socket.connect((self.master_addr, self.master_port))
+                break
+            except (ConnectionRefusedError, socket.timeout):
+                time.sleep(0.1)
+
+        try:
+            worker_socket.sendall(struct.pack("!I", len(data)))
+            worker_socket.sendall(data)
+
+            all_data: list[bytes] = []
+            data_lens = struct.unpack("!" + "I" * self.world_size, worker_socket.recv(4 * self.world_size))
+            for data_len in data_lens:
+                all_data.append(worker_socket.recv(data_len))
+        finally:
+            worker_socket.close()
+        return all_data
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(prog="allgather")
+    parser.add_argument("msg", type=str, help="msg to allgather")
+    parser.add_argument("-0", action="store_true", help="separate by \\0", dest="zero")
+
+    args = parser.parse_args()
+
+    allgather = AllGather()
+    gathered = allgather.allgather(args.msg.encode())
+    gathered_str = (s.decode() for s in gathered)
+    if args.zero:
+        print("\0".join(gathered_str), flush=True)
+    else:
+        print("\n".join(gathered_str), flush=True)
+
+
+if __name__ == "__main__":
+    main()

ladyrick/cli/multi_ssh.py

@@ -1,17 +1,9 @@
-import argparse
-import dataclasses
-import itertools
 import json
 import os
-import pathlib
-import random
 import select
-import shlex
 import signal
 import subprocess
 import sys
-import time
-import uuid
 
 
 def log(msg):
@@ -37,19 +29,18 @@ def remote_head():
     extra_envs = json.loads(sys.argv[2])
     os.environ.update(extra_envs)
 
-    parser = argparse.ArgumentParser(prog=REMOTE_HEAD_PROG_NAME)
-    parser.add_argument("cmd", type=str, nargs="+", help="cmd")
-    args = parser.parse_args(sys.argv[3:])
+    cmd = sys.argv[3:]
     try:
         from setproctitle import setproctitle
 
-        setproctitle(" ".join([REMOTE_HEAD_PROG_NAME] + args.cmd))
+        setproctitle(" ".join([REMOTE_HEAD_PROG_NAME] + cmd))
     except ImportError:
         pass
 
     # start child process
     child = subprocess.Popen(
-        args.cmd,
+        cmd,
+        stdin=subprocess.PIPE,
         stdout=sys.stdout,
         stderr=sys.stderr,
         start_new_session=True,
@@ -97,6 +88,15 @@ if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] == REMOTE_HEAD_P
 
 
 # ----- remote_head end ----- #
+if True:
+    import argparse
+    import dataclasses
+    import itertools
+    import pathlib
+    import random
+    import shlex
+    import time
+    import uuid
 
 
 @dataclasses.dataclass
@@ -158,18 +158,20 @@ class RemoteExecutor:
     def set_envs(cls, executors: list["RemoteExecutor"]):
         assert executors
         envs = {}
-
-        cmd = cls.make_ssh_cmd(executors[0].host, "hostname -I")
-        master_ips = subprocess.check_output(cmd).decode().split()
-        priority = {"172": 0, "192": 1, "10": 2}
-        master_addr, cur_p = None, -1
-        for ip in master_ips:
-            prefix = ip.split(".", 1)[0]
-            p = priority.get(prefix, 3)
-            if p > cur_p:
-                master_addr, cur_p = ip, p
-        assert master_addr is not None
-        envs["MASTER_ADDR"] = master_addr
+        if len(executors) > 1:
+            cmd = cls.make_ssh_cmd(executors[0].host, "hostname -I")
+            master_ips = subprocess.check_output(cmd).decode().split()
+            priority = {"172": 0, "192": 1, "10": 2}
+            master_addr, cur_p = None, -1
+            for ip in master_ips:
+                prefix = ip.split(".", 1)[0]
+                p = priority.get(prefix, 3)
+                if p > cur_p:
+                    master_addr, cur_p = ip, p
+            assert master_addr is not None
+            envs["MASTER_ADDR"] = master_addr
+        else:
+            envs["MASTER_ADDR"] = "127.0.0.1"
         envs["MASTER_PORT"] = str(random.randint(20000, 40000))
         envs["WORLD_SIZE"] = str(len(executors))
 
@@ -213,10 +215,6 @@ def signal_repeat_checker(sig_to_check, count, duration):
     return checker
 
 
-def get_common_envs():
-    pass
-
-
 def main():
     parser = argparse.ArgumentParser(prog="multi-ssh", add_help=False)
     parser.add_argument("-h", type=str, action="append", help="hosts to connect. order is 1")
@@ -225,13 +223,18 @@ def main():
     parser.add_argument("-l", type=str, help="ssh login User")
     parser.add_argument("-o", type=str, action="append", help="ssh options")
     parser.add_argument("-F", type=str, help="ssh config file")
-    parser.add_argument("cmd", type=str, nargs="+", help="cmd")
     parser.add_argument("--hosts-config", type=str, action="append", help="hosts config string. order is 2")
     parser.add_argument("--hosts-config-file", type=str, action="append", help="hosts config file. order is 3")
     parser.add_argument("--help", action="help", default=argparse.SUPPRESS, help="show this help message and exit")
+    parser.add_argument("cmd", type=str, nargs=argparse.REMAINDER, help="cmd")
 
     args = parser.parse_args()
 
+    if not args.cmd:
+        print("cmd is required\n")
+        parser.print_help()
+        sys.exit(1)
+
     hosts = [
         Host(hn, args.F, args.l, args.p, args.i, args.o)
         for hn in itertools.chain.from_iterable(h.split(",") for h in args.h or [])
@@ -255,6 +258,11 @@ def main():
             )
         )
 
+    if not hosts:
+        print("hosts is required. specify hosts by -h, --hosts-config or --hosts-config-file\n")
+        parser.print_help()
+        sys.exit(1)
+
     executors = [RemoteExecutor(host, args.cmd) for host in hosts]
 
     RemoteExecutor.set_envs(executors)

ladyrick/cli/psf.py

@@ -5,7 +5,7 @@ import sys
 
 def main():
     if os.uname().sysname != "Linux":
-        print("Only support uname Linux")
+        print("only support uname Linux")
         sys.exit(1)
     verbose = ""
     root_pid = 0
@@ -15,7 +15,7 @@ def main():
         elif arg.isdigit():
             root_pid = int(arg)
         else:
-            print(f"Invalid args: {arg}")
+            print(f"invalid args: {arg}")
             sys.exit(1)
     if root_pid in (0, 1):
         cmd = ["ps", verbose + "afxopid,user,cmd"]

{ladyrick-0.5.0.dist-info → ladyrick-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ladyrick
-Version: 0.5.0
+Version: 0.5.2
 Summary: ladyrick's tools
 Author-email: ladyrick <ladyrick@qq.com>
 License-Expression: MIT

{ladyrick-0.5.0.dist-info → ladyrick-0.5.2.dist-info}/RECORD

@@ -1,4 +1,5 @@
 ladyrick/__init__.py,sha256=udDiOLu2L1jIvjgSCLdDsTnS2OA7kdj2BA9kit2maVo,1343
+ladyrick/allgather.py,sha256=jhu6eYABIwdl2WTKbK9UR9ji8imhUIdGd-ITPfluYm0,3531
 ladyrick/debug.py,sha256=jMVdL9cg1yFaX6Yfdw0Yq6LZMOl7av-DnpaR51eEV38,3083
 ladyrick/loader.py,sha256=Ykg4yxK-UlhAGpBJnnHQLecEMkDpQETOenIs2okwdGY,2263
 ladyrick/pickle.py,sha256=AjIm4H2_3bPOW-NPd9ow_sA3Z520LZwf1VeX6742ADs,2147
@@ -8,8 +9,8 @@ ladyrick/torch.py,sha256=CEdHYaOZ00StZettf4MoIB3tMF0fbzSIH-3pOqOMIZM,1977
 ladyrick/typing.py,sha256=YQeApe63dk7yL4NS5ytlR6v3dLCii2-qsXNlUvjK-zw,203
 ladyrick/utils.py,sha256=jRRaqC6kNbCJPGeE0YisFgis-wiuINLik1mcUQtytow,608
 ladyrick/vars.py,sha256=VbFh2u7XybUaBuiYEXBa4sOmoS99vc2AIXdYLBh8vjk,3763
-ladyrick/cli/multi_ssh.py,sha256=T585QpD2aSt-XKELHdiF67ukc-jKf1Q43ipA1jyOZU8,8923
-ladyrick/cli/psf.py,sha256=L37rvhODzcdRsqWM_G-xkem_f783ggiYfc95SHcELp0,1146
+ladyrick/cli/multi_ssh.py,sha256=faRIgradz3e6ZjwpPd77OB_u5FEJFCfVzGq3qrwJsgk,9220
+ladyrick/cli/psf.py,sha256=JLk3gbPn7E3uuPBbzGvLgJmFQlilA6zg_Xlg7xW5jik,1146
 ladyrick/cli/tee.py,sha256=UMJxSJLOEfbV43auVKRTIJ5ZAMAkAfj8byiFLk5PUHE,3579
 ladyrick/cli/test_args.py,sha256=f5sUPDlcf6nbNf6UfLwZQI5g5LN8wlFBQZ10GLw22cg,212
 ladyrick/cli/test_signal.py,sha256=uvVLbHlvpmYgZcfRF9Lcte0OJJrMDDsmId3BpSxeaOA,490
@@ -19,9 +20,9 @@ ladyrick/patch/rich_print.py,sha256=z3Ea1VCunXZvNvEDFHpoyWc8ydINmh-gOIJ1ssscs6s,
 ladyrick/patch/python/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ladyrick/patch/python/__main__.py,sha256=BAGy1phd26WWcGM9TKHbqIpeZliVofopBndtMIPtDQ0,651
 ladyrick/patch/python/usercustomize.py,sha256=8mYpcZ8p-l41fiSJue727n8cAmcEmUktObDYZDdLJfs,218
-ladyrick-0.5.0.dist-info/licenses/LICENSE,sha256=EeNAFxYAOYEmo2YEM7Zk5Oknq4RI0XMAbk4Rgoem6fs,1065
-ladyrick-0.5.0.dist-info/METADATA,sha256=LqfRIxtvCt1XMtpvjjacYbq7qBPoSiZMK8sUdmOs8Pw,883
-ladyrick-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ladyrick-0.5.0.dist-info/entry_points.txt,sha256=RlkKx2ecXAbDLfeEI_lIEvC21ZyEGli8hZmu8hIYE-E,159
-ladyrick-0.5.0.dist-info/top_level.txt,sha256=RIC3-Jty2qzLYXSOr7fOu1loTwlMU9cF6MFeGIROxWU,9
-ladyrick-0.5.0.dist-info/RECORD,,
+ladyrick-0.5.2.dist-info/licenses/LICENSE,sha256=EeNAFxYAOYEmo2YEM7Zk5Oknq4RI0XMAbk4Rgoem6fs,1065
+ladyrick-0.5.2.dist-info/METADATA,sha256=K3kyX7wNOmge7CX0mZl7HtOWl3ai9GVFFfOdN7gkUtA,883
+ladyrick-0.5.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ladyrick-0.5.2.dist-info/entry_points.txt,sha256=WVddlz7Wn9t1kQkA-HluVeeyKRJB3-7KrRpmWxfR2cM,195
+ladyrick-0.5.2.dist-info/top_level.txt,sha256=RIC3-Jty2qzLYXSOr7fOu1loTwlMU9cF6MFeGIROxWU,9
+ladyrick-0.5.2.dist-info/RECORD,,

{ladyrick-0.5.0.dist-info → ladyrick-0.5.2.dist-info}/entry_points.txt

@@ -1,4 +1,5 @@
 [console_scripts]
+allgather = ladyrick.allgather:main
 ladyrick-tee = ladyrick.cli.tee:main
 multi-ssh = ladyrick.cli.multi_ssh:main
 pretty-print = ladyrick.pprint:main