ominfra 0.0.0.dev152__py3-none-any.whl → 0.0.0.dev153__py3-none-any.whl

ominfra/manage/bootstrap.py

@@ -2,6 +2,7 @@ import dataclasses as dc
 
 from .config import MainConfig
 from .remote.config import RemoteConfig
+from .system.config import SystemConfig
 
 
 @dc.dataclass(frozen=True)
@@ -9,3 +10,5 @@ class MainBootstrap:
     main_config: MainConfig = MainConfig()
 
     remote_config: RemoteConfig = RemoteConfig()
+
+    system_config: SystemConfig = SystemConfig()

ominfra/manage/bootstrap_.py

@@ -13,6 +13,7 @@ def main_bootstrap(bs: MainBootstrap) -> Injector:
     injector = inj.create_injector(bind_main(  # noqa
         main_config=bs.main_config,
         remote_config=bs.remote_config,
+        system_config=bs.system_config,
     ))
 
     return injector

ominfra/manage/commands/interp.py

@@ -25,9 +25,6 @@ class InterpCommand(Command['InterpCommand.Output']):
         opts: InterpOpts
 
 
-##
-
-
 class InterpCommandExecutor(CommandExecutor[InterpCommand, InterpCommand.Output]):
     async def execute(self, cmd: InterpCommand) -> InterpCommand.Output:
         i = InterpSpecifier.parse(check.not_none(cmd.spec))

ominfra/manage/commands/subprocess.py

@@ -48,9 +48,6 @@ class SubprocessCommand(Command['SubprocessCommand.Output']):
         stderr: ta.Optional[bytes] = None
 
 
-##
-
-
 class SubprocessCommandExecutor(CommandExecutor[SubprocessCommand, SubprocessCommand.Output]):
     async def execute(self, cmd: SubprocessCommand) -> SubprocessCommand.Output:
         proc: asyncio.subprocess.Process

ominfra/manage/deploy/{command.py → commands.py}

@@ -17,9 +17,6 @@ class DeployCommand(Command['DeployCommand.Output']):
         pass
 
 
-##
-
-
 class DeployCommandExecutor(CommandExecutor[DeployCommand, DeployCommand.Output]):
     async def execute(self, cmd: DeployCommand) -> DeployCommand.Output:
         log.info('Deploying!')

ominfra/manage/deploy/inject.py

@@ -6,8 +6,8 @@ from omlish.lite.inject import InjectorBindings
 from omlish.lite.inject import inj
 
 from ..commands.inject import bind_command
-from .command import DeployCommand
-from .command import DeployCommandExecutor
+from .commands import DeployCommand
+from .commands import DeployCommandExecutor
 
 
 def bind_deploy(

ominfra/manage/inject.py

@@ -13,6 +13,8 @@ from .marshal import ObjMarshalerInstaller
 from .marshal import ObjMarshalerInstallers
 from .remote.config import RemoteConfig
 from .remote.inject import bind_remote
+from .system.config import SystemConfig
+from .system.inject import bind_system
 
 
 ##
@@ -22,6 +24,7 @@ def bind_main(
         *,
         main_config: MainConfig,
         remote_config: RemoteConfig,
+        system_config: SystemConfig,
 ) -> InjectorBindings:
     lst: ta.List[InjectorBindingOrBindings] = [
         inj.bind(main_config),
@@ -30,11 +33,15 @@ def bind_main(
             main_config=main_config,
         ),
 
+        bind_deploy(),
+
         bind_remote(
             remote_config=remote_config,
         ),
 
-        bind_deploy(),
+        bind_system(
+            system_config=system_config,
+        ),
     ]
 
     #

ominfra/manage/main.py

@@ -68,6 +68,8 @@ class MainCli(ArgparseCli):
                 ) if self.args.pycharm_debug_port is not None else None,
 
                 timebomb_delay_s=self.args.remote_timebomb_delay_s,
+
+                use_in_process_remote_executor=True,
             ),
         )
 

ominfra/manage/remote/_main.py

@@ -8,8 +8,8 @@ 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.asyncs.asyncio.streams import asyncio_open_stream_reader
+from omlish.asyncs.asyncio.streams import asyncio_open_stream_writer
 from omlish.lite.cached import cached_nullary
 from omlish.lite.check import check
 from omlish.lite.inject import Injector

