fix: Update CLI with backend.ai-cli (#94)

* feat: Update CLI to use the new cli pkg with interrupt-handling mixin
* setup: Upgrade backend.ai-cli to 0.4.1
* fix: Restore old code because now we use mixins
* fix: Ensure proper SSL termination and proxy context termination
 - Currently Python's asyncio does not wait for processing of
    "connection lost" event for all open SSL transports.
    This generates a very long and messy exception traces, which would
    confuse the end-users.
    We apply a hack to mitigate this issue.
  - When we call other context managers inside context managers and get
    exceptions during an outer enter handler,
    we should call exit handlers of the inner context managers properly
    once their enter handlers are called.

Author: achimnol

CHANGELOG.rst

@@ -1,6 +1,13 @@
 Changes
 =======
 
+20.03.0b2 (2020-04-xx)
+----------------------
+
+* FIX: Improve exception handling in ``backend.ai app`` command
+  and internally update backend.ai-cli package to work consistently
+  with #93 (#94)
+
 19.12.0b1 (2020-01-11)
 ----------------------
 

setup.py

@@ -6,7 +6,7 @@
     'setuptools>=45.2.0',
 ]
 install_requires = [
-    'backend.ai-cli~=0.3',
+    'backend.ai-cli~=0.4.1',
     'aiohttp~=3.6.2',
     'appdirs~=1.4.3',
     'async_timeout~=3.0',  # to avoid pip10 resolver issue
@@ -98,6 +98,7 @@ def read_src_version():
     install_requires=install_requires,
     extras_require={
         'dev': dev_requires,
+        'build': build_requires,
         'test': test_requires,
         'lint': lint_requires,
         'typecheck': typecheck_requires,

src/ai/backend/client/cli/__init__.py

@@ -1,18 +1,13 @@
-from pathlib import Path
-import os
-import signal
-import sys
 import warnings
 
 import click
-from click.exceptions import ClickException, Abort
 
 from .. import __version__
 from ..config import APIConfig, set_config
-from ai.backend.cli.extensions import AliasGroup
+from ai.backend.cli.extensions import ExtendedCommandGroup
 
 
-@click.group(cls=AliasGroup,
+@click.group(cls=ExtendedCommandGroup,
              context_settings={'help_option_names': ['-h', '--help']})
 @click.option('--skip-sslcert-validation',
               help='Skip SSL certificate validation for all API requests.',
@@ -29,24 +24,6 @@ def main(skip_sslcert_validation):
     warnings.showwarning = show_warning
 
 
-@click.command(context_settings=dict(ignore_unknown_options=True,
-                                     allow_extra_args=True))
-def run_alias():
-    """
-    Quick aliases for run command.
-    """
-    mode = Path(sys.argv[0]).stem
-    help = True if len(sys.argv) <= 1 else False
-    if mode == 'lcc':
-        sys.argv.insert(1, 'c')
-    elif mode == 'lpython':
-        sys.argv.insert(1, 'python')
-    sys.argv.insert(1, 'run')
-    if help:
-        sys.argv.append('--help')
-    run_main()
-
-
 def _attach_command():
     from . import admin, config, app, files, logs, manager, proxy, ps, run  # noqa
     from . import vfolder       # noqa
@@ -55,48 +32,3 @@ def _attach_command():
 
 
 _attach_command()
-
-
-def run_main():
-    try:
-        _interrupted = False
-        main.main(
-            standalone_mode=False,
-            prog_name='backend.ai',
-        )
-    except KeyboardInterrupt:
-        # For interruptions outside the Click's exception handling block.
-        print("Interrupted!", end="", file=sys.stderr)
-        sys.stderr.flush()
-        _interrupted = True
-    except Abort as e:
-        # Click wraps unhandled KeyboardInterrupt with a plain
-        # sys.exit(1) call and prints "Aborted!" message
-        # (which would look non-sense to users).
-        # This is *NOT* what we want.
-        # Instead of relying on Click, mark the _interrupted
-        # flag to perform our own exit routines.
-        if isinstance(e.__context__, KeyboardInterrupt):
-            print("Interrupted!", end="", file=sys.stderr)
-            sys.stderr.flush()
-            _interrupted = True
-        else:
-            print("Aborted!", end="", file=sys.stderr)
-            sys.stderr.flush()
-            sys.exit(1)
-    except ClickException as e:
-        e.show()
-        sys.exit(e.exit_code)
-    finally:
-        if _interrupted:
-            # Override the exit code when it's interrupted,
-            # referring https://github.com/python/cpython/pull/11862
-            if sys.platform.startswith('win'):
-                # Use STATUS_CONTROL_C_EXIT to notify cmd.exe
-                # for interrupted exit
-                sys.exit(-1073741510)
-            else:
-                # Use the default signal handler to set the exit
-                # code properly for interruption.
-                signal.signal(signal.SIGINT, signal.SIG_DFL)
-                os.kill(os.getpid(), signal.SIGINT)

src/ai/backend/client/cli/__main__.py

@@ -1,4 +1,4 @@
-from . import run_main
+from . import main
 
 
-run_main()
+main()

src/ai/backend/client/cli/app.py

@@ -1,8 +1,6 @@
 import asyncio
 import json
-import os
 import shlex
-import signal
 import sys
 from typing import (
     Union, Optional,
@@ -197,34 +195,37 @@ async def __aenter__(self) -> None:
         await self.api_session.__aenter__()
 
         user_url_template = "{protocol}://{host}:{port}"
-        compute_session = self.api_session.ComputeSession(self.session_name)
-        all_apps = await compute_session.stream_app_info()
-        for app_info in all_apps:
-            if app_info['name'] == self.app_name:
-                if 'url_template' in app_info.keys():
-                    user_url_template = app_info['url_template']
-                break
-        else:
-            print_fail(f'The app "{self.app_name}" is not supported by the session.')
-            self.exit_code = 1
-            os.kill(0, signal.SIGINT)
-            return
-
-        self.local_server = await asyncio.start_server(
-            self.handle_connection, self.host, self.port)
-        user_url = user_url_template.format(
-            protocol=self.protocol,
-            host=self.host,
-            port=self.port,
-        )
-        print_info(
-            "A local proxy to the application \"{0}\" ".format(self.app_name) +
-            "provided by the session \"{0}\" ".format(self.session_name) +
-            "is available at:\n{0}".format(user_url)
-        )
-        if self.host == '0.0.0.0':
-            print_warn('NOTE: Replace "0.0.0.0" with the actual hostname you use '
-                       'to connect with the CLI app proxy.')
+        try:
+            compute_session = self.api_session.ComputeSession(self.session_name)
+            all_apps = await compute_session.stream_app_info()
+            for app_info in all_apps:
+                if app_info['name'] == self.app_name:
+                    if 'url_template' in app_info.keys():
+                        user_url_template = app_info['url_template']
+                    break
+            else:
+                print_fail(f'The app "{self.app_name}" is not supported by the session.')
+                self.exit_code = 1
+                return
+
+            self.local_server = await asyncio.start_server(
+                self.handle_connection, self.host, self.port)
+            user_url = user_url_template.format(
+                protocol=self.protocol,
+                host=self.host,
+                port=self.port,
+            )
+            print_info(
+                "A local proxy to the application \"{0}\" ".format(self.app_name) +
+                "provided by the session \"{0}\" ".format(self.session_name) +
+                "is available at:\n{0}".format(user_url)
+            )
+            if self.host == '0.0.0.0':
+                print_warn('NOTE: Replace "0.0.0.0" with the actual hostname you use '
+                        'to connect with the CLI app proxy.')
+        except Exception:
+            await self.api_session.__aexit__(*sys.exc_info())
+            raise
 
     async def __aexit__(self, *exc_info) -> None:
         if self.local_server is not None:
@@ -267,15 +268,19 @@ def app(session_name, app, protocol, bind, arg, env):
     elif len(bind_parts) == 2:
         host = bind_parts[0]
         port = int(bind_parts[1])
-    proxy_ctx = ProxyRunnerContext(
-        host, port,
-        session_name, app,
-        protocol=protocol,
-        args=arg,
-        envs=env,
-    )
-    asyncio_run_forever(proxy_ctx)
-    sys.exit(proxy_ctx.exit_code)
+    try:
+        proxy_ctx = ProxyRunnerContext(
+            host, port,
+            session_name, app,
+            protocol=protocol,
+            args=arg,
+            envs=env,
+        )
+        asyncio_run_forever(proxy_ctx)
+        sys.exit(proxy_ctx.exit_code)
+    except Exception as e:
+        print_error(e)
+        sys.exit(1)
 
 
 @main.command()

src/ai/backend/client/cli/dotfile.py

@@ -3,12 +3,12 @@
 import click
 from tabulate import tabulate
 
-from . import AliasGroup, main
+from . import main
 from .pretty import print_info, print_warn, print_error
 from ..session import Session
 
 
-@main.group(cls=AliasGroup)
+@main.group()
 def dotfile():
     '''Provides dotfile operations.'''
 

src/ai/backend/client/cli/session_template.py

@@ -3,13 +3,12 @@
 import click
 from tabulate import tabulate
 
-from ai.backend.cli.extensions import AliasGroup
 from . import main
 from .pretty import print_info, print_warn, print_error
 from ..session import Session
 
 
-@main.group(cls=AliasGroup, aliases=['sesstpl'])
+@main.group(aliases=['sesstpl'])
 def session_template():
     '''Provides task template operations'''
 

src/ai/backend/client/cli/vfolder.py

@@ -6,13 +6,12 @@
 import click
 from tabulate import tabulate
 
-from ai.backend.cli.extensions import AliasGroup
 from . import main
 from .pretty import print_wait, print_done, print_error, print_fail
 from ..session import Session
 
 
-@main.group(cls=AliasGroup)
+@main.group()
 def vfolder():
     '''Provides virtual folder operations.'''
 

src/ai/backend/client/request.py

@@ -5,6 +5,7 @@
 import io
 import logging
 from pathlib import Path
+import sys
 from typing import Any, Callable, Mapping, Sequence, Union
 
 import aiohttp
@@ -460,6 +461,7 @@ async def __aenter__(self):
                 raw_resp = await self._rqst_ctx.__aenter__()
                 if self.check_status and raw_resp.status // 100 != 2:
                     msg = await raw_resp.text()
+                    await raw_resp.__aexit__(None, None, None)
                     raise BackendAPIError(raw_resp.status, raw_resp.reason, msg)
                 return self.response_cls(self.session, raw_resp,
                                          async_mode=self._async_mode)
@@ -475,6 +477,7 @@ async def __aenter__(self):
             except aiohttp.ClientResponseError as e:
                 msg = 'API endpoint response error.\n' \
                       '\u279c {!r}'.format(e)
+                await raw_resp.__aexit__(*sys.exc_info())
                 raise BackendClientError(msg) from e
 
     def __exit__(self, *args):

src/ai/backend/client/session.py

@@ -63,6 +63,46 @@ async def _negotiate_api_version(
         return client_version
 
 
+async def _close_aiohttp_session(session: aiohttp.ClientSession):
+    # This is a hacky workaround for premature closing of SSL transports
+    # on Windows Proactor event loops.
+    # Thanks to Vadim Markovtsev's comment on the aiohttp issue #1925.
+    # (https://github.com/aio-libs/aiohttp/issues/1925#issuecomment-592596034)
+    transports = 0
+    all_is_lost = asyncio.Event()
+    if len(session.connector._conns) == 0:
+        all_is_lost.set()
+    for conn in session.connector._conns.values():
+        for handler, _ in conn:
+            proto = getattr(handler.transport, "_ssl_protocol", None)
+            if proto is None:
+                continue
+            transports += 1
+            orig_lost = proto.connection_lost
+            orig_eof_received = proto.eof_received
+
+            def connection_lost(exc):
+                orig_lost(exc)
+                nonlocal transports
+                transports -= 1
+                if transports == 0:
+                    all_is_lost.set()
+
+            def eof_received():
+                try:
+                    orig_eof_received()
+                except AttributeError:
+                    # It may happen that eof_received() is called after
+                    # _app_protocol and _transport are set to None.
+                    pass
+
+            proto.connection_lost = connection_lost
+            proto.eof_received = eof_received
+    await session.close()
+    if transports > 0:
+        await all_is_lost.wait()
+
+
 class _SyncWorkerThread(threading.Thread):
 
     sentinel = object()
@@ -374,7 +414,7 @@ def close(self):
         if self._closed:
             return
         self._closed = True
-        self._worker_thread.work_queue.put(self.aiohttp_session.close())
+        self._worker_thread.work_queue.put(_close_aiohttp_session(self.aiohttp_session))
         self._worker_thread.work_queue.put(self.worker_thread.sentinel)
         self._worker_thread.join()
 
@@ -600,7 +640,7 @@ async def close(self):
         if self._closed:
             return
         self._closed = True
-        await self.aiohttp_session.close()
+        await _close_aiohttp_session(self.aiohttp_session)
 
     async def __aenter__(self):
         assert not self.closed, 'Cannot reuse closed session'