ominfra 0.0.0.dev142__py3-none-any.whl → 0.0.0.dev143__py3-none-any.whl

ominfra/manage/remote/execution.py

@@ -0,0 +1,193 @@
+# ruff: noqa: UP006 UP007
+import contextlib
+import dataclasses as dc
+import logging
+import typing as ta
+
+from omlish.lite.cached import cached_nullary
+from omlish.lite.check import check_not_none
+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 .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
+
+
+##
+
+
+def _remote_execution_main() -> None:
+    rt = pyremote_bootstrap_finalize()  # noqa
+
+    chan = RemoteChannel(
+        rt.input,
+        rt.output,
+    )
+
+    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])
+
+    ce = injector[CommandExecutor]
+
+    while True:
+        i = chan.recv_obj(Command)
+        if i is None:
+            break
+
+        r = ce.try_execute(
+            i,
+            log=log,
+            omit_exc_object=True,
+        )
+
+        chan.send_obj(r)
+
+
+##
+
+
+@dc.dataclass()
+class RemoteCommandError(Exception):
+    e: CommandException
+
+
+class RemoteCommandExecutor(CommandExecutor):
+    def __init__(self, chan: RemoteChannel) -> None:
+        super().__init__()
+
+        self._chan = chan
+
+    def _remote_execute(self, cmd: Command) -> CommandOutputOrException:
+        self._chan.send_obj(cmd, Command)
+
+        if (r := self._chan.recv_obj(CommandOutputOrExceptionData)) is None:
+            raise EOFError
+
+        return r
+
+    # @ta.override
+    def execute(self, cmd: Command) -> Command.Output:
+        r = 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(
+            self,
+            cmd: Command,
+            *,
+            log: ta.Optional[logging.Logger] = None,
+            omit_exc_object: bool = False,
+    ) -> CommandOutputOrException:
+        try:
+            r = self._remote_execute(cmd)
+
+        except Exception as e:  # noqa
+            if log is not None:
+                log.exception('Exception executing remote command: %r', type(cmd))
+
+            return CommandOutputOrExceptionData(exception=CommandException.of(
+                e,
+                omit_exc_object=omit_exc_object,
+                cmd=cmd,
+            ))
+
+        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

@@ -0,0 +1,29 @@
+# ruff: noqa: UP006 UP007
+import typing as ta
+
+from omlish.lite.inject import InjectorBindingOrBindings
+from omlish.lite.inject import InjectorBindings
+from omlish.lite.inject import inj
+
+from .config import RemoteConfig
+from .execution import RemoteExecution
+from .payload import RemoteExecutionPayloadFile
+from .spawning import RemoteSpawning
+
+
+def bind_remote(
+        *,
+        remote_config: RemoteConfig,
+) -> InjectorBindings:
+    lst: ta.List[InjectorBindingOrBindings] = [
+        inj.bind(remote_config),
+
+        inj.bind(RemoteSpawning, singleton=True),
+
+        inj.bind(RemoteExecution, singleton=True),
+    ]
+
+    if (pf := remote_config.payload_file) is not None:
+        lst.append(inj.bind(pf, to_key=RemoteExecutionPayloadFile))
+
+    return inj.as_bindings(*lst)

ominfra/manage/{payload.py → remote/payload.py}

@@ -6,6 +6,9 @@ import typing as ta
 from omlish.lite.cached import cached_nullary
 
 
+RemoteExecutionPayloadFile = ta.NewType('RemoteExecutionPayloadFile', str)
+
+
 @cached_nullary
 def _get_self_src() -> str:
     return inspect.getsource(sys.modules[__name__])
@@ -23,7 +26,10 @@ def _is_self_amalg() -> bool:
     return _is_src_amalg(_get_self_src())
 
 
-def get_payload_src(*, file: ta.Optional[str]) -> str:
+def get_remote_payload_src(
+        *,
+        file: ta.Optional[RemoteExecutionPayloadFile],
+) -> str:
     if file is not None:
         with open(file) as f:
             return f.read()

ominfra/manage/{spawning.py → remote/spawning.py}

@@ -11,25 +11,16 @@ from omlish.lite.subprocesses import SubprocessChannelOption
 from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
 
 
-class PySpawner:
-    DEFAULT_PYTHON = 'python3'
+class RemoteSpawning:
+    @dc.dataclass(frozen=True)
+    class Target:
+        shell: ta.Optional[str] = None
+        shell_quote: bool = False
 
-    def __init__(
-            self,
-            src: str,
-            *,
-            shell: ta.Optional[str] = None,
-            shell_quote: bool = False,
-            python: str = DEFAULT_PYTHON,
-            stderr: ta.Optional[SubprocessChannelOption] = None,
-    ) -> None:
-        super().__init__()
-
-        self._src = src
-        self._shell = shell
-        self._shell_quote = shell_quote
-        self._python = python
-        self._stderr = stderr
+        DEFAULT_PYTHON: ta.ClassVar[str] = 'python3'
+        python: str = DEFAULT_PYTHON
+
+        stderr: ta.Optional[str] = None  # SubprocessChannelOption
 
     #
 