ominfra/manage/remote/config.py

@@ -20,3 +20,5 @@ class RemoteConfig:
     timebomb_delay_s: ta.Optional[float] = 60 * 60.
 
     heartbeat_interval_s: float = 3.
+
+    use_in_process_remote_executor: bool = False

ominfra/manage/remote/connection.py

@@ -1,8 +1,10 @@
 # ruff: noqa: UP006 UP007
 import abc
+import asyncio
 import contextlib
 import typing as ta
 
+from omlish.asyncs.asyncio.channels import asyncio_create_bytes_channel
 from omlish.lite.cached import cached_nullary
 from omlish.lite.marshal import ObjMarshalerManager
 
@@ -10,9 +12,11 @@ from ...pyremote import PyremoteBootstrapDriver
 from ...pyremote import PyremoteBootstrapOptions
 from ...pyremote import pyremote_build_bootstrap_cmd
 from ..bootstrap import MainBootstrap
+from ..commands.execution import LocalCommandExecutor
 from ._main import _remote_execution_main  # noqa
 from .channel import RemoteChannelImpl
 from .execution import RemoteCommandExecutor
+from .execution import _RemoteCommandHandler
 from .payload import RemoteExecutionPayloadFile
 from .payload import get_remote_payload_src
 from .spawning import RemoteSpawning
@@ -104,3 +108,47 @@ class PyremoteRemoteExecutionConnector(RemoteExecutionConnector):
                 await rce.start()
 
                 yield rce
+
+
+##
+
+
+class InProcessRemoteExecutionConnector(RemoteExecutionConnector):
+    def __init__(
+            self,
+            *,
+            msh: ObjMarshalerManager,
+            local_executor: LocalCommandExecutor,
+    ) -> None:
+        super().__init__()
+
+        self._msh = msh
+        self._local_executor = local_executor
+
+    @contextlib.asynccontextmanager
+    async def connect(
+            self,
+            tgt: RemoteSpawning.Target,
+            bs: MainBootstrap,
+    ) -> ta.AsyncGenerator[RemoteCommandExecutor, None]:
+        r0, w0 = asyncio_create_bytes_channel()
+        r1, w1 = asyncio_create_bytes_channel()
+
+        remote_chan = RemoteChannelImpl(r0, w1, msh=self._msh)
+        local_chan = RemoteChannelImpl(r1, w0, msh=self._msh)
+
+        rch = _RemoteCommandHandler(
+            remote_chan,
+            self._local_executor,
+        )
+        rch_task = asyncio.create_task(rch.run())  # noqa
+        try:
+            rce: RemoteCommandExecutor
+            async with contextlib.aclosing(RemoteCommandExecutor(local_chan)) as rce:
+                await rce.start()
+
+                yield rce
+
+        finally:
+            rch.stop()
+            await rch_task

ominfra/manage/remote/execution.py

@@ -1,9 +1,14 @@
 # ruff: noqa: UP006 UP007
+"""
+TODO:
+ - sequence all messages
+"""
 import abc
 import asyncio
 import dataclasses as dc
 import itertools
 import logging
+import time
 import typing as ta
 
 from omlish.lite.check import check
@@ -96,38 +101,80 @@ class _RemoteLogHandler(logging.Handler):
 
 
 class _RemoteCommandHandler:
+    DEFAULT_PING_INTERVAL_S: float = 3.
+
     def __init__(
             self,
             chan: RemoteChannel,
             executor: CommandExecutor,
             *,
             stop: ta.Optional[asyncio.Event] = None,
+            ping_interval_s: float = DEFAULT_PING_INTERVAL_S,
     ) -> None:
         super().__init__()
 
         self._chan = chan
         self._executor = executor
         self._stop = stop if stop is not None else asyncio.Event()
+        self._ping_interval_s = ping_interval_s
 
         self._cmds_by_seq: ta.Dict[int, _RemoteCommandHandler._Command] = {}
 
+        self._last_ping_send: float = 0.
+        self._ping_in_flight: bool = False
+        self._last_ping_recv: ta.Optional[float] = None
+
+    def stop(self) -> None:
+        self._stop.set()
+
     @dc.dataclass(frozen=True)
     class _Command:
         req: _RemoteProtocol.CommandRequest
         fut: asyncio.Future
 
     async def run(self) -> None:
