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

ominfra/manage/remote/connection.py

@@ -0,0 +1,106 @@
+# ruff: noqa: UP006 UP007
+import abc
+import contextlib
+import typing as ta
+
+from omlish.lite.cached import cached_nullary
+from omlish.lite.marshal import ObjMarshalerManager
+
+from ...pyremote import PyremoteBootstrapDriver
+from ...pyremote import PyremoteBootstrapOptions
+from ...pyremote import pyremote_build_bootstrap_cmd
+from ..bootstrap import MainBootstrap
+from ._main import _remote_execution_main  # noqa
+from .channel import RemoteChannelImpl
+from .execution import RemoteCommandExecutor
+from .payload import RemoteExecutionPayloadFile
+from .payload import get_remote_payload_src
+from .spawning import RemoteSpawning
+
+
+##
+
+
+class RemoteExecutionConnector(abc.ABC):
+    @abc.abstractmethod
+    def connect(
+            self,
+            tgt: RemoteSpawning.Target,
+            bs: MainBootstrap,
+    ) -> ta.AsyncContextManager[RemoteCommandExecutor]:
+        raise NotImplementedError
+
+
+##
+
+
+class PyremoteRemoteExecutionConnector(RemoteExecutionConnector):
+    def __init__(
+            self,
+            *,
+            spawning: RemoteSpawning,
+            msh: ObjMarshalerManager,
+            payload_file: ta.Optional[RemoteExecutionPayloadFile] = None,
+    ) -> None:
+        super().__init__()
+
+        self._spawning = spawning
+        self._msh = msh
+        self._payload_file = payload_file
+
+    #
+
+    @cached_nullary
+    def _payload_src(self) -> str:
+        return get_remote_payload_src(file=self._payload_file)
+
+    @cached_nullary
+    def _remote_src(self) -> ta.Sequence[str]:
+        return [
+            self._payload_src(),
+            '_remote_execution_main()',
+        ]
+
+    @cached_nullary
+    def _spawn_src(self) -> str:
+        return pyremote_build_bootstrap_cmd(__package__ or 'manage')
+
+    #
+
+    @contextlib.asynccontextmanager
+    async def connect(
+            self,
+            tgt: RemoteSpawning.Target,
+            bs: MainBootstrap,
+    ) -> ta.AsyncGenerator[RemoteCommandExecutor, None]:
+        spawn_src = self._spawn_src()
+        remote_src = self._remote_src()
+
+        async with self._spawning.spawn(
+                tgt,
+                spawn_src,
+                debug=bs.main_config.debug,
+        ) as proc:
+            res = await PyremoteBootstrapDriver(  # noqa
+                remote_src,
+                PyremoteBootstrapOptions(
+                    debug=bs.main_config.debug,
+                ),
+            ).async_run(
+                proc.stdout,
+                proc.stdin,
+            )
+
+            chan = RemoteChannelImpl(
+                proc.stdout,
+                proc.stdin,
+                msh=self._msh,
+            )
+
+            await chan.send_obj(bs)
+
+            rce: RemoteCommandExecutor
+            async with contextlib.aclosing(RemoteCommandExecutor(chan)) as rce:
+                await rce.start()
+
+                yield rce

ominfra/manage/remote/execution.py

@@ -1,123 +1,182 @@
 # ruff: noqa: UP006 UP007
-import contextlib
+import abc
+import asyncio
 import dataclasses as dc
+import itertools
 import logging
-import threading
 import typing as ta
 
-from omlish.lite.cached import cached_nullary
+from omlish.lite.check import check_isinstance
+from omlish.lite.check import check_none
 from omlish.lite.check import check_not_none
+from omlish.lite.check import check_state
 from omlish.lite.logs import log
