ladyrick 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

ladyrick/cli/multi_ssh.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 import select
 import signal
 import subprocess
@@ -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,18 +192,27 @@ 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 terminate(self):
+        if self.process is not None and self.process.poll() is None:
+            log("terminate RemoteExecutor")
+            self.process.terminate()
+
+    def poll(self):
+        assert self.process is not None
+        return self.process.poll()
 
 
-def signal_repeat_checker(sig_to_check, count, duration):
+def signal_repeat_checker(sig_to_check, duration: float):
     last_int_signal_time = []
 
     def checker(sig: signal.Signals):
@@ -206,11 +222,8 @@ def signal_repeat_checker(sig_to_check, count, duration):
             threadhold = cur_time - duration
             last_int_signal_time = [t for t in last_int_signal_time if t >= threadhold]
             last_int_signal_time.append(cur_time)
-
-            if len(last_int_signal_time) >= count:
-                log(f"received {sig_to_check.name} for {count} times or more in {duration} second(s)")
-                return True
-        return False
+            return len(last_int_signal_time)
+        return 0
 
     return checker
 
@@ -223,6 +236,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,28 +277,47 @@ 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)
 
     for executor in executors:
         executor.start()
 
-    checker = signal_repeat_checker(signal.SIGINT, count=3, duration=1)
+    import rich
+
+    checker = signal_repeat_checker(signal.SIGINT, duration=1)
 
     def handle_signal(sig, frame):
-        if checker(sig):
+        log(f"received signal {sig}")
+        sig_count = checker(sig)
+        if sig_count >= 3:
             sig = signal.SIGUSR2
-            log("\nTry to froce kill remote processes")
+            if sig_count == 3:
+                rich.print("\n[bold magenta]Can't wait. Try to froce kill remote processes...[/bold magenta]")
         else:
-            log(f"\nReceived {signal.Signals(sig).name}, forwarding...")
+            rich.print(
+                f"\n[bold green]Received {signal.Signals(sig).name}, forwarding to remote processes...[/bold green]"
+            )
         for executor in executors:
             executor.send_signal(sig)
+        if sig_count >= 4:
+            rich.print("\n[bold red]Really Can't wait!!! Froce kill local processes and exiting right now![/bold red]")
+            for executor in executors:
+                executor.terminate()
 
     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-0.5.4.dist-info → ladyrick-0.5.6.dist-info}/METADATA

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

{ladyrick-0.5.4.dist-info → ladyrick-0.5.6.dist-info}/RECORD

@@ -10,7 +10,7 @@ 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=5_UhSSaL_K9FBCHg3zuCk41CHJXy3QAh__qCabbffQY,1438
-ladyrick/cli/multi_ssh.py,sha256=wbO7woW0roAd6Q5LW861c8h5r7sfkCKW2qvCq-HssTY,9230
+ladyrick/cli/multi_ssh.py,sha256=d_7zPvmfVUQZVcxhCFj37ktTOb0-AzxSEr0z6ipMY9U,10403
 ladyrick/cli/psf.py,sha256=JLk3gbPn7E3uuPBbzGvLgJmFQlilA6zg_Xlg7xW5jik,1146
 ladyrick/cli/tee.py,sha256=UMJxSJLOEfbV43auVKRTIJ5ZAMAkAfj8byiFLk5PUHE,3579
 ladyrick/cli/test_args.py,sha256=f5sUPDlcf6nbNf6UfLwZQI5g5LN8wlFBQZ10GLw22cg,212
@@ -21,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.4.dist-info/licenses/LICENSE,sha256=EeNAFxYAOYEmo2YEM7Zk5Oknq4RI0XMAbk4Rgoem6fs,1065
-ladyrick-0.5.4.dist-info/METADATA,sha256=lAnPzvF3RUZk9942XC-UDUMZuiD77o_h6LFRZx2jj-c,883
-ladyrick-0.5.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ladyrick-0.5.4.dist-info/entry_points.txt,sha256=28kidI1OCOAi7S1aHBr8veak49VNZ0tWFvf-Ty8UmSU,225
-ladyrick-0.5.4.dist-info/top_level.txt,sha256=RIC3-Jty2qzLYXSOr7fOu1loTwlMU9cF6MFeGIROxWU,9
-ladyrick-0.5.4.dist-info/RECORD,,
+ladyrick-0.5.6.dist-info/licenses/LICENSE,sha256=EeNAFxYAOYEmo2YEM7Zk5Oknq4RI0XMAbk4Rgoem6fs,1065
+ladyrick-0.5.6.dist-info/METADATA,sha256=s-vKvkA34C7DS3PXt2QygfVzhmBxSJC42bawDe_QC8Q,883
+ladyrick-0.5.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ladyrick-0.5.6.dist-info/entry_points.txt,sha256=28kidI1OCOAi7S1aHBr8veak49VNZ0tWFvf-Ty8UmSU,225
+ladyrick-0.5.6.dist-info/top_level.txt,sha256=RIC3-Jty2qzLYXSOr7fOu1loTwlMU9cF6MFeGIROxWU,9
+ladyrick-0.5.6.dist-info/RECORD,,