ominfra 0.0.0.dev147__py3-none-any.whl → 0.0.0.dev148__py3-none-any.whl

ominfra/configs.py

@@ -1,6 +1,7 @@
 # ruff: noqa: UP006 UP007
 # @omlish-lite
 import json
+import os.path
 import typing as ta
 
 from omdev.toml.parser import toml_loads
@@ -14,6 +15,33 @@ T = ta.TypeVar('T')
 ConfigMapping = ta.Mapping[str, ta.Any]
 
 
+def parse_config_file(
+        name: str,
+        f: ta.TextIO,
+) -> ConfigMapping:
+    if name.endswith('.toml'):
+        return toml_loads(f.read())
+
+    elif any(name.endswith(e) for e in ('.yml', '.yaml')):
+        yaml = __import__('yaml')
+        return yaml.safe_load(f)
+
+    elif name.endswith('.ini'):
+        import configparser
+        cp = configparser.ConfigParser()
+        cp.read_file(f)
+        config_dct: ta.Dict[str, ta.Any] = {}
+        for sec in cp.sections():
+            cd = config_dct
+            for k in sec.split('.'):
+                cd = cd.setdefault(k, {})
+            cd.update(cp.items(sec))
+        return config_dct
+
+    else:
+        return json.loads(f.read())
+
+
 def read_config_file(
         path: str,
         cls: ta.Type[T],
@@ -21,13 +49,10 @@ def read_config_file(
         prepare: ta.Optional[ta.Callable[[ConfigMapping], ConfigMapping]] = None,
 ) -> T:
     with open(path) as cf:
-        if path.endswith('.toml'):
-            config_dct = toml_loads(cf.read())
-        else:
-            config_dct = json.loads(cf.read())
+        config_dct = parse_config_file(os.path.basename(path), cf)
 
     if prepare is not None:
-        config_dct = prepare(config_dct)  # type: ignore
+        config_dct = prepare(config_dct)
 
     return unmarshal_obj(config_dct, cls)
 

ominfra/manage/commands/base.py

@@ -23,8 +23,8 @@ class Command(abc.ABC, ta.Generic[CommandOutputT]):
         pass
 
     @ta.final
-    def execute(self, executor: 'CommandExecutor') -> CommandOutputT:
-        return check_isinstance(executor.execute(self), self.Output)  # type: ignore[return-value]
+    async def execute(self, executor: 'CommandExecutor') -> CommandOutputT:
+        return check_isinstance(await executor.execute(self), self.Output)  # type: ignore[return-value]
 
 
 ##
@@ -85,10 +85,10 @@ class CommandOutputOrExceptionData(CommandOutputOrException):
 
 class CommandExecutor(abc.ABC, ta.Generic[CommandT, CommandOutputT]):
     @abc.abstractmethod
-    def execute(self, cmd: CommandT) -> CommandOutputT:
+    def execute(self, cmd: CommandT) -> ta.Awaitable[CommandOutputT]:
         raise NotImplementedError
 
-    def try_execute(
+    async def try_execute(
             self,
             cmd: CommandT,
             *,
@@ -96,7 +96,7 @@ class CommandExecutor(abc.ABC, ta.Generic[CommandT, CommandOutputT]):
             omit_exc_object: bool = False,
     ) -> CommandOutputOrException[CommandOutputT]:
         try:
-            o = self.execute(cmd)
+            o = await self.execute(cmd)
 
         except Exception as e:  # noqa
             if log is not None:

ominfra/manage/commands/execution.py

@@ -18,6 +18,6 @@ class LocalCommandExecutor(CommandExecutor):
 
         self._command_executors = command_executors
 
-    def execute(self, cmd: Command) -> Command.Output:
+    async def execute(self, cmd: Command) -> Command.Output:
         ce: CommandExecutor = self._command_executors[type(cmd)]
-        return ce.execute(cmd)
+        return await ce.execute(cmd)

ominfra/manage/commands/inject.py

@@ -54,7 +54,7 @@ def bind_command(
 class _FactoryCommandExecutor(CommandExecutor):
     factory: ta.Callable[[], CommandExecutor]
 
-    def execute(self, i: Command) -> Command.Output:
+    def execute(self, i: Command) -> ta.Awaitable[Command.Output]:
         return self.factory().execute(i)
 
 

ominfra/manage/commands/interp.py

@@ -29,9 +29,9 @@ class InterpCommand(Command['InterpCommand.Output']):
 
 
 class InterpCommandExecutor(CommandExecutor[InterpCommand, InterpCommand.Output]):
-    def execute(self, cmd: InterpCommand) -> InterpCommand.Output:
+    async def execute(self, cmd: InterpCommand) -> InterpCommand.Output:
         i = InterpSpecifier.parse(check_not_none(cmd.spec))
-        o = check_not_none(DEFAULT_INTERP_RESOLVER.resolve(i, install=cmd.install))
+        o = check_not_none(await DEFAULT_INTERP_RESOLVER.resolve(i, install=cmd.install))
         return InterpCommand.Output(
             exe=o.exe,
             version=str(o.version.version),

ominfra/manage/commands/subprocess.py

@@ -1,11 +1,15 @@
-# ruff: noqa: UP006 UP007
+# ruff: noqa: TC003 UP006 UP007
+import asyncio.subprocess
 import dataclasses as dc
 import os
 import subprocess
 import time
 import typing as ta
 
+from omlish.lite.asyncio.subprocesses import asyncio_subprocess_communicate
+from omlish.lite.asyncio.subprocesses import asyncio_subprocess_popen
 from omlish.lite.check import check_not_isinstance
+from omlish.lite.check import check_not_none
 from omlish.lite.subprocesses import SUBPROCESS_CHANNEL_OPTION_VALUES
 from omlish.lite.subprocesses import SubprocessChannelOption
 from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
@@ -49,27 +53,31 @@ class SubprocessCommand(Command['SubprocessCommand.Output']):
 
 
 class SubprocessCommandExecutor(CommandExecutor[SubprocessCommand, SubprocessCommand.Output]):
-    def execute(self, inp: SubprocessCommand) -> SubprocessCommand.Output:
-        with subprocess.Popen(
-            subprocess_maybe_shell_wrap_exec(*inp.cmd),
+    async def execute(self, cmd: SubprocessCommand) -> SubprocessCommand.Output:
+        proc: asyncio.subprocess.Process
+        async with asyncio_subprocess_popen(
+            *subprocess_maybe_shell_wrap_exec(*cmd.cmd),
 
-            shell=inp.shell,
-            cwd=inp.cwd,
-            env={**os.environ, **(inp.env or {})},
+            shell=cmd.shell,
+            cwd=cmd.cwd,
+            env={**os.environ, **(cmd.env or {})},
 
-            stdin=subprocess.PIPE if inp.input is not None else None,
-            stdout=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, inp.stdout)],
-            stderr=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, inp.stderr)],
+            stdin=subprocess.PIPE if cmd.input is not None else None,
+            stdout=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, cmd.stdout)],
+            stderr=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, cmd.stderr)],
+
+            timeout=cmd.timeout,
         ) as proc:
             start_time = time.time()
-            stdout, stderr = proc.communicate(
-                input=inp.input,
-                timeout=inp.timeout,
+            stdout, stderr = await asyncio_subprocess_communicate(
+                proc,
+                input=cmd.input,
+                timeout=cmd.timeout,
             )
             end_time = time.time()
 
         return SubprocessCommand.Output(
-            rc=proc.returncode,
+            rc=check_not_none(proc.returncode),
             pid=proc.pid,
 
             elapsed_s=end_time - start_time,

ominfra/manage/deploy/command.py

@@ -21,7 +21,7 @@ class DeployCommand(Command['DeployCommand.Output']):
 
 
 class DeployCommandExecutor(CommandExecutor[DeployCommand, DeployCommand.Output]):
-    def execute(self, cmd: DeployCommand) -> DeployCommand.Output:
+    async def execute(self, cmd: DeployCommand) -> DeployCommand.Output:
         log.info('Deploying!')
 
         return DeployCommand.Output()

ominfra/manage/main.py

@@ -5,6 +5,7 @@
 manage.py -s 'docker run -i python:3.12'
 manage.py -s 'ssh -i /foo/bar.pem foo@bar.baz' -q --python=python3.8
 """
+import asyncio
 import contextlib
 import json
 import typing as ta
@@ -21,38 +22,14 @@ from .commands.base import CommandExecutor
 from .commands.execution import LocalCommandExecutor
 from .config import MainConfig
 from .remote.config import RemoteConfig
-from .remote.execution import RemoteExecution
+from .remote.connection import RemoteExecutionConnector
 from .remote.spawning import RemoteSpawning
 
 
 ##
 
 
-def _main() -> None:
-    import argparse
-
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument('--_payload-file')
-
-    parser.add_argument('-s', '--shell')
-    parser.add_argument('-q', '--shell-quote', action='store_true')
-    parser.add_argument('--python', default='python3')
-
-    parser.add_argument('--pycharm-debug-port', type=int)
-    parser.add_argument('--pycharm-debug-host')
-    parser.add_argument('--pycharm-debug-version')
-
-    parser.add_argument('--debug', action='store_true')
-
-    parser.add_argument('--local', action='store_true')
-
-    parser.add_argument('command', nargs='+')
-
-    args = parser.parse_args()
-
-    #
-
+async def _async_main(args: ta.Any) -> None:
     bs = MainBootstrap(
         main_config=MainConfig(
             log_level='DEBUG' if args.debug else 'INFO',
@@ -65,12 +42,16 @@ def _main() -> None:
 
             pycharm_remote_debug=PycharmRemoteDebug(
                 port=args.pycharm_debug_port,
-                host=args.pycharm_debug_host,
+                **(dict(host=args.pycharm_debug_host) if args.pycharm_debug_host is not None else {}),
                 install_version=args.pycharm_debug_version,
             ) if args.pycharm_debug_port is not None else None,
+
+            timebomb_delay_s=args.remote_timebomb_delay_s,
         ),
     )
 
+    #
+
     injector = main_bootstrap(
         bs,
     )
@@ -89,7 +70,7 @@ def _main() -> None:
 
     #
 
-    with contextlib.ExitStack() as es:
+    async with contextlib.AsyncExitStack() as es:
         ce: CommandExecutor
 
         if args.local:
@@ -102,16 +83,51 @@ def _main() -> None:
                 python=args.python,
             )
 
-            ce = es.enter_context(injector[RemoteExecution].connect(tgt, bs))  # noqa
+            ce = await es.enter_async_context(injector[RemoteExecutionConnector].connect(tgt, bs))  # noqa
 
-        for cmd in cmds:
-            r = ce.try_execute(
+        async def run_command(cmd: Command) -> None:
+            res = await ce.try_execute(
                 cmd,
                 log=log,
                 omit_exc_object=True,
             )
 
-            print(msh.marshal_obj(r, opts=ObjMarshalOptions(raw_bytes=True)))
+            print(msh.marshal_obj(res, opts=ObjMarshalOptions(raw_bytes=True)))
+
+        await asyncio.gather(*[
+            run_command(cmd)
+            for cmd in cmds
+        ])
+
+
+def _main() -> None:
+    import argparse
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--_payload-file')
+
+    parser.add_argument('-s', '--shell')
+    parser.add_argument('-q', '--shell-quote', action='store_true')
+    parser.add_argument('--python', default='python3')
+
+    parser.add_argument('--pycharm-debug-port', type=int)
+    parser.add_argument('--pycharm-debug-host')
+    parser.add_argument('--pycharm-debug-version')
+
+    parser.add_argument('--remote-timebomb-delay-s', type=float)
+
+    parser.add_argument('--debug', action='store_true')
+
+    parser.add_argument('--local', action='store_true')
+
+    parser.add_argument('command', nargs='+')
+
+    args = parser.parse_args()
+
+    #
+
+    asyncio.run(_async_main(args))
 
 
 if __name__ == '__main__':

ominfra/manage/remote/_main.py

@@ -0,0 +1,172 @@
+# ruff: noqa: UP006 UP007
+import asyncio
+import functools
+import logging
+import os
+import signal
+import threading
+import time
+import typing as ta
+
+from omlish.lite.asyncio.asyncio import asyncio_open_stream_reader
+from omlish.lite.asyncio.asyncio import asyncio_open_stream_writer
+from omlish.lite.cached import cached_nullary
+from omlish.lite.check import check_none
+from omlish.lite.check import check_not_none
+from omlish.lite.deathsig import set_process_deathsig
+from omlish.lite.inject import Injector
+from omlish.lite.logs import log
+from omlish.lite.marshal import ObjMarshalerManager
+from omlish.lite.pycharm import pycharm_debug_connect
+
+from ...pyremote import pyremote_bootstrap_finalize
+from ..bootstrap import MainBootstrap
+from ..commands.execution import LocalCommandExecutor
+from .channel import RemoteChannel
+from .channel import RemoteChannelImpl
+from .execution import _RemoteCommandHandler
+from .execution import _RemoteLogHandler
+
+
+if ta.TYPE_CHECKING:
+    from ..bootstrap_ import main_bootstrap
+else:
+    main_bootstrap: ta.Any = None
+
+
+##
+
+
+class _RemoteExecutionLogHandler(logging.Handler):
+    def __init__(self, fn: ta.Callable[[str], None]) -> None:
+        super().__init__()
+        self._fn = fn
+
+    def emit(self, record):
+        msg = self.format(record)
+        self._fn(msg)
+
+
+##
+
+
+class _RemoteExecutionMain:
+    def __init__(
+            self,
+            chan: RemoteChannel,
+    ) -> None:
+        super().__init__()
+
+        self._chan = chan
+
+        self.__bootstrap: ta.Optional[MainBootstrap] = None
+        self.__injector: ta.Optional[Injector] = None
+
+    @property
+    def _bootstrap(self) -> MainBootstrap:
+        return check_not_none(self.__bootstrap)
+
+    @property
+    def _injector(self) -> Injector:
+        return check_not_none(self.__injector)
+
+    #
+
+    def _timebomb_main(
+            self,
+            delay_s: float,
+            *,
+            sig: int = signal.SIGINT,
+            code: int = 1,
+    ) -> None:
+        time.sleep(delay_s)
+
+        if (pgid := os.getpgid(0)) == os.getpid():
+            os.killpg(pgid, sig)
+
+        os._exit(code)  # noqa
+
+    @cached_nullary
+    def _timebomb_thread(self) -> ta.Optional[threading.Thread]:
+        if (tbd := self._bootstrap.remote_config.timebomb_delay_s) is None:
+            return None
+
+        thr = threading.Thread(
+            target=functools.partial(self._timebomb_main, tbd),
+            name=f'{self.__class__.__name__}.timebomb',
+            daemon=True,
+        )
+
+        thr.start()
+
+        log.debug('Started timebomb thread: %r', thr)
+
+        return thr
+
+    #
+
+    @cached_nullary
+    def _log_handler(self) -> _RemoteLogHandler:
+        return _RemoteLogHandler(self._chan)
+
+    #
+
+    async def _setup(self) -> None:
+        check_none(self.__bootstrap)
+        check_none(self.__injector)
+
+        # Bootstrap
+
+        self.__bootstrap = check_not_none(await self._chan.recv_obj(MainBootstrap))
+
+        if (prd := self._bootstrap.remote_config.pycharm_remote_debug) is not None:
+            pycharm_debug_connect(prd)
+
+        self.__injector = main_bootstrap(self._bootstrap)
+
+        self._chan.set_marshaler(self._injector[ObjMarshalerManager])
+
+        # Post-bootstrap
+
+        if self._bootstrap.remote_config.set_pgid:
+            if os.getpgid(0) != os.getpid():
+                log.debug('Setting pgid')
+                os.setpgid(0, 0)
+
+        if (ds := self._bootstrap.remote_config.deathsig) is not None:
+            log.debug('Setting deathsig: %s', ds)
+            set_process_deathsig(int(signal.Signals[f'SIG{ds.upper()}']))
+
+        self._timebomb_thread()
+
+        if self._bootstrap.remote_config.forward_logging:
+            log.debug('Installing log forwarder')
+            logging.root.addHandler(self._log_handler())
+
+    #
+
+    async def run(self) -> None:
+        await self._setup()
+
+        executor = self._injector[LocalCommandExecutor]
+
+        handler = _RemoteCommandHandler(self._chan, executor)
+
+        await handler.run()
+
+
+def _remote_execution_main() -> None:
+    rt = pyremote_bootstrap_finalize()  # noqa
+
+    async def inner() -> None:
+        input = await asyncio_open_stream_reader(rt.input)  # noqa
+        output = await asyncio_open_stream_writer(rt.output)
+
+        chan = RemoteChannelImpl(
+            input,
+            output,
+        )
+
+        await _RemoteExecutionMain(chan).run()
+
+    asyncio.run(inner())

ominfra/manage/remote/channel.py

@@ -1,7 +1,8 @@
 # ruff: noqa: UP006 UP007
+import abc
+import asyncio
 import json
 import struct
-import threading
 import typing as ta
 
 from omlish.lite.json import json_dumps_compact
@@ -12,11 +13,30 @@ from omlish.lite.marshal import ObjMarshalerManager
 T = ta.TypeVar('T')
 
 
-class RemoteChannel:
+##
+
+
+class RemoteChannel(abc.ABC):
+    @abc.abstractmethod
+    def send_obj(self, o: ta.Any, ty: ta.Any = None) -> ta.Awaitable[None]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def recv_obj(self, ty: ta.Type[T]) -> ta.Awaitable[ta.Optional[T]]:
+        raise NotImplementedError
+
+    def set_marshaler(self, msh: ObjMarshalerManager) -> None:  # noqa
+        pass
+
+
+##
+
+
+class RemoteChannelImpl(RemoteChannel):
     def __init__(
             self,
-            input: ta.IO,  # noqa
-            output: ta.IO,
+            input: asyncio.StreamReader,  # noqa
+            output: asyncio.StreamWriter,
             *,
             msh: ObjMarshalerManager = OBJ_MARSHALER_MANAGER,
     ) -> None:
@@ -26,38 +46,43 @@ class RemoteChannel:
         self._output = output
         self._msh = msh
 
-        self._lock = threading.RLock()
+        self._input_lock = asyncio.Lock()
+        self._output_lock = asyncio.Lock()
 
     def set_marshaler(self, msh: ObjMarshalerManager) -> None:
         self._msh = msh
 
-    def _send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
+    #
+
+    async def _send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
         j = json_dumps_compact(self._msh.marshal_obj(o, ty))
         d = j.encode('utf-8')
 
         self._output.write(struct.pack('<I', len(d)))
         self._output.write(d)
-        self._output.flush()
+        await self._output.drain()
+
+    async def send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
+        async with self._output_lock:
+            return await self._send_obj(o, ty)
 
-    def send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
-        with self._lock:
-            return self._send_obj(o, ty)
+    #
 
-    def _recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
-        d = self._input.read(4)
+    async def _recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
+        d = await self._input.read(4)
         if not d:
             return None
         if len(d) != 4:
             raise EOFError
 
         sz = struct.unpack('<I', d)[0]
-        d = self._input.read(sz)
+        d = await self._input.read(sz)
         if len(d) != sz:
             raise EOFError
 
         j = json.loads(d.decode('utf-8'))
         return self._msh.unmarshal_obj(j, ty)
 
-    def recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
-        with self._lock:
-            return self._recv_obj(ty)
+    async def recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
+        async with self._input_lock:
+            return await self._recv_obj(ty)

ominfra/manage/remote/config.py

@@ -9,4 +9,14 @@ from omlish.lite.pycharm import PycharmRemoteDebug
 class RemoteConfig:
     payload_file: ta.Optional[str] = None
 
+    set_pgid: bool = True
+
+    deathsig: ta.Optional[str] = 'KILL'
+
     pycharm_remote_debug: ta.Optional[PycharmRemoteDebug] = None
+
+    forward_logging: bool = True
+
+    timebomb_delay_s: ta.Optional[float] = 60 * 60.
+
+    heartbeat_interval_s: float = 3.