-from omlish.lite.marshal import ObjMarshalerManager
-from omlish.lite.pycharm import pycharm_debug_connect
-
-from ...pyremote import PyremoteBootstrapDriver
-from ...pyremote import PyremoteBootstrapOptions
-from ...pyremote import pyremote_bootstrap_finalize
-from ...pyremote import pyremote_build_bootstrap_cmd
-from ..bootstrap import MainBootstrap
+
 from ..commands.base import Command
 from ..commands.base import CommandException
 from ..commands.base import CommandExecutor
 from ..commands.base import CommandOutputOrException
 from ..commands.base import CommandOutputOrExceptionData
-from ..commands.execution import LocalCommandExecutor
 from .channel import RemoteChannel
-from .payload import RemoteExecutionPayloadFile
-from .payload import get_remote_payload_src
-from .spawning import RemoteSpawning
 
 
-if ta.TYPE_CHECKING:
-    from ..bootstrap_ import main_bootstrap
-else:
-    main_bootstrap: ta.Any = None
+T = ta.TypeVar('T')
 
 
 ##
 
 
-class _RemoteExecutionLogHandler(logging.Handler):
-    def __init__(self, fn: ta.Callable[[str], None]) -> None:
-        super().__init__()
-        self._fn = fn
+class _RemoteProtocol:
+    class Message(abc.ABC):  # noqa
+        async def send(self, chan: RemoteChannel) -> None:
+            await chan.send_obj(self, _RemoteProtocol.Message)
 
-    def emit(self, record):
-        msg = self.format(record)
-        self._fn(msg)
+        @classmethod
+        async def recv(cls: ta.Type[T], chan: RemoteChannel) -> ta.Optional[T]:
+            return await chan.recv_obj(cls)
 
+    #
 
-@dc.dataclass(frozen=True)
-class _RemoteExecutionRequest:
-    c: Command
+    class Request(Message, abc.ABC):  # noqa
+        pass
 
+    @dc.dataclass(frozen=True)
+    class CommandRequest(Request):
+        seq: int
+        cmd: Command
 
-@dc.dataclass(frozen=True)
-class _RemoteExecutionLog:
-    s: str
+    @dc.dataclass(frozen=True)
+    class PingRequest(Request):
+        time: float
 
+    #
 
-@dc.dataclass(frozen=True)
-class _RemoteExecutionResponse:
-    r: ta.Optional[CommandOutputOrExceptionData] = None
-    l: ta.Optional[_RemoteExecutionLog] = None
+    class Response(Message, abc.ABC):  # noqa
+        pass
 
+    @dc.dataclass(frozen=True)
+    class LogResponse(Response):
+        s: str
 
-def _remote_execution_main() -> None:
-    rt = pyremote_bootstrap_finalize()  # noqa
+    @dc.dataclass(frozen=True)
+    class CommandResponse(Response):
+        seq: int
+        res: CommandOutputOrExceptionData
 
-    chan = RemoteChannel(
-        rt.input,
-        rt.output,
-    )
+    @dc.dataclass(frozen=True)
+    class PingResponse(Response):
+        time: float
 
-    bs = check_not_none(chan.recv_obj(MainBootstrap))
 
-    if (prd := bs.remote_config.pycharm_remote_debug) is not None:
-        pycharm_debug_connect(prd)
+##
 
-    injector = main_bootstrap(bs)
 
-    chan.set_marshaler(injector[ObjMarshalerManager])
+class _RemoteLogHandler(logging.Handler):
+    def __init__(
+            self,
+            chan: RemoteChannel,
+            loop: ta.Any = None,
+    ) -> None:
+        super().__init__()
 
-    #
+        self._chan = chan
+        self._loop = loop
 
-    log_lock = threading.RLock()
-    send_logs = False
+    def emit(self, record):
+        msg = self.format(record)
 
-    def log_fn(s: str) -> None:
-        with log_lock:
-            if send_logs:
-                chan.send_obj(_RemoteExecutionResponse(l=_RemoteExecutionLog(s)))
+        async def inner():
+            await _RemoteProtocol.LogResponse(msg).send(self._chan)
 