@@ -37,20 +28,24 @@ class PySpawner:
         cmd: ta.Sequence[str]
         shell: bool
 
-    def _prepare_cmd(self) -> _PreparedCmd:
-        if self._shell is not None:
-            sh_src = f'{self._python} -c {shlex.quote(self._src)}'
-            if self._shell_quote:
+    def _prepare_cmd(
+            self,
+            tgt: Target,
+            src: str,
+    ) -> _PreparedCmd:
+        if tgt.shell is not None:
+            sh_src = f'{tgt.python} -c {shlex.quote(src)}'
+            if tgt.shell_quote:
                 sh_src = shlex.quote(sh_src)
-            sh_cmd = f'{self._shell} {sh_src}'
-            return PySpawner._PreparedCmd(
+            sh_cmd = f'{tgt.shell} {sh_src}'
+            return RemoteSpawning._PreparedCmd(
                 cmd=[sh_cmd],
                 shell=True,
             )
 
         else:
-            return PySpawner._PreparedCmd(
-                cmd=[self._python, '-c', self._src],
+            return RemoteSpawning._PreparedCmd(
+                cmd=[tgt.python, '-c', src],
                 shell=False,
             )
 
@@ -65,23 +60,28 @@ class PySpawner:
     @contextlib.contextmanager
     def spawn(
             self,
+            tgt: Target,
+            src: str,
             *,
             timeout: ta.Optional[float] = None,
     ) -> ta.Generator[Spawned, None, None]:
-        pc = self._prepare_cmd()
+        pc = self._prepare_cmd(tgt, src)
 
         with subprocess.Popen(
             subprocess_maybe_shell_wrap_exec(*pc.cmd),
             shell=pc.shell,
             stdin=subprocess.PIPE,
             stdout=subprocess.PIPE,
-            stderr=SUBPROCESS_CHANNEL_OPTION_VALUES[self._stderr] if self._stderr is not None else None,
+            stderr=(
+                SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, tgt.stderr)]
+                if tgt.stderr is not None else None
+            ),
         ) as proc:
             stdin = check_not_none(proc.stdin)
             stdout = check_not_none(proc.stdout)
 
             try:
-                yield PySpawner.Spawned(
+                yield RemoteSpawning.Spawned(
                     stdin=stdin,
                     stdout=stdout,
                     stderr=proc.stderr,

ominfra/pyremote.py

@@ -2,6 +2,9 @@
 # @omlish-lite
 """
 Basically this: https://mitogen.networkgenomics.com/howitworks.html
+
+TODO:
+ - log: ta.Optional[logging.Logger] = None + log.debug's
 """
 import base64
 import dataclasses as dc
@@ -23,6 +26,9 @@ import zlib
 class PyremoteBootstrapOptions:
     debug: bool = False
 
+    DEFAULT_MAIN_NAME_OVERRIDE: ta.ClassVar[str] = '__pyremote__'
+    main_name_override: ta.Optional[str] = DEFAULT_MAIN_NAME_OVERRIDE
+
 
 ##
 
@@ -339,6 +345,10 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
     os.dup2(nfd := os.open('/dev/null', os.O_WRONLY), 1)
     os.close(nfd)
 
+    if (mn := options.main_name_override) is not None:
+        # Inspections like typing.get_type_hints need an entry in sys.modules.
+        sys.modules[mn] = sys.modules['__main__']
+
     # Write fourth ack
     output.write(_PYREMOTE_BOOTSTRAP_ACK3)
 
@@ -357,14 +367,41 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
 
 
 class PyremoteBootstrapDriver:
-    def __init__(self, main_src: str, options: PyremoteBootstrapOptions = PyremoteBootstrapOptions()) -> None:
+    def __init__(
+            self,
+            main_src: ta.Union[str, ta.Sequence[str]],
+            options: PyremoteBootstrapOptions = PyremoteBootstrapOptions(),
+    ) -> None:
         super().__init__()
 
         self._main_src = main_src
-        self._main_z = zlib.compress(main_src.encode('utf-8'))
-
         self._options = options
+
+        self._prepared_main_src = self._prepare_main_src(main_src, options)
+        self._main_z = zlib.compress(self._prepared_main_src.encode('utf-8'))
+
         self._options_json = json.dumps(dc.asdict(options), indent=None, separators=(',', ':')).encode('utf-8')  # noqa
+    #
+
+    @classmethod
+    def _prepare_main_src(
+            cls,
+            main_src: ta.Union[str, ta.Sequence[str]],
+            options: PyremoteBootstrapOptions,
+    ) -> str:
+        parts: ta.List[str]
+        if isinstance(main_src, str):
+            parts = [main_src]
+        else:
+            parts = list(main_src)
+
+        if (mn := options.main_name_override) is not None:
+            parts.insert(0, f'__name__ = {mn!r}')
+
+        if len(parts) == 1:
+            return parts[0]
+        else:
+            return '\n\n'.join(parts)
 
     #