ladyrick 0.5.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

ladyrick/__init__.py

@@ -6,7 +6,20 @@ from typing import TYPE_CHECKING
 def __getattr__(name):
     from importlib import import_module
 
-    if name in ("pretty_print",):
+    if name in (
+        "allgather",
+        "debug",
+        "loader",
+        "pickle",
+        "pprint",
+        "print_utils",
+        "torch",
+        "typing",
+        "utils",
+        "vars",
+    ):
+        return import_module(f"ladyrick.{name}")
+    elif name in ("pretty_print",):
         m = import_module("ladyrick.pprint")
     elif name in ("class_name", "utc_8_now", "get_timestr"):
         m = import_module("ladyrick.utils")
@@ -33,3 +46,13 @@ if TYPE_CHECKING:
     from ladyrick.typing import type_like  # noqa
     from ladyrick.utils import class_name, get_timestr, utc_8_now  # noqa
     from ladyrick.vars import Dump, V, Vars, dump  # noqa
+    from ladyrick import allgather  # noqa
+    from ladyrick import debug  # noqa
+    from ladyrick import loader  # noqa
+    from ladyrick import pickle  # noqa
+    from ladyrick import pprint  # noqa
+    from ladyrick import print_utils  # noqa
+    from ladyrick import torch  # noqa
+    from ladyrick import typing  # noqa
+    from ladyrick import utils  # noqa
+    from ladyrick import vars  # noqa

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/calc.py

@@ -0,0 +1,42 @@
+import os
+import sys
+
+
+def initialize(interactive_mode: bool):
+    imports = {}
+    exec("import ladyrick.patch.rich_print")
+    exec("import math; from math import *", locals=imports)
+    exec("import itertools; from itertools import *", locals=imports)
+    exec("import functools; from functools import cache, lru_cache, partial, reduce", locals=imports)
+    exec("import os, sys, time, re, random", locals=imports)
+    exec("from pathlib import Path", locals=imports)
+    exec("from time import sleep", locals=imports)
+    exec("from subprocess import check_output", locals=imports)
+    exec("from random import randint, choice, random as rand", locals=imports)
+    exec("import ladyrick", locals=imports)
+    if interactive_mode:
+        imports["ladyrick"].pprint.reg_ipython_cmd()
+    globals().update(imports)
+
+
+def main():
+    if len(sys.argv) == 1 and not os.getenv("__ladyrick_calc__"):
+        import setproctitle
+
+        os.environ["__ladyrick_calc__"] = setproctitle.getproctitle()
+        os.execl(sys.executable, sys.executable, "-m", "IPython", "-i", "--no-banner", "--no-confirm-exit", __file__)
+    else:
+        if len(sys.argv) == 1:
+            import setproctitle
+
+            setproctitle.setproctitle(os.environ["__ladyrick_calc__"])
+            del setproctitle
+
+        initialize(len(sys.argv) == 1)
+
+        if len(sys.argv) > 1:
+            print(eval(" ".join(sys.argv[1:])))
+
+
+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/pprint.py

@@ -251,6 +251,17 @@ def pretty_print(model):
     Printer().print(model, FakeRichStream())
 
 
+def reg_ipython_cmd():
+    from IPython.core.magic import register_line_magic
+
+    @register_line_magic
+    def pp(line: str):
+        if line.strip():
+            pretty_print(eval(line))
+        else:
+            rich_print("pp is pretty_print.\nusage: pp <exps>")
+
+
 @dataclasses.dataclass
 class load_failed:
     file: str
@@ -321,12 +332,7 @@ def main():
         if "torch" in sys.modules:
             add_to_globals["torch"] = sys.modules["torch"]
         globals().update(add_to_globals)
-
-        from IPython.core.magic import register_line_magic
-
-        @register_line_magic
-        def pp(line):
-            pretty_print(eval(line))
+        reg_ipython_cmd()
 
 
 if __name__ == "__main__":

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

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

{ladyrick-0.5.1.dist-info → ladyrick-0.5.3.dist-info}/RECORD

@@ -1,14 +1,16 @@
-ladyrick/__init__.py,sha256=udDiOLu2L1jIvjgSCLdDsTnS2OA7kdj2BA9kit2maVo,1343
+ladyrick/__init__.py,sha256=wkastx1reQfUUtx9n5cW-CM-9f0F0uX7Qek2vex782s,2004
+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
-ladyrick/pprint.py,sha256=E_oginnkLmYbvDKzm82KujkqSSjIcFhjU21-l1Wo-Qk,11350
+ladyrick/pprint.py,sha256=K9nzrO9pkQ0smgLDEMmSAm9_MZCgqYGFT-XCwVZFn_o,11496
 ladyrick/print_utils.py,sha256=NRalqrTy8PWj1tL2K2wAlXbQpiX1_N7P0-rQ9kmmHEQ,3472
 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/calc.py,sha256=szg4dyCOsWHNGKPcBcL3Tg4dCCZzBFiGow1O98pyxH4,1447
+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 +21,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.3.dist-info/licenses/LICENSE,sha256=EeNAFxYAOYEmo2YEM7Zk5Oknq4RI0XMAbk4Rgoem6fs,1065
+ladyrick-0.5.3.dist-info/METADATA,sha256=JseJXyRyBHz1PrCo-kcHp_EoezLHCyt3yR6qwyQ9VDE,883
+ladyrick-0.5.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ladyrick-0.5.3.dist-info/entry_points.txt,sha256=28kidI1OCOAi7S1aHBr8veak49VNZ0tWFvf-Ty8UmSU,225
+ladyrick-0.5.3.dist-info/top_level.txt,sha256=RIC3-Jty2qzLYXSOr7fOu1loTwlMU9cF6MFeGIROxWU,9
+ladyrick-0.5.3.dist-info/RECORD,,

{ladyrick-0.5.1.dist-info → ladyrick-0.5.3.dist-info}/entry_points.txt

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