+        log.debug('_RemoteCommandHandler loop start: %r', self)
+
         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))
+                recv_task = asyncio.create_task(_RemoteProtocol.Message.recv(self._chan))
 
-            done, pending = await asyncio.wait([
-                stop_task,
-                recv_task,
-            ], return_when=asyncio.FIRST_COMPLETED)
+            if not self._ping_in_flight:
+                if not self._last_ping_recv:
+                    ping_wait_time = 0.
+                else:
+                    ping_wait_time = self._ping_interval_s - (time.time() - self._last_ping_recv)
+            else:
+                ping_wait_time = float('inf')
+            wait_time = min(self._ping_interval_s, ping_wait_time)
+            log.debug('_RemoteCommandHandler loop wait: %f', wait_time)
+
+            done, pending = await asyncio.wait(
+                [
+                    stop_task,
+                    recv_task,
+                ],
+                return_when=asyncio.FIRST_COMPLETED,
+                timeout=wait_time,
+            )
+
+            #
+
+            if (
+                    (time.time() - self._last_ping_send >= self._ping_interval_s) and
+                    not self._ping_in_flight
+            ):
+                now = time.time()
+                self._last_ping_send = now
+                self._ping_in_flight = True
+                await _RemoteProtocol.PingRequest(
+                    time=now,
+                ).send(self._chan)
+
+            #
 
             if recv_task in done:
                 msg: ta.Optional[_RemoteProtocol.Message] = check.isinstance(
@@ -141,6 +188,20 @@ class _RemoteCommandHandler:
 
                 await self._handle_message(msg)
 
+        log.debug('_RemoteCommandHandler loop stopping: %r', self)
+
+        for task in [
+            stop_task,
+            recv_task,
+        ]:
+            if task is not None and not task.done():
+                task.cancel()
+
+        for cmd in self._cmds_by_seq.values():
+            cmd.fut.cancel()
+
+        log.debug('_RemoteCommandHandler loop exited: %r', self)
+
     async def _handle_message(self, msg: _RemoteProtocol.Message) -> None:
         if isinstance(msg, _RemoteProtocol.PingRequest):
             log.debug('Ping: %r', msg)
@@ -148,6 +209,12 @@ class _RemoteCommandHandler:
                 time=msg.time,
             ).send(self._chan)
 
+        elif isinstance(msg, _RemoteProtocol.PingResponse):
+            latency_s = time.time() - msg.time
+            log.debug('Pong: %0.2f ms %r', latency_s * 1000., msg)
+            self._last_ping_recv = time.time()
+            self._ping_in_flight = False
+
         elif isinstance(msg, _RemoteProtocol.CommandRequest):
             fut = asyncio.create_task(self._handle_command_request(msg))
             self._cmds_by_seq[msg.seq] = _RemoteCommandHandler._Command(
@@ -229,16 +296,23 @@ class RemoteCommandExecutor(CommandExecutor):
             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)
