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

ladyrick/cli/calc.py

@@ -5,15 +5,15 @@ 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)
+    exec("import math; from math import *", None, imports)
+    exec("import itertools; from itertools import *", None, imports)
+    exec("import functools; from functools import cache, lru_cache, partial, reduce", None, imports)
+    exec("import os, sys, time, re, random", None, imports)
+    exec("from pathlib import Path", None, imports)
+    exec("from time import sleep", None, imports)
+    exec("from subprocess import check_output", None, imports)
+    exec("from random import randint, choice, random as rand", None, imports)
+    exec("import ladyrick", None, imports)
     if interactive_mode:
         imports["ladyrick"].pprint.reg_ipython_cmd()
     globals().update(imports)

ladyrick/cli/multi_ssh.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 import select
 import signal
 import subprocess
@@ -21,7 +22,7 @@ def force_kill(child: subprocess.Popen, child_pgid):
         child.kill()  # kill -9
 
 
-REMOTE_HEAD_PROG_NAME = "ladyrick-remote-head"
+REMOTE_HEAD_PROG_NAME = "ladyrick/multi-ssh/remote-head"
 
 
 def remote_head():
@@ -51,14 +52,14 @@ def remote_head():
     def handle_signal(sig, frame=None):
         if sig == signal.SIGUSR2:
             # SIGUSR2 trigger force_kill manually
-            log("SIGUSR2 received. force kill")
+            print("SIGUSR2 received. force kill")
             force_kill(child, child_pgid)
         else:
-            log(f"forward signal {sig} to {child.pid}")
+            print(f"forward signal {sig} to {child.pid}")
             try:
                 os.kill(child.pid, sig)
             except ProcessLookupError as e:
-                log(str(e))
+                print(str(e))
 
     for sig in [signal.SIGHUP, signal.SIGINT, signal.SIGTERM, signal.SIGUSR1, signal.SIGUSR2]:
         signal.signal(sig, handle_signal)
@@ -110,15 +111,20 @@ class Host:
 
 
 class RemoteExecutor:
-    def __init__(self, host: Host, command: list[str]):
+    def __init__(self, host: Host, command: list[str], envs: dict | None = None):
         self.host = host
         self.command = command
         self.envs = {}
+        if envs:
+            for k, v in envs.items():
+                if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", k):
+                    raise ValueError(f"invalid env name: {k!r}")
+                self.envs[k] = v
         self.process = None
 
     @classmethod
     def make_ssh_cmd(cls, host: Host, cmd: str):
-        opts = ["/usr/bin/env", "ssh", "-qT", "-oStrictHostKeyChecking=no"]
+        opts = ["/usr/bin/env", "ssh", "-T", "-oStrictHostKeyChecking=no"]
         if host.config_file is not None:
             opts.append(f"-F{host.config_file}")
         if host.User is not None:
@@ -134,6 +140,7 @@ class RemoteExecutor:
         return opts
 
     def start(self):
