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

ominfra/scripts/manage.py

@@ -68,7 +68,7 @@ VersionCmpLocalType = ta.Union['NegativeInfinityVersionType', _VersionCmpLocalTy
 VersionCmpKey = ta.Tuple[int, ta.Tuple[int, ...], VersionCmpPrePostDevType, VersionCmpPrePostDevType, VersionCmpPrePostDevType, VersionCmpLocalType]  # noqa
 VersionComparisonMethod = ta.Callable[[VersionCmpKey, VersionCmpKey], bool]
 
-# ../../omlish/lite/asyncio/asyncio.py
+# ../../omlish/asyncs/asyncio/timeouts.py
 AwaitableT = ta.TypeVar('AwaitableT', bound=ta.Awaitable)
 
 # ../../omlish/lite/cached.py
@@ -105,6 +105,9 @@ InjectorBindingOrBindings = ta.Union['InjectorBinding', 'InjectorBindings']
 # ../../omlish/lite/subprocesses.py
 SubprocessChannelOption = ta.Literal['pipe', 'stdout', 'devnull']  # ta.TypeAlias
 
+# system/packages.py
+SystemPackageOrStr = ta.Union['SystemPackage', str]
+
 
 ########################################
 # ../../../omdev/packaging/versions.py
@@ -524,6 +527,22 @@ class MainConfig:
     debug: bool = False
 
 
+########################################
+# ../system/config.py
+
+
+@dc.dataclass(frozen=True)
+class SystemConfig:
+    platform: ta.Optional[str] = None
+
+
+########################################
+# ../system/types.py
+
+
+SystemPlatform = ta.NewType('SystemPlatform', str)
+
+
 ########################################
 # ../../pyremote.py
 """
@@ -1044,10 +1063,47 @@ class PyremoteBootstrapDriver:
 
 
 ########################################
-# ../../../omlish/lite/asyncio/asyncio.py
+# ../../../omlish/asyncs/asyncio/channels.py
 
 
-##
+class AsyncioBytesChannelTransport(asyncio.Transport):
+    def __init__(self, reader: asyncio.StreamReader) -> None:
+        super().__init__()
+
+        self.reader = reader
+        self.closed: asyncio.Future = asyncio.Future()
+
+    # @ta.override
+    def write(self, data: bytes) -> None:
+        self.reader.feed_data(data)
+
+    # @ta.override
+    def close(self) -> None:
+        self.reader.feed_eof()
+        if not self.closed.done():
+            self.closed.set_result(True)
+
+    # @ta.override
+    def is_closing(self) -> bool:
+        return self.closed.done()
+
+
+def asyncio_create_bytes_channel(
+        loop: ta.Any = None,
+) -> ta.Tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+    if loop is None:
+        loop = asyncio.get_running_loop()
+
+    reader = asyncio.StreamReader()
+    protocol = asyncio.StreamReaderProtocol(reader)
+    transport = AsyncioBytesChannelTransport(reader)
+    writer = asyncio.StreamWriter(transport, protocol, reader, loop)
+
+    return reader, writer
+
+
+########################################
+# ../../../omlish/asyncs/asyncio/streams.py
 
 
 ASYNCIO_DEFAULT_BUFFER_LIMIT = 2 ** 16
@@ -1091,7 +1147,8 @@ async def asyncio_open_stream_writer(
     )
 
 
-##
+########################################
+# ../../../omlish/asyncs/asyncio/timeouts.py
 
 
 def asyncio_maybe_timeout(
@@ -2572,6 +2629,8 @@ class RemoteConfig:
 
     heartbeat_interval_s: float = 3.
 
+    use_in_process_remote_executor: bool = False
+
 
 ########################################
 # ../remote/payload.py
@@ -4606,6 +4665,8 @@ class MainBootstrap:
 
     remote_config: RemoteConfig = RemoteConfig()
 
+    system_config: SystemConfig = SystemConfig()
+
 
 ########################################
 # ../commands/execution.py
@@ -4655,7 +4716,7 @@ def install_command_marshaling(
 
 
 ########################################
-# ../deploy/command.py
+# ../deploy/commands.py
 
 
 ##
@@ -4668,9 +4729,6 @@ class DeployCommand(Command['DeployCommand.Output']):
         pass
 
 
-##
-
-
 class DeployCommandExecutor(CommandExecutor[DeployCommand, DeployCommand.Output]):
     async def execute(self, cmd: DeployCommand) -> DeployCommand.Output:
         log.info('Deploying!')
@@ -4769,6 +4827,27 @@ class RemoteChannelImpl(RemoteChannel):
             return await self._recv_obj(ty)
 
 
+########################################
+# ../system/commands.py
+
+
+##
+
+
+@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()
+
+
 ########################################
 # ../../../omlish/lite/subprocesses.py
 
@@ -4955,6 +5034,10 @@ def subprocess_close(
 
 ########################################
 # ../remote/execution.py
+"""
+TODO:
+ - sequence all messages
+"""
 
 
 ##
@@ -5033,38 +5116,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))
+
+            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)
 
-            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(
@@ -5078,6 +5203,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)
@@ -5085,6 +5224,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(
@@ -5166,16 +5311,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(
@@ -5205,7 +5357,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,
@@ -5219,6 +5371,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)
 
@@ -5435,22 +5591,25 @@ async def asyncio_subprocess_communicate(
     return await AsyncioProcessCommunicator(proc).communicate(input, timeout)  # noqa
 
 
-##
-
-
-async def _asyncio_subprocess_check_run(
+async def asyncio_subprocess_run(
         *args: str,
         input: ta.Any = None,  # noqa
         timeout: ta.Optional[float] = None,
+        check: bool = False,  # noqa
+        capture_output: ta.Optional[bool] = None,
         **kwargs: ta.Any,
 ) -> ta.Tuple[ta.Optional[bytes], ta.Optional[bytes]]:
+    if capture_output:
+        kwargs.setdefault('stdout', subprocess.PIPE)
+        kwargs.setdefault('stderr', subprocess.PIPE)
+
     args, kwargs = prepare_subprocess_invocation(*args, **kwargs)
 
     proc: asyncio.subprocess.Process
     async with asyncio_subprocess_popen(*args, **kwargs) as proc:
         stdout, stderr = await asyncio_subprocess_communicate(proc, input, timeout)
 
-    if proc.returncode:
+    if check and proc.returncode:
         raise subprocess.CalledProcessError(
             proc.returncode,
             args,
@@ -5461,6 +5620,9 @@ async def _asyncio_subprocess_check_run(
     return stdout, stderr
 
 
+##
+
+
 async def asyncio_subprocess_check_call(
         *args: str,
         stdout: ta.Any = sys.stderr,
@@ -5468,11 +5630,12 @@ async def asyncio_subprocess_check_call(
         timeout: ta.Optional[float] = None,
         **kwargs: ta.Any,
 ) -> None:
-    _, _ = await _asyncio_subprocess_check_run(
+    _, _ = await asyncio_subprocess_run(
         *args,
         stdout=stdout,
         input=input,
         timeout=timeout,
+        check=True,
         **kwargs,
     )
 
@@ -5483,11 +5646,12 @@ async def asyncio_subprocess_check_output(
         timeout: ta.Optional[float] = None,
         **kwargs: ta.Any,
 ) -> bytes:
-    stdout, stderr = await _asyncio_subprocess_check_run(
+    stdout, stderr = await asyncio_subprocess_run(
         *args,
         stdout=asyncio.subprocess.PIPE,
         input=input,
         timeout=timeout,
+        check=True,
         **kwargs,
     )
 
@@ -5682,9 +5846,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
@@ -5966,6 +6127,102 @@ class SubprocessRemoteSpawning(RemoteSpawning):
                     pass
 
 
+########################################
+# ../system/packages.py
+"""
+TODO:
+ - yum/rpm
+"""
+
+
+@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
+
+
 ########################################
 # ../../../omdev/interp/providers.py
 """
@@ -6119,6 +6376,50 @@ class PyremoteRemoteExecutionConnector(RemoteExecutionConnector):
                 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
+
+
 ########################################
 # ../../../omdev/interp/pyenv.py
 """
@@ -6689,13 +6990,27 @@ 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)
 
@@ -6817,9 +7132,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))
@@ -6947,6 +7259,48 @@ def bind_deploy(
     return inj.as_bindings(*lst)
 
 
+########################################
+# ../system/inject.py
+
+
+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)
+
+
 ########################################
 # ../inject.py
 
@@ -6958,6 +7312,7 @@ def bind_main(
         *,
         main_config: MainConfig,
         remote_config: RemoteConfig,
+        system_config: SystemConfig,
 ) -> InjectorBindings:
     lst: ta.List[InjectorBindingOrBindings] = [
         inj.bind(main_config),
@@ -6966,11 +7321,15 @@ def bind_main(
             main_config=main_config,
         ),
 
+        bind_deploy(),
+
         bind_remote(
             remote_config=remote_config,
         ),
 
-        bind_deploy(),
+        bind_system(
+            system_config=system_config,
+        ),
     ]
 
     #
@@ -7005,6 +7364,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
@@ -7052,6 +7412,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,
             ),
         )