ominfra 0.0.0.dev139__tar.gz → 0.0.0.dev141__tar.gz

{ominfra-0.0.0.dev139/ominfra.egg-info → ominfra-0.0.0.dev141}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ominfra
-Version: 0.0.0.dev139
+Version: 0.0.0.dev141
 Summary: ominfra
 Author: wrmsr
 License: BSD-3-Clause
@@ -12,8 +12,8 @@ Classifier: Operating System :: OS Independent
 Classifier: Operating System :: POSIX
 Requires-Python: >=3.12
 License-File: LICENSE
-Requires-Dist: omdev==0.0.0.dev139
-Requires-Dist: omlish==0.0.0.dev139
+Requires-Dist: omdev==0.0.0.dev141
+Requires-Dist: omlish==0.0.0.dev141
 Provides-Extra: all
 Requires-Dist: paramiko~=3.5; extra == "all"
 Requires-Dist: asyncssh~=2.18; extra == "all"

{ominfra-0.0.0.dev139 → ominfra-0.0.0.dev141}/ominfra/manage/commands/base.py

@@ -4,22 +4,24 @@ import dataclasses as dc
 import typing as ta
 
 
-CommandInputT = ta.TypeVar('CommandInputT', bound='Command.Input')
+CommandT = ta.TypeVar('CommandT', bound='Command')
 CommandOutputT = ta.TypeVar('CommandOutputT', bound='Command.Output')
 
 
 ##
 
 
-class Command(abc.ABC, ta.Generic[CommandInputT, CommandOutputT]):
-    @dc.dataclass(frozen=True)
-    class Input(abc.ABC):  # noqa
-        pass
-
+@dc.dataclass(frozen=True)
+class Command(abc.ABC, ta.Generic[CommandOutputT]):
     @dc.dataclass(frozen=True)
     class Output(abc.ABC):  # noqa
         pass
 
+
+##
+
+
+class CommandExecutor(abc.ABC, ta.Generic[CommandT, CommandOutputT]):
     @abc.abstractmethod
-    def _execute(self, inp: CommandInputT) -> CommandOutputT:
+    def execute(self, i: CommandT) -> CommandOutputT:
         raise NotImplementedError

{ominfra-0.0.0.dev139 → ominfra-0.0.0.dev141}/ominfra/manage/commands/subprocess.py

@@ -8,29 +8,29 @@ import typing as ta
 from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
 
 from .base import Command
+from .base import CommandExecutor
 
 
 ##
 
 
-class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Output']):
-    @dc.dataclass(frozen=True)
-    class Input(Command.Input):
-        args: ta.Sequence[str]
+@dc.dataclass(frozen=True)
+class SubprocessCommand(Command['SubprocessCommand.Output']):
+    args: ta.Sequence[str]
 
-        shell: bool = False
-        cwd: ta.Optional[str] = None
-        env: ta.Optional[ta.Mapping[str, str]] = None
+    shell: bool = False
+    cwd: ta.Optional[str] = None
+    env: ta.Optional[ta.Mapping[str, str]] = None
 
-        capture_stdout: bool = False
-        capture_stderr: bool = False
+    capture_stdout: bool = False
+    capture_stderr: bool = False
 
-        input: ta.Optional[bytes] = None
-        timeout: ta.Optional[float] = None
+    input: ta.Optional[bytes] = None
+    timeout: ta.Optional[float] = None
 
-        def __post_init__(self) -> None:
-            if isinstance(self.args, str):
-                raise TypeError(self.args)
+    def __post_init__(self) -> None:
+        if isinstance(self.args, str):
+            raise TypeError(self.args)
 
     @dc.dataclass(frozen=True)
     class Output(Command.Output):
@@ -42,7 +42,12 @@ class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Ou
         stdout: ta.Optional[bytes] = None
         stderr: ta.Optional[bytes] = None
 