+            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)
+                await self._handle_queued_request(req)
+
+            #
 
             if recv_task in done:
                 msg: ta.Optional[_RemoteProtocol.Message] = check.isinstance(
@@ -268,7 +342,7 @@ class RemoteCommandExecutor(CommandExecutor):
 
         log.debug('RemoteCommandExecutor loop exited: %r', self)
 
-    async def _handle_request(self, req: _Request) -> None:
+    async def _handle_queued_request(self, req: _Request) -> None:
         self._reqs_by_seq[req.seq] = req
         await _RemoteProtocol.CommandRequest(
             seq=req.seq,
@@ -282,6 +356,10 @@ class RemoteCommandExecutor(CommandExecutor):
                 time=msg.time,
             ).send(self._chan)
 
+        elif isinstance(msg, _RemoteProtocol.PingResponse):
+            latency_s = time.time() - msg.time
+            log.debug('Pong: %0.2f ms %r', latency_s * 1000., msg)
+
         elif isinstance(msg, _RemoteProtocol.LogResponse):
             log.info(msg.s)
 

ominfra/manage/remote/inject.py

@@ -6,6 +6,7 @@ from omlish.lite.inject import InjectorBindings
 from omlish.lite.inject import inj
 
 from .config import RemoteConfig
+from .connection import InProcessRemoteExecutionConnector
 from .connection import PyremoteRemoteExecutionConnector
 from .connection import RemoteExecutionConnector
 from .payload import RemoteExecutionPayloadFile
@@ -22,12 +23,26 @@ def bind_remote(
 
         inj.bind(SubprocessRemoteSpawning, singleton=True),
         inj.bind(RemoteSpawning, to_key=SubprocessRemoteSpawning),
-
-        inj.bind(PyremoteRemoteExecutionConnector, singleton=True),
-        inj.bind(RemoteExecutionConnector, to_key=PyremoteRemoteExecutionConnector),
     ]
 
+    #
+
+    if remote_config.use_in_process_remote_executor:
+        lst.extend([
+            inj.bind(InProcessRemoteExecutionConnector, singleton=True),
+            inj.bind(RemoteExecutionConnector, to_key=InProcessRemoteExecutionConnector),
+        ])
+    else:
+        lst.extend([
+            inj.bind(PyremoteRemoteExecutionConnector, singleton=True),
+            inj.bind(RemoteExecutionConnector, to_key=PyremoteRemoteExecutionConnector),
+        ])
+
+    #
+
     if (pf := remote_config.payload_file) is not None:
-        lst.append(inj.bind(pf, to_key=RemoteExecutionPayloadFile))
+        lst.append(inj.bind(pf, key=RemoteExecutionPayloadFile))
+
+    #
 
     return inj.as_bindings(*lst)

ominfra/manage/system/commands.py

@@ -0,0 +1,24 @@
+# ruff: noqa: UP006 UP007
+import dataclasses as dc
+
+from omlish.lite.logs import log
+
+from ..commands.base import Command
+from ..commands.base import CommandExecutor
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class CheckSystemPackageCommand(Command['CheckSystemPackageCommand.Output']):
+    @dc.dataclass(frozen=True)
+    class Output(Command.Output):
+        pass
+
+
+class CheckSystemPackageCommandExecutor(CommandExecutor[CheckSystemPackageCommand, CheckSystemPackageCommand.Output]):
+    async def execute(self, cmd: CheckSystemPackageCommand) -> CheckSystemPackageCommand.Output:
+        log.info('Checking system package!')
+
+        return CheckSystemPackageCommand.Output()

ominfra/manage/system/config.py

@@ -0,0 +1,8 @@
+# ruff: noqa: UP006 UP007
+import dataclasses as dc
+import typing as ta
+
+
+@dc.dataclass(frozen=True)
+class SystemConfig:
+    platform: ta.Optional[str] = None

ominfra/manage/system/inject.py

@@ -0,0 +1,54 @@
+# ruff: noqa: UP006 UP007
+import sys
+import typing as ta
+
+from omlish.lite.inject import InjectorBindingOrBindings
+from omlish.lite.inject import InjectorBindings
+from omlish.lite.inject import inj
+
+from ..commands.inject import bind_command
+from .commands import CheckSystemPackageCommand
+from .commands import CheckSystemPackageCommandExecutor
+from .config import SystemConfig
+from .packages import AptSystemPackageManager
+from .packages import BrewSystemPackageManager
+from .packages import SystemPackageManager
+from .types import SystemPlatform
+
+
+def bind_system(
+        *,
+        system_config: SystemConfig,
+) -> InjectorBindings:
+    lst: ta.List[InjectorBindingOrBindings] = [
+        inj.bind(system_config),
+    ]
+
+    #
+
+    platform = system_config.platform or sys.platform
+    lst.append(inj.bind(platform, key=SystemPlatform))
+
+    #
+
+    if platform == 'linux':
+        lst.extend([
+            inj.bind(AptSystemPackageManager, singleton=True),
+            inj.bind(SystemPackageManager, to_key=AptSystemPackageManager),
+        ])
+
+    elif platform == 'darwin':
+        lst.extend([
+            inj.bind(BrewSystemPackageManager, singleton=True),
+            inj.bind(SystemPackageManager, to_key=BrewSystemPackageManager),
+        ])
+
+    #
+
+    lst.extend([
+        bind_command(CheckSystemPackageCommand, CheckSystemPackageCommandExecutor),
+    ])
+
+    #
+
+    return inj.as_bindings(*lst)

ominfra/manage/system/packages.py

@@ -0,0 +1,106 @@
+# ruff: noqa: UP006 UP007
+"""
+TODO:
+ - yum/rpm
+"""
+import abc
+import dataclasses as dc
+import json
+import os
+import typing as ta
+
+from omlish.lite.asyncio.subprocesses import asyncio_subprocess_check_call
+from omlish.lite.asyncio.subprocesses import asyncio_subprocess_check_output
+from omlish.lite.asyncio.subprocesses import asyncio_subprocess_run
+from omlish.lite.check import check
+
+
+SystemPackageOrStr = ta.Union['SystemPackage', str]
+
+
+@dc.dataclass(frozen=True)
+class SystemPackage:
+    name: str
+    version: ta.Optional[str] = None
+
+
+class SystemPackageManager(abc.ABC):
+    @abc.abstractmethod
+    def update(self) -> ta.Awaitable[None]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def upgrade(self) -> ta.Awaitable[None]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def install(self, *packages: SystemPackageOrStr) -> ta.Awaitable[None]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def query(self, *packages: SystemPackageOrStr) -> ta.Awaitable[ta.Mapping[str, SystemPackage]]:
+        raise NotImplementedError
+
+
+class BrewSystemPackageManager(SystemPackageManager):
+    async def update(self) -> None:
+        await asyncio_subprocess_check_call('brew', 'update')
+
+    async def upgrade(self) -> None:
+        await asyncio_subprocess_check_call('brew', 'upgrade')
+
+    async def install(self, *packages: SystemPackageOrStr) -> None:
+        es: ta.List[str] = []
+        for p in packages:
+            if isinstance(p, SystemPackage):
+                es.append(p.name + (f'@{p.version}' if p.version is not None else ''))
+            else:
+                es.append(p)
+        await asyncio_subprocess_check_call('brew', 'install', *es)
+
+    async def query(self, *packages: SystemPackageOrStr) -> ta.Mapping[str, SystemPackage]:
+        pns = [p.name if isinstance(p, SystemPackage) else p for p in packages]
+        o = await asyncio_subprocess_check_output('brew', 'info', '--json', *pns)
+        j = json.loads(o.decode())
+        d: ta.Dict[str, SystemPackage] = {}
+        for e in j:
+            if not e['installed']:
+                continue
+            d[e['name']] = SystemPackage(
+                name=e['name'],
+                version=e['installed'][0]['version'],
+            )
+        return d
+
+
+class AptSystemPackageManager(SystemPackageManager):
+    _APT_ENV: ta.ClassVar[ta.Mapping[str, str]] = {
+        'DEBIAN_FRONTEND': 'noninteractive',
+    }
+
+    async def update(self) -> None:
+        await asyncio_subprocess_check_call('apt', 'update', env={**os.environ, **self._APT_ENV})
+
+    async def upgrade(self) -> None:
+        await asyncio_subprocess_check_call('apt', 'upgrade', '-y', env={**os.environ, **self._APT_ENV})
+
+    async def install(self, *packages: SystemPackageOrStr) -> None:
+        pns = [p.name if isinstance(p, SystemPackage) else p for p in packages]  # FIXME: versions
+        await asyncio_subprocess_check_call('apt', 'install', '-y', *pns, env={**os.environ, **self._APT_ENV})
+
+    async def query(self, *packages: SystemPackageOrStr) -> ta.Mapping[str, SystemPackage]:
+        pns = [p.name if isinstance(p, SystemPackage) else p for p in packages]
+        cmd = ['dpkg-query', '-W', '-f=${Package}=${Version}\n', *pns]
+        stdout, stderr = await asyncio_subprocess_run(
+            *cmd,
+            capture_output=True,
+            check=False,
+        )
+        d: ta.Dict[str, SystemPackage] = {}
+        for l in check.not_none(stdout).decode('utf-8').strip().splitlines():
+            n, v = l.split('=', 1)
+            d[n] = SystemPackage(
+                name=n,
+                version=v,
+            )
+        return d

ominfra/manage/system/types.py

@@ -0,0 +1,5 @@
+# ruff: noqa: UP006 UP007
+import typing as ta
+
+
+SystemPlatform = ta.NewType('SystemPlatform', str)