ominfra 0.0.0.dev139__py3-none-any.whl → 0.0.0.dev141__py3-none-any.whl

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/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/manage/main.py

@@ -3,25 +3,19 @@
 # ruff: noqa: UP006 UP007
 """
 manage.py -s 'docker run -i python:3.12'
-manage.py -qs 'ssh -i foo/bar foo@bar.baz' --python=python3.8
+manage.py -s 'ssh -i /foo/bar.pem foo@bar.baz' -q --python=python3.8
 """
-import inspect
 import json
-import shlex
 import struct
-import subprocess
-import sys
 import typing as ta
 
-from omlish.lite.cached import cached_nullary
-from omlish.lite.check import check_not_none
+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 omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
 
 from ..pyremote import PyremoteBootstrapDriver
 from ..pyremote import PyremoteBootstrapOptions
@@ -29,6 +23,9 @@ 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
 
 
 ##
@@ -39,29 +36,23 @@ _COMMAND_TYPES = {
 }
 
 
-register_opj_marshaler(
-    Command.Input,
-    PolymorphicObjMarshaler.of([
-        PolymorphicObjMarshaler.Impl(
-            cty.Input,
-            k,
-            get_obj_marshaler(cty.Input),
-        )
-        for k, cty in _COMMAND_TYPES.items()
-    ]),
-)
-
-register_opj_marshaler(
-    Command.Output,
-    PolymorphicObjMarshaler.of([
-        PolymorphicObjMarshaler.Impl(
-            cty.Output,
-            k,
-            get_obj_marshaler(cty.Output),
+@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()
+            ]),
         )
-        for k, cty in _COMMAND_TYPES.items()
-    ]),
-)
 
 
 ##
@@ -99,12 +90,12 @@ def _remote_main() -> None:
     rt = pyremote_bootstrap_finalize()  # noqa
 
     while True:
-        i = _recv_obj(rt.input, Command.Input)
+        i = _recv_obj(rt.input, Command)
         if i is None:
             break
 
-        if isinstance(i, SubprocessCommand.Input):
-            o = SubprocessCommand()._execute(i)  # noqa
+        if isinstance(i, SubprocessCommand):
+            o = SubprocessCommandExecutor().execute(i)  # noqa
         else:
             raise TypeError(i)
 
@@ -114,38 +105,6 @@ def _remote_main() -> None:
 ##
 
 
-@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_amalg_src(*, amalg_file: ta.Optional[str]) -> str:
-    if amalg_file is not None:
-        with open(amalg_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()
-
-
-##
-
-
 def _main() -> None:
     import argparse
 
@@ -154,19 +113,20 @@ def _main() -> None:
     parser.add_argument('-s', '--shell')
     parser.add_argument('-q', '--shell-quote', action='store_true')
     parser.add_argument('--python', default='python3')
-    parser.add_argument('--_amalg-file')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--_payload-file')
 
     args = parser.parse_args()
 
     #
 
-    amalg_src = _get_amalg_src(amalg_file=args._amalg_file)  # noqa
+    payload_src = get_payload_src(file=args._payload_file)  # noqa
 
     #
 
     remote_src = '\n\n'.join([
         '__name__ = "__remote__"',
-        amalg_src,
+        payload_src,
         '_remote_main()',
     ])
 
@@ -174,60 +134,41 @@ def _main() -> None:
 
     bs_src = pyremote_build_bootstrap_cmd(__package__ or 'manage')
 
-    if args.shell is not None:
-        sh_src = f'{args.python} -c {shlex.quote(bs_src)}'
-        if args.shell_quote:
-            sh_src = shlex.quote(sh_src)
-        sh_cmd = f'{args.shell} {sh_src}'
-        cmd = [sh_cmd]
-        shell = True
-    else:
-        cmd = [args.python, '-c', bs_src]
-        shell = False
-
-    proc = subprocess.Popen(
-        subprocess_maybe_shell_wrap_exec(*cmd),
-        shell=shell,
-        stdin=subprocess.PIPE,
-        stdout=subprocess.PIPE,
-    )
-
-    stdin = check_not_none(proc.stdin)
-    stdout = check_not_none(proc.stdout)
-
-    res = PyremoteBootstrapDriver(  # noqa
-        remote_src,
-        PyremoteBootstrapOptions(
-            # debug=True,
-        ),
-    ).run(stdin, stdout)
-    # print(res)
-
     #
 
-    for ci in [
-        SubprocessCommand.Input(
-            args=['python3', '-'],
-            input=b'print(1)\n',
-            capture_stdout=True,
-        ),
-        SubprocessCommand.Input(
-            args=['uname'],
-            capture_stdout=True,
-        ),
-    ]:
-        _send_obj(stdin, ci, Command.Input)
-
-        o = _recv_obj(stdout, Command.Output)
-
-        print(o)
-
-    try:
-        stdin.close()
-    except BrokenPipeError:
-        pass
-
-    proc.wait()
+    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__':

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/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/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/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
 

ominfra/scripts/manage.py

@@ -6,7 +6,7 @@
 # ruff: noqa: N802 UP006 UP007 UP036
 """
 manage.py -s 'docker run -i python:3.12'
-manage.py -qs 'ssh -i foo/bar foo@bar.baz' --python=python3.8
+manage.py -s 'ssh -i /foo/bar.pem foo@bar.baz' -q --python=python3.8
 """
 import abc
 import base64
@@ -49,11 +49,12 @@ if sys.version_info < (3, 8):
 
 
 # commands/base.py
-CommandInputT = ta.TypeVar('CommandInputT', bound='Command.Input')
+CommandT = ta.TypeVar('CommandT', bound='Command')
 CommandOutputT = ta.TypeVar('CommandOutputT', bound='Command.Output')
 
 # ../../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)
@@ -66,17 +67,19 @@ SizedT = ta.TypeVar('SizedT', bound=ta.Sized)
 ##
 
 
-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
 
 
@@ -203,6 +206,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'
@@ -245,11 +250,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,)))
@@ -300,6 +304,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',
@@ -335,14 +340,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')
@@ -366,11 +372,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
@@ -378,12 +387,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)
 