-    def _execute(self, inp: Input) -> Output:
+
+##
+
+
+class SubprocessCommandExecutor(CommandExecutor[SubprocessCommand, SubprocessCommand.Output]):
+    def execute(self, inp: SubprocessCommand) -> SubprocessCommand.Output:
         proc = subprocess.Popen(
             subprocess_maybe_shell_wrap_exec(*inp.args),
 

ominfra-0.0.0.dev141/ominfra/manage/main.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+# @omlish-amalg ../scripts/manage.py
+# ruff: noqa: UP006 UP007
+"""
+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 json
+import struct
+import typing as ta
+
+from omlish.lite.cached import static_init
+from omlish.lite.json import json_dumps_compact
+from omlish.lite.marshal import PolymorphicObjMarshaler
+from omlish.lite.marshal import get_obj_marshaler
+from omlish.lite.marshal import marshal_obj
+from omlish.lite.marshal import register_opj_marshaler
+from omlish.lite.marshal import unmarshal_obj
+
+from ..pyremote import PyremoteBootstrapDriver
+from ..pyremote import PyremoteBootstrapOptions
+from ..pyremote import pyremote_bootstrap_finalize
+from ..pyremote import pyremote_build_bootstrap_cmd
+from .commands.base import Command
+from .commands.subprocess import SubprocessCommand
+from .commands.subprocess import SubprocessCommandExecutor
+from .payload import get_payload_src
+from .spawning import PySpawner
+
+
+##
+
+
+_COMMAND_TYPES = {
+    'subprocess': SubprocessCommand,
+}
+
+
+@static_init
+def _register_command_marshaling() -> None:
+    for fn in [
+        lambda c: c,
+        lambda c: c.Output,
+    ]:
+        register_opj_marshaler(
+            fn(Command),
+            PolymorphicObjMarshaler.of([
+                PolymorphicObjMarshaler.Impl(
+                    fn(cty),
+                    k,
+                    get_obj_marshaler(fn(cty)),
+                )
+                for k, cty in _COMMAND_TYPES.items()
+            ]),
+        )
+
+
+##
+
+
+def _send_obj(f: ta.IO, o: ta.Any, ty: ta.Any = None) -> None:
+    j = json_dumps_compact(marshal_obj(o, ty))
+    d = j.encode('utf-8')
+
+    f.write(struct.pack('<I', len(d)))
+    f.write(d)
+    f.flush()
+
+
+def _recv_obj(f: ta.IO, ty: ta.Any) -> ta.Any:
+    d = f.read(4)
+    if not d:
+        return None
+    if len(d) != 4:
+        raise EOFError
+
+    sz = struct.unpack('<I', d)[0]
+    d = f.read(sz)
+    if len(d) != sz:
+        raise EOFError
+
+    j = json.loads(d.decode('utf-8'))
+    return unmarshal_obj(j, ty)
+
+
+##
+
+
+def _remote_main() -> None:
+    rt = pyremote_bootstrap_finalize()  # noqa
+
+    while True:
+        i = _recv_obj(rt.input, Command)
+        if i is None:
+            break
+
+        if isinstance(i, SubprocessCommand):
+            o = SubprocessCommandExecutor().execute(i)  # noqa
+        else:
+            raise TypeError(i)
+
+        _send_obj(rt.output, o, Command.Output)
+
+
+##
+
+
+def _main() -> None:
+    import argparse
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-s', '--shell')
+    parser.add_argument('-q', '--shell-quote', action='store_true')
+    parser.add_argument('--python', default='python3')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--_payload-file')
+
+    args = parser.parse_args()
+
+    #
+
+    payload_src = get_payload_src(file=args._payload_file)  # noqa
+
+    #
+
+    remote_src = '\n\n'.join([
+        '__name__ = "__remote__"',
+        payload_src,
+        '_remote_main()',
+    ])
+
+    #
+
+    bs_src = pyremote_build_bootstrap_cmd(__package__ or 'manage')
+
+    #
+
+    spawner = PySpawner(
+        bs_src,
+        shell=args.shell,
+        shell_quote=args.shell_quote,
+        python=args.python,
+    )
+    with spawner.spawn() as proc:
+        res = PyremoteBootstrapDriver(  # noqa
+            remote_src,
+            PyremoteBootstrapOptions(
+                debug=args.debug,
+            ),
+        ).run(proc.stdin, proc.stdout)
+        # print(res)
+
+        #
+
+        for ci in [
+            SubprocessCommand(
+                args=['python3', '-'],
+                input=b'print(1)\n',
+                capture_stdout=True,
+            ),
+            SubprocessCommand(
+                args=['uname'],
+                capture_stdout=True,
+            ),
+        ]:
+            _send_obj(proc.stdin, ci, Command)
+
+            o = _recv_obj(proc.stdout, Command.Output)
+
+            print(o)
+
+
+if __name__ == '__main__':
+    _main()

ominfra-0.0.0.dev141/ominfra/manage/payload.py

@@ -0,0 +1,35 @@
+# ruff: noqa: UP006 UP007
+import inspect
+import sys
+import typing as ta
+
+from omlish.lite.cached import cached_nullary
+
+
+@cached_nullary
+def _get_self_src() -> str:
+    return inspect.getsource(sys.modules[__name__])
+
+
+def _is_src_amalg(src: str) -> bool:
+    for l in src.splitlines():  # noqa
+        if l.startswith('# @omlish-amalg-output '):
+            return True
+    return False
+
+
+@cached_nullary
+def _is_self_amalg() -> bool:
+    return _is_src_amalg(_get_self_src())
+
+
+def get_payload_src(*, file: ta.Optional[str]) -> str:
+    if file is not None:
+        with open(file) as f:
+            return f.read()
+
+    if _is_self_amalg():
+        return _get_self_src()
+
+    import importlib.resources
+    return importlib.resources.files(__package__.split('.')[0] + '.scripts').joinpath('manage.py').read_text()

ominfra-0.0.0.dev141/ominfra/manage/spawning.py

@@ -0,0 +1,100 @@
+# ruff: noqa: UP006 UP007
+import contextlib
+import dataclasses as dc
+import shlex
+import subprocess
+import typing as ta
+
+from omlish.lite.check import check_not_none
+from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
+
+
+class PySpawner:
+    DEFAULT_PYTHON = 'python3'
+
+    def __init__(
+            self,
+            src: str,
+            *,
+            shell: ta.Optional[str] = None,
+            shell_quote: bool = False,
+            python: str = DEFAULT_PYTHON,
+            stderr: ta.Optional[ta.Literal['pipe', 'stdout', 'devnull']] = None,
+    ) -> None:
+        super().__init__()
+
+        self._src = src
+        self._shell = shell
+        self._shell_quote = shell_quote
+        self._python = python
+        self._stderr = stderr
+
+    #
+
+    class _PreparedCmd(ta.NamedTuple):
+        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:
+                sh_src = shlex.quote(sh_src)
+            sh_cmd = f'{self._shell} {sh_src}'
+            return PySpawner._PreparedCmd(
+                cmd=[sh_cmd],
+                shell=True,
+            )
+
+        else:
+            return PySpawner._PreparedCmd(
+                cmd=[self._python, '-c', self._src],
+                shell=False,
+            )
+
+    #
+
+    _STDERR_KWARG_MAP: ta.Mapping[str, int] = {
+        'pipe': subprocess.PIPE,
+        'stdout': subprocess.STDOUT,
+        'devnull': subprocess.DEVNULL,
+    }
+
+    @dc.dataclass(frozen=True)
+    class Spawned:
+        stdin: ta.IO
+        stdout: ta.IO
+        stderr: ta.Optional[ta.IO]
+
+    @contextlib.contextmanager
+    def spawn(
+            self,
+            *,
+            timeout: ta.Optional[float] = None,
+    ) -> ta.Generator[Spawned, None, None]:
+        pc = self._prepare_cmd()
+
+        with subprocess.Popen(
+            subprocess_maybe_shell_wrap_exec(*pc.cmd),
+            shell=pc.shell,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=self._STDERR_KWARG_MAP[self._stderr] if self._stderr is not None else None,
+        ) as proc:
+            stdin = check_not_none(proc.stdin)
+            stdout = check_not_none(proc.stdout)
+
+            try:
+                yield PySpawner.Spawned(
+                    stdin=stdin,
+                    stdout=stdout,
+                    stderr=proc.stderr,
+                )
+
+            finally:
+                try:
+                    stdin.close()
+                except BrokenPipeError:
+                    pass
+
+                proc.wait(timeout)

{ominfra-0.0.0.dev139 → ominfra-0.0.0.dev141}/ominfra/pyremote.py

@@ -132,6 +132,8 @@ _PYREMOTE_BOOTSTRAP_SRC_FD = 101
 
 _PYREMOTE_BOOTSTRAP_CHILD_PID_VAR = '_OPYR_CHILD_PID'
 _PYREMOTE_BOOTSTRAP_ARGV0_VAR = '_OPYR_ARGV0'
+_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR = '_OPYR_CONTEXT_NAME'
+_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR = '_OPYR_SRC_FILE'
 _PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR = '_OPYR_OPTIONS_JSON'
 
 _PYREMOTE_BOOTSTRAP_ACK0 = b'OPYR000\n'
@@ -174,11 +176,10 @@ def _pyremote_bootstrap_main(context_name: str) -> None:
         for f in [r0, w0, r1, w1]:
             os.close(f)
 
-        # Save child pid to close after relaunch
+        # Save vars
         os.environ[_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR] = str(cp)
-
-        # Save original argv0
         os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR] = sys.executable
+        os.environ[_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR] = context_name
 
         # Start repl reading stdin from r0
         os.execl(sys.executable, sys.executable + (_PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT % (context_name,)))
@@ -229,6 +230,7 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
 
         '_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR',
         '_PYREMOTE_BOOTSTRAP_ARGV0_VAR',
+        '_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR',
 
         '_PYREMOTE_BOOTSTRAP_ACK0',
         '_PYREMOTE_BOOTSTRAP_ACK1',
@@ -264,14 +266,15 @@ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
 class PyremotePayloadRuntime:
     input: ta.BinaryIO
     output: ta.BinaryIO
+    context_name: str
     main_src: str
     options: PyremoteBootstrapOptions
     env_info: PyremoteEnvInfo
 
 
 def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
-    # If json options var is not present we need to do initial finalization
-    if _PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR not in os.environ:
+    # If src file var is not present we need to do initial finalization
+    if _PYREMOTE_BOOTSTRAP_SRC_FILE_VAR not in os.environ:
         # Read second copy of main src
         r1 = os.fdopen(_PYREMOTE_BOOTSTRAP_SRC_FD, 'rb', 0)
         main_src = r1.read().decode('utf-8')
@@ -295,11 +298,14 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
             os.write(tfd, main_src.encode('utf-8'))
             os.close(tfd)
 
-            # Set json options var
+            # Set vars
+            os.environ[_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR] = tfn
             os.environ[_PYREMOTE_BOOTSTRAP_OPTIONS_JSON_VAR] = options_json.decode('utf-8')
 
             # Re-exec temp file
-            os.execl(os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR], sys.orig_argv[0], tfn)
+            exe = os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR]
+            context_name = os.environ[_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR]
+            os.execl(exe, exe + (_PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT % (context_name,)), tfn)
 
     else:
         # Load options json var
@@ -307,12 +313,15 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
         options = PyremoteBootstrapOptions(**json.loads(options_json_str))
 
         # Read temp source file
-        with open(sys.orig_argv[1]) as sf:
+        with open(os.environ.pop(_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR)) as sf:
             main_src = sf.read()
 
     # Restore original argv0
     sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
 
+    # Grab context name
+    context_name = os.environ.pop(_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR)
+
     # Write third ack
     os.write(1, _PYREMOTE_BOOTSTRAP_ACK2)
 
@@ -335,6 +344,7 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
     return PyremotePayloadRuntime(
         input=input,
         output=output,
+        context_name=context_name,
         main_src=main_src,
         options=options,
         env_info=env_info,

{ominfra-0.0.0.dev139 → ominfra-0.0.0.dev141}/ominfra/scripts/journald2aws.py

@@ -58,6 +58,7 @@ TomlPos = int  # ta.TypeAlias
 
 # ../../../../omlish/lite/cached.py
 T = ta.TypeVar('T')
+CallableT = ta.TypeVar('CallableT', bound=ta.Callable)
 
 # ../../../../omlish/lite/check.py
 SizedT = ta.TypeVar('SizedT', bound=ta.Sized)
@@ -915,6 +916,12 @@ def cached_nullary(fn):  # ta.Callable[..., T]) -> ta.Callable[..., T]:
     return _cached_nullary(fn)
 
 
+def static_init(fn: CallableT) -> CallableT:
+    fn = cached_nullary(fn)
+    fn()
+    return fn
+
+
 ########################################
 # ../../../../../omlish/lite/check.py