+        assert self.process is None
         code = pathlib.Path(__file__).read_text().split("# ----- remote_head end ----- #")[0].strip()
         remote_cmd = shlex.join(
             [
@@ -185,15 +192,19 @@ class RemoteExecutor:
             e.envs["RANK"] = str(i)
 
     def send_signal(self, sig):
-        if self.process:
+        assert self.process is not None
+        if self.process.poll() is None and self.process.stdin and not self.process.stdin.closed:
             sig_name = signal.Signals(sig).name
-            if self.process.stdin and not self.process.stdin.closed:
-                log(f"writing to stdin: SIGNAL {sig_name}")
-                try:
-                    self.process.stdin.write(f"SIGNAL {sig_name}\n".encode())
-                    self.process.stdin.flush()
-                except (BrokenPipeError, OSError) as e:
-                    log(e)
+            log(f"writing to stdin: SIGNAL {sig_name}")
+            try:
+                self.process.stdin.write(f"SIGNAL {sig_name}\n".encode())
+                self.process.stdin.flush()
+            except (BrokenPipeError, OSError) as e:
+                log(e)
+
+    def poll(self):
+        assert self.process is not None
+        return self.process.poll()
 
 
 def signal_repeat_checker(sig_to_check, count, duration):
@@ -223,6 +234,7 @@ 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("-e", "--env", type=str, action="append", help="extra envs")
     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")
@@ -263,7 +275,15 @@ def main():
         parser.print_help()
         sys.exit(1)
 
-    executors = [RemoteExecutor(host, args.cmd) for host in hosts]
+    envs = {}
+    if args.env:
+        for e in args.env:
+            p = e.split("=", 1)
+            if len(p) == 1:
+                p.append("")
+            envs[p[0]] = p[1]
+
+    executors = [RemoteExecutor(host, args.cmd, envs) for host in hosts]
 
     RemoteExecutor.set_envs(executors)
 
@@ -284,7 +304,7 @@ def main():
     for sig in [signal.SIGHUP, signal.SIGINT, signal.SIGTERM, signal.SIGUSR1, signal.SIGUSR2]:
         signal.signal(sig, handle_signal)
 
-    while any(e.process.poll() is None for e in executors if e.process):
+    while any([e.poll() is None for e in executors if e.process]):
         time.sleep(0.5)
     log("finished")
 

ladyrick/cli/test_signal.py

@@ -5,8 +5,8 @@ import time
 
 def handle_signal(sig, frame):
     print(f"received signal {sig}, {signal.Signals(sig).name}", flush=True)
-    if sig in {signal.SIGTERM, signal.SIGINT, signal.SIGHUP}:
-        sys.exit(0)
+    # if sig in {signal.SIGTERM, signal.SIGINT, signal.SIGHUP}:
+        # sys.exit(0)
 
 
 for sig in [s for s in vars(signal).values() if isinstance(s, signal.Signals)]:

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

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

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

@@ -9,21 +9,21 @@ 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/calc.py,sha256=szg4dyCOsWHNGKPcBcL3Tg4dCCZzBFiGow1O98pyxH4,1447
-ladyrick/cli/multi_ssh.py,sha256=faRIgradz3e6ZjwpPd77OB_u5FEJFCfVzGq3qrwJsgk,9220
+ladyrick/cli/calc.py,sha256=5_UhSSaL_K9FBCHg3zuCk41CHJXy3QAh__qCabbffQY,1438
+ladyrick/cli/multi_ssh.py,sha256=02DmcotlKGlaAEIgIi5KsnjBxkNnf60OiimpH20JfYU,9903
 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
+ladyrick/cli/test_signal.py,sha256=zGPGMR81j1_vNyGEAcRtO2WjhDXiCj2xHMDATk-Z33M,494
 ladyrick/patch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ladyrick/patch/print.py,sha256=Ks_ONpeLOyC9pFKLmAiNliwGPbJyt9qz-cnRUPVto34,845
 ladyrick/patch/rich_print.py,sha256=z3Ea1VCunXZvNvEDFHpoyWc8ydINmh-gOIJ1ssscs6s,841
 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.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.5.dist-info/licenses/LICENSE,sha256=EeNAFxYAOYEmo2YEM7Zk5Oknq4RI0XMAbk4Rgoem6fs,1065
+ladyrick-0.5.5.dist-info/METADATA,sha256=5F-dTbYGPoApk--A2gPS5cS_Ev4f3WRRRLKbPrYq4S8,883
+ladyrick-0.5.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ladyrick-0.5.5.dist-info/entry_points.txt,sha256=28kidI1OCOAi7S1aHBr8veak49VNZ0tWFvf-Ty8UmSU,225
+ladyrick-0.5.5.dist-info/top_level.txt,sha256=RIC3-Jty2qzLYXSOr7fOu1loTwlMU9cF6MFeGIROxWU,9
+ladyrick-0.5.5.dist-info/RECORD,,