@@ -406,6 +418,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,
@@ -550,6 +563,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
 
@@ -733,6 +752,39 @@ def deep_subclasses(cls: ta.Type[T]) -> ta.Iterator[ta.Type[T]]:
         todo.extend(reversed(cur.__subclasses__()))
 
 
+########################################
+# ../payload.py
+
+
+@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()
+
+
 ########################################
 # ../../../omlish/lite/logs.py
 """
@@ -1498,24 +1550,23 @@ def subprocess_close(
 ##
 
 
-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):
@@ -1527,7 +1578,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),
 
@@ -1558,6 +1614,101 @@ class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Ou
         )
 
 
+########################################
+# ../spawning.py
+
+
+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)
+
+
 ########################################
 # main.py
 
@@ -1570,29 +1721,23 @@ _COMMAND_TYPES = {
 }
 
 
-register_opj_marshaler(
-    Command.Input,
-    PolymorphicObjMarshaler.of([
-        PolymorphicObjMarshaler.Impl(
-            cty.Input,
-            k,
-            get_obj_marshaler(cty.Input),
-        )
-        for k, cty in _COMMAND_TYPES.items()
-    ]),
-)
-
-register_opj_marshaler(
-    Command.Output,
-    PolymorphicObjMarshaler.of([
-        PolymorphicObjMarshaler.Impl(
-            cty.Output,
-            k,
-            get_obj_marshaler(cty.Output),
+@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()
+            ]),
         )
-        for k, cty in _COMMAND_TYPES.items()
-    ]),
-)
 
 
 ##
@@ -1630,12 +1775,12 @@ def _remote_main() -> None:
     rt = pyremote_bootstrap_finalize()  # noqa
 
     while True:
-        i = _recv_obj(rt.input, Command.Input)
+        i = _recv_obj(rt.input, Command)
         if i is None:
             break
 
-        if isinstance(i, SubprocessCommand.Input):
-            o = SubprocessCommand()._execute(i)  # noqa
+        if isinstance(i, SubprocessCommand):
+            o = SubprocessCommandExecutor().execute(i)  # noqa
         else:
             raise TypeError(i)
 
@@ -1645,38 +1790,6 @@ def _remote_main() -> None:
 ##
 
 
-@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_amalg_src(*, amalg_file: ta.Optional[str]) -> str:
-    if amalg_file is not None:
-        with open(amalg_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()
-
-
-##
-
-
 def _main() -> None:
     import argparse
 
@@ -1685,19 +1798,20 @@ def _main() -> None:
     parser.add_argument('-s', '--shell')
     parser.add_argument('-q', '--shell-quote', action='store_true')
     parser.add_argument('--python', default='python3')
-    parser.add_argument('--_amalg-file')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--_payload-file')
 
     args = parser.parse_args()
 
     #
 
-    amalg_src = _get_amalg_src(amalg_file=args._amalg_file)  # noqa
+    payload_src = get_payload_src(file=args._payload_file)  # noqa
 
     #
 
     remote_src = '\n\n'.join([
         '__name__ = "__remote__"',
-        amalg_src,
+        payload_src,
         '_remote_main()',
     ])
 
@@ -1705,60 +1819,41 @@ def _main() -> None:
 
     bs_src = pyremote_build_bootstrap_cmd(__package__ or 'manage')
 
-    if args.shell is not None:
-        sh_src = f'{args.python} -c {shlex.quote(bs_src)}'
-        if args.shell_quote:
-            sh_src = shlex.quote(sh_src)
-        sh_cmd = f'{args.shell} {sh_src}'
-        cmd = [sh_cmd]
-        shell = True
-    else:
-        cmd = [args.python, '-c', bs_src]
-        shell = False
-
-    proc = subprocess.Popen(
-        subprocess_maybe_shell_wrap_exec(*cmd),
-        shell=shell,
-        stdin=subprocess.PIPE,
-        stdout=subprocess.PIPE,
-    )
-
-    stdin = check_not_none(proc.stdin)
-    stdout = check_not_none(proc.stdout)
-
-    res = PyremoteBootstrapDriver(  # noqa
-        remote_src,
-        PyremoteBootstrapOptions(
-            # debug=True,
-        ),
-    ).run(stdin, stdout)
-    # print(res)
-
     #
 
-    for ci in [
-        SubprocessCommand.Input(
-            args=['python3', '-'],
-            input=b'print(1)\n',
-            capture_stdout=True,
-        ),
-        SubprocessCommand.Input(
-            args=['uname'],
-            capture_stdout=True,
-        ),
-    ]:
-        _send_obj(stdin, ci, Command.Input)
-
-        o = _recv_obj(stdout, Command.Output)
-
-        print(o)
+    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)
 
-    try:
-        stdin.close()
-    except BrokenPipeError:
-        pass
+        #
 
-    proc.wait()
+        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__':

ominfra/scripts/supervisor.py

@@ -101,6 +101,7 @@ V = ta.TypeVar('V')
 
 # ../../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)
@@ -1454,6 +1455,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
 

{ominfra-0.0.0.dev139.dist-info → ominfra-0.0.0.dev141.dist-info}/METADATA

@@ -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.dist-info → ominfra-0.0.0.dev141.dist-info}/RECORD

@@ -3,7 +3,7 @@ ominfra/__about__.py,sha256=6i1AoruFYQCd-PyhhbDQDWY2d1tiQu9nkwWr-fXAqfY,705
 ominfra/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ominfra/cmds.py,sha256=E0AfnvEmnKntXWvmLW5L05_NeDpBET1VBXn7vV6EwBQ,2083
 ominfra/configs.py,sha256=8aU1Qmbr-qjaE2iP3gAbA2SWJYMPZ-uGK007L01PoOI,1727
-ominfra/pyremote.py,sha256=vIEUncFgYvcFj3aK0dvJkRuAuaaRyqXk7c8rl2WPOR4,12599
+ominfra/pyremote.py,sha256=vuLEZSjYtil7Knl4qlXYI4dqqiEGJvuO7EIcqr8Kabk,13142
 ominfra/ssh.py,sha256=jQpc4WvkMckIfk4vILda8zFaeharRqc_6wxW50b0OjQ,5431
 ominfra/threadworkers.py,sha256=oX4ubZn7h932saXpRIJu2MNhBExgGGMuGhdXarZxLJw,4948
 ominfra/clouds/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -28,14 +28,16 @@ ominfra/journald/genmessages.py,sha256=rLTS-K2v7otNOtTz4RoOEVYCm0fQuuBzf47e0T61t
 ominfra/journald/messages.py,sha256=Wr7TjWMOySc0WnKwp_-idR8RfeKlbyJQ_KkELr0VB70,2192
 ominfra/journald/tailer.py,sha256=5abcFMfgi7fnY9ZEQe2ZVobaJxjQkeu6d9Kagw33a1w,33525
 ominfra/manage/__init__.py,sha256=YqHcYMfmo04OoVNYc6Ykc1I8WaJot44skwhm9DcnmXE,166
-ominfra/manage/main.py,sha256=D-wxBsGi-EBBpKBC1hhLCr8PkJxLs1Vj9CAsSwisn-k,5141
+ominfra/manage/main.py,sha256=rTY98MrDqGrpwRxo93z5TkfE82Ze3XrL1TqEZ--nkzk,3996
+ominfra/manage/payload.py,sha256=NPuQtMGe1q3kDjG5cgOFmd6M9ITyKRc4IyMRY9PHzRI,818
+ominfra/manage/spawning.py,sha256=EeV4hNcUZOWlM1VuhfGolLHSMWJdz6Dn24OLFtcbisE,2665
 ominfra/manage/commands/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ominfra/manage/commands/base.py,sha256=TTrHL213jf-ClVqToiNHuxQey1Yf6052E8u3E9hAf7Q,574
-ominfra/manage/commands/subprocess.py,sha256=GpbD-cTorgCRg203lCl81HAh-NBYA6ObKa5ZP2ss9rg,1884
+ominfra/manage/commands/base.py,sha256=nXvqvMYeiuya0hcLnVjQNX3jUdIQN6Brt_UGIQBfNvs,559
+ominfra/manage/commands/subprocess.py,sha256=7d-q8hOVaI5YY-dMU0R1KbCIo3R2mlgBSZ2g8NKEn9A,1941
 ominfra/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ominfra/scripts/journald2aws.py,sha256=kCyeyCo53GRoCKAdVAe0aGkgmINxi8IBcUfRXphACek,131618
-ominfra/scripts/manage.py,sha256=xgDqnBWLdLogvu2R5e3Pgg3Dc9M2P_6QN1-7MQE5dnY,43817
-ominfra/scripts/supervisor.py,sha256=3rhnR5gT3MxqFW9ew5ePYioTDLvEcM2nlYmvbGBCvzI,245884
+ominfra/scripts/journald2aws.py,sha256=wALg2habRTCsXEEACYqzB8qIxmZ0gHSL2NoPfu_iKDk,131771
+ominfra/scripts/manage.py,sha256=nKidMTzCGV0Clfaw3-VlAIpy36LK3afr5gGkPUMaUQo,46670
+ominfra/scripts/supervisor.py,sha256=q8d1LsBIthlv2-cTi1vb9z_2AfZ0f8mE-Xym7TLeiLU,246037
 ominfra/supervisor/LICENSE.txt,sha256=yvqaMNsDhWxziHa9ien6qCW1SkZv-DQlAg96XjfSee8,1746
 ominfra/supervisor/__init__.py,sha256=Y3l4WY4JRi2uLG6kgbGp93fuGfkxkKwZDvhsa0Rwgtk,15
 ominfra/supervisor/__main__.py,sha256=I0yFw-C08OOiZ3BF6lF1Oiv789EQXu-_j6whDhQUTEA,66
@@ -77,9 +79,9 @@ ominfra/tailscale/api.py,sha256=C5-t_b6jZXUWcy5k8bXm7CFnk73pSdrlMOgGDeGVrpw,1370
 ominfra/tailscale/cli.py,sha256=DSGp4hn5xwOW-l_u_InKlSF6kIobxtUtVssf_73STs0,3567
 ominfra/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ominfra/tools/listresources.py,sha256=4qVg5txsb10EHhvqXXeM6gJ2jx9LbroEnPydDv1uXs0,6176
-ominfra-0.0.0.dev139.dist-info/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
-ominfra-0.0.0.dev139.dist-info/METADATA,sha256=f4A84661QYB6jEMis-4rQWqm5ZajlCBBJ7y367ojJVM,731
-ominfra-0.0.0.dev139.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-ominfra-0.0.0.dev139.dist-info/entry_points.txt,sha256=kgecQ2MgGrM9qK744BoKS3tMesaC3yjLnl9pa5CRczg,37
-ominfra-0.0.0.dev139.dist-info/top_level.txt,sha256=E-b2OHkk_AOBLXHYZQ2EOFKl-_6uOGd8EjeG-Zy6h_w,8
-ominfra-0.0.0.dev139.dist-info/RECORD,,
+ominfra-0.0.0.dev141.dist-info/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
+ominfra-0.0.0.dev141.dist-info/METADATA,sha256=CRa6olAYdI24dF7R_HKUHw0aebJQzvHwvACjG82Mff4,731
+ominfra-0.0.0.dev141.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+ominfra-0.0.0.dev141.dist-info/entry_points.txt,sha256=kgecQ2MgGrM9qK744BoKS3tMesaC3yjLnl9pa5CRczg,37
+ominfra-0.0.0.dev141.dist-info/top_level.txt,sha256=E-b2OHkk_AOBLXHYZQ2EOFKl-_6uOGd8EjeG-Zy6h_w,8
+ominfra-0.0.0.dev141.dist-info/RECORD,,