ladyrick 0.5.1__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

@@ -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))
 

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

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

{ladyrick-0.5.1.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,7 +9,7 @@ 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=uS2843LPUZTtr08CrXPUyw4PomagMmuLeZNgM-dAUZk,9086
+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
@@ -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.1.dist-info/licenses/LICENSE,sha256=EeNAFxYAOYEmo2YEM7Zk5Oknq4RI0XMAbk4Rgoem6fs,1065
-ladyrick-0.5.1.dist-info/METADATA,sha256=YfwVxt9gm68kRBlwdY0DQKIqvuky_uAXw82hiVW5nZQ,883
-ladyrick-0.5.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ladyrick-0.5.1.dist-info/entry_points.txt,sha256=RlkKx2ecXAbDLfeEI_lIEvC21ZyEGli8hZmu8hIYE-E,159
-ladyrick-0.5.1.dist-info/top_level.txt,sha256=RIC3-Jty2qzLYXSOr7fOu1loTwlMU9cF6MFeGIROxWU,9
-ladyrick-0.5.1.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.1.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