-    log_handler = _RemoteExecutionLogHandler(log_fn)
-    logging.root.addHandler(log_handler)
+        loop = self._loop
+        if loop is None:
+            loop = asyncio.get_running_loop()
+        if loop is not None:
+            asyncio.run_coroutine_threadsafe(inner(), loop)
 
-    #
 
-    ce = injector[LocalCommandExecutor]
+##
 
-    while True:
-        req = chan.recv_obj(_RemoteExecutionRequest)
-        if req is None:
-            break
 
-        with log_lock:
-            send_logs = True
+class _RemoteCommandHandler:
+    def __init__(
+            self,
+            chan: RemoteChannel,
+            executor: CommandExecutor,
+            *,
+            stop: ta.Optional[asyncio.Event] = None,
+    ) -> None:
+        super().__init__()
 
-        r = ce.try_execute(
-            req.c,
+        self._chan = chan
+        self._executor = executor
+        self._stop = stop if stop is not None else asyncio.Event()
+
+        self._cmds_by_seq: ta.Dict[int, _RemoteCommandHandler._Command] = {}
+
+    @dc.dataclass(frozen=True)
+    class _Command:
+        req: _RemoteProtocol.CommandRequest
+        fut: asyncio.Future
+
+    async def run(self) -> None:
+        stop_task = asyncio.create_task(self._stop.wait())
+        recv_task: ta.Optional[asyncio.Task] = None
+
+        while not self._stop.is_set():
+            if recv_task is None:
+                recv_task = asyncio.create_task(_RemoteProtocol.Request.recv(self._chan))
+
+            done, pending = await asyncio.wait([
+                stop_task,
+                recv_task,
+            ], return_when=asyncio.FIRST_COMPLETED)
+
+            if recv_task in done:
+                msg: ta.Optional[_RemoteProtocol.Message] = check_isinstance(
+                    recv_task.result(),
+                    (_RemoteProtocol.Message, type(None)),
+                )
+                recv_task = None
+
+                if msg is None:
+                    break
+
+                await self._handle_message(msg)
+
+    async def _handle_message(self, msg: _RemoteProtocol.Message) -> None:
+        if isinstance(msg, _RemoteProtocol.PingRequest):
+            log.debug('Ping: %r', msg)
+            await _RemoteProtocol.PingResponse(
+                time=msg.time,
+            ).send(self._chan)
+
+        elif isinstance(msg, _RemoteProtocol.CommandRequest):
+            fut = asyncio.create_task(self._handle_command_request(msg))
+            self._cmds_by_seq[msg.seq] = _RemoteCommandHandler._Command(
+                req=msg,
+                fut=fut,
+            )
+
+        else:
+            raise TypeError(msg)
+
+    async def _handle_command_request(self, req: _RemoteProtocol.CommandRequest) -> None:
+        res = await self._executor.try_execute(
+            req.cmd,
             log=log,
             omit_exc_object=True,
         )
 
-        with log_lock:
-            send_logs = False
+        await _RemoteProtocol.CommandResponse(
+            seq=req.seq,
+            res=CommandOutputOrExceptionData(
+                output=res.output,
+                exception=res.exception,
+            ),
+        ).send(self._chan)
 
-        chan.send_obj(_RemoteExecutionResponse(r=CommandOutputOrExceptionData(
-            output=r.output,
-            exception=r.exception,
-        )))
+        self._cmds_by_seq.pop(req.seq)  # noqa
 
 
 ##
@@ -134,29 +193,129 @@ class RemoteCommandExecutor(CommandExecutor):
 
         self._chan = chan
 
-    def _remote_execute(self, cmd: Command) -> CommandOutputOrException:
-        self._chan.send_obj(_RemoteExecutionRequest(cmd))
+        self._cmd_seq = itertools.count()
+        self._queue: asyncio.Queue = asyncio.Queue()  # asyncio.Queue[RemoteCommandExecutor._Request]
+        self._stop = asyncio.Event()
+        self._loop_task: ta.Optional[asyncio.Task] = None
+        self._reqs_by_seq: ta.Dict[int, RemoteCommandExecutor._Request] = {}
+
+    #
+
+    async def start(self) -> None:
+        check_none(self._loop_task)
+        check_state(not self._stop.is_set())
+        self._loop_task = asyncio.create_task(self._loop())
+
+    async def aclose(self) -> None:
+        self._stop.set()
+        if self._loop_task is not None:
+            await self._loop_task
 
-        while True:
-            if (r := self._chan.recv_obj(_RemoteExecutionResponse)) is None:
-                raise EOFError
+    #
+
+    @dc.dataclass(frozen=True)
+    class _Request:
+        seq: int
+        cmd: Command
+        fut: asyncio.Future
+
+    async def _loop(self) -> None:
+        log.debug('RemoteCommandExecutor loop start: %r', self)
+
+        stop_task = asyncio.create_task(self._stop.wait())
+        queue_task: ta.Optional[asyncio.Task] = None
+        recv_task: ta.Optional[asyncio.Task] = None
+
+        while not self._stop.is_set():
+            if queue_task is None:
+                queue_task = asyncio.create_task(self._queue.get())
+            if recv_task is None:
+                recv_task = asyncio.create_task(_RemoteProtocol.Message.recv(self._chan))
+
+            done, pending = await asyncio.wait([
+                stop_task,
+                queue_task,
+                recv_task,
+            ], return_when=asyncio.FIRST_COMPLETED)
+
+            if queue_task in done:
+                req = check_isinstance(queue_task.result(), RemoteCommandExecutor._Request)
+                queue_task = None
+                await self._handle_request(req)
+
+            if recv_task in done:
+                msg: ta.Optional[_RemoteProtocol.Message] = check_isinstance(
+                    recv_task.result(),
+                    (_RemoteProtocol.Message, type(None)),
+                )
+                recv_task = None
+
+                if msg is None:
+                    log.debug('RemoteCommandExecutor got eof: %r', self)
+                    break
+
+                await self._handle_message(msg)
+
+        log.debug('RemoteCommandExecutor loop stopping: %r', self)
+
+        for task in [
+            stop_task,
+            queue_task,
+            recv_task,
+        ]:
+            if task is not None and not task.done():
+                task.cancel()
+
+        for req in self._reqs_by_seq.values():
+            req.fut.cancel()
+
+        log.debug('RemoteCommandExecutor loop exited: %r', self)
+
+    async def _handle_request(self, req: _Request) -> None:
+        self._reqs_by_seq[req.seq] = req
+        await _RemoteProtocol.CommandRequest(
+            seq=req.seq,
+            cmd=req.cmd,
+        ).send(self._chan)
+
+    async def _handle_message(self, msg: _RemoteProtocol.Message) -> None:
+        if isinstance(msg, _RemoteProtocol.PingRequest):
+            log.debug('Ping: %r', msg)
+            await _RemoteProtocol.PingResponse(
+                time=msg.time,
+            ).send(self._chan)
+
+        elif isinstance(msg, _RemoteProtocol.LogResponse):
+            log.info(msg.s)
+
+        elif isinstance(msg, _RemoteProtocol.CommandResponse):
+            req = self._reqs_by_seq.pop(msg.seq)
+            req.fut.set_result(msg.res)
+
+        else:
+            raise TypeError(msg)
 
-            if r.l is not None:
-                log.info(r.l.s)
+    #
 
-            if r.r is not None:
-                return r.r
+    async def _remote_execute(self, cmd: Command) -> CommandOutputOrException:
+        req = RemoteCommandExecutor._Request(
+            seq=next(self._cmd_seq),
+            cmd=cmd,
+            fut=asyncio.Future(),
+        )
+        await self._queue.put(req)
+        return await req.fut
 
     # @ta.override
-    def execute(self, cmd: Command) -> Command.Output:
-        r = self._remote_execute(cmd)
+    async def execute(self, cmd: Command) -> Command.Output:
+        r = await self._remote_execute(cmd)
         if (e := r.exception) is not None:
             raise RemoteCommandError(e)
         else:
             return check_not_none(r.output)
 
     # @ta.override
-    def try_execute(
+    async def try_execute(
             self,
             cmd: Command,
             *,
@@ -164,7 +323,7 @@ class RemoteCommandExecutor(CommandExecutor):
             omit_exc_object: bool = False,
     ) -> CommandOutputOrException:
         try:
-            r = self._remote_execute(cmd)
+            r = await self._remote_execute(cmd)
 
         except Exception as e:  # noqa
             if log is not None:
@@ -178,73 +337,3 @@ class RemoteCommandExecutor(CommandExecutor):
 
         else:
             return r
-
-
-##
-
-
-class RemoteExecution:
-    def __init__(
-            self,
-            *,
-            spawning: RemoteSpawning,
-            msh: ObjMarshalerManager,
-            payload_file: ta.Optional[RemoteExecutionPayloadFile] = None,
-    ) -> None:
-        super().__init__()
-
-        self._spawning = spawning
-        self._msh = msh
-        self._payload_file = payload_file
-
-    #
-
-    @cached_nullary
-    def _payload_src(self) -> str:
-        return get_remote_payload_src(file=self._payload_file)
-
-    @cached_nullary
-    def _remote_src(self) -> ta.Sequence[str]:
-        return [
-            self._payload_src(),
-            '_remote_execution_main()',
-        ]
-
-    @cached_nullary
-    def _spawn_src(self) -> str:
-        return pyremote_build_bootstrap_cmd(__package__ or 'manage')
-
-    #
-
-    @contextlib.contextmanager
-    def connect(
-            self,
-            tgt: RemoteSpawning.Target,
-            bs: MainBootstrap,
-    ) -> ta.Generator[RemoteCommandExecutor, None, None]:
-        spawn_src = self._spawn_src()
-        remote_src = self._remote_src()
-
-        with self._spawning.spawn(
-                tgt,
-                spawn_src,
-        ) as proc:
-            res = PyremoteBootstrapDriver(  # noqa
-                remote_src,
-                PyremoteBootstrapOptions(
-                    debug=bs.main_config.debug,
-                ),
-            ).run(
-                proc.stdout,
-                proc.stdin,
-            )
-
-            chan = RemoteChannel(
-                proc.stdout,
-                proc.stdin,
-                msh=self._msh,
-            )
-
-            chan.send_obj(bs)
-
-            yield RemoteCommandExecutor(chan)

ominfra/manage/remote/inject.py

@@ -6,9 +6,11 @@ from omlish.lite.inject import InjectorBindings
 from omlish.lite.inject import inj
 
 from .config import RemoteConfig
-from .execution import RemoteExecution
+from .connection import PyremoteRemoteExecutionConnector
+from .connection import RemoteExecutionConnector
 from .payload import RemoteExecutionPayloadFile
 from .spawning import RemoteSpawning
+from .spawning import SubprocessRemoteSpawning
 
 
 def bind_remote(
@@ -18,9 +20,11 @@ def bind_remote(
     lst: ta.List[InjectorBindingOrBindings] = [
         inj.bind(remote_config),
 
-        inj.bind(RemoteSpawning, singleton=True),
+        inj.bind(SubprocessRemoteSpawning, singleton=True),
+        inj.bind(RemoteSpawning, to_key=SubprocessRemoteSpawning),
 
-        inj.bind(RemoteExecution, singleton=True),
+        inj.bind(PyremoteRemoteExecutionConnector, singleton=True),
+        inj.bind(RemoteExecutionConnector, to_key=PyremoteRemoteExecutionConnector),
     ]
 
     if (pf := remote_config.payload_file) is not None: