refactor: Clean up BaseFunction inheritance using contextvars

* Previously, to bind the current session with API function classes,
  we generated new type objects at runtime.

  - This has confused IDEs and type checkers.

  - Now type checkers can statically deduce the types for individual API
    function classes.

  - TODO: many many type errors are there still...

* Now we use contextvars (ai.backend.client.session.api_session)
  to keep the reference to the current session.

  - There are no public Session/AsyncSession API changes!

  - Only the API function classes need to be rewritten.

  - For synchronous Session, we pass the context to the separate worker
    thread using copy_context() whenever calling API functions, which is
    a light-weight operation.

* Remove redundant src/ai/backend/client/etcd.py which had been already
  copied to src/ai/backend/client/func/etcd.py

BREAKING-CHANGE: Dropped Python 3.6 support. Now it requires Python 3.7 or higher.

Author: achimnol

.travis.yml

@@ -10,7 +10,6 @@ stages:
 
 # build matrix for test stage
 python:
-  - "3.6"
   - "3.7"
   - "3.8"
 os:
@@ -53,7 +52,7 @@ jobs:
   fast_finish: true
   # exclude the duplicate default test stage
   exclude:
-    - python: "3.6"
+    - python: "3.7"
 
 notifications:
   webhooks:

setup.cfg

@@ -13,3 +13,7 @@ norecursedirs = venv virtualenv .git
 timeout = 5
 markers =
     integration: Test cases that require real manager (and agents) to be running on http://localhost:8081.
+
+[mypy]
+ignore_missing_imports = true
+namespace_packages = true

setup.py

@@ -82,7 +82,6 @@ def read_src_version():
         'Intended Audience :: Developers',
         'Programming Language :: Python',
         'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',
         'Operating System :: POSIX',
@@ -94,7 +93,7 @@ def read_src_version():
     ],
     package_dir={'': 'src'},
     packages=find_namespace_packages(where='src', include='ai.backend.*'),
-    python_requires='>=3.6',
+    python_requires='>=3.7',
     setup_requires=setup_requires,
     install_requires=install_requires,
     extras_require={

src/ai/backend/client/etcd.py

@@ -1,33 +1,31 @@
 from typing import Any, Mapping, Optional
 
-from .base import api_function
+from .base import api_function, BaseFunction
 from ..request import Request
+from ..session import api_session
 
 __all__ = (
     'Admin',
 )
 
 
-class Admin:
-    '''
+class Admin(BaseFunction):
+    """
     Provides the function interface for making admin GrapQL queries.
 
     .. note::
 
       Depending on the privilege of your API access key, you may or may not
       have access to querying/mutating server-side resources of other
       users.
-    '''
-
-    session = None
-    '''The client session instance that this function class is bound to.'''
+    """
 
     @api_function
     @classmethod
     async def query(cls, query: str,
                     variables: Optional[Mapping[str, Any]] = None,
                     ) -> Any:
-        '''
+        """
         Sends the GraphQL query and returns the response.
 
         :param query: The GraphQL query string.
@@ -36,12 +34,12 @@ async def query(cls, query: str,
             in the query.
 
         :returns: The object parsed from the response JSON string.
-        '''
+        """
         gql_query = {
             'query': query,
             'variables': variables if variables else {},
         }
-        rqst = Request(cls.session, 'POST', '/admin/graphql')
+        rqst = Request(api_session.get(), 'POST', '/admin/graphql')
         rqst.set_json(gql_query)
         async with rqst.fetch() as resp:
             return await resp.json()

src/ai/backend/client/func/admin.py

@@ -1,17 +1,18 @@
 import textwrap
 from typing import Iterable, Sequence
 
-from .base import api_function
+from .base import api_function, BaseFunction
 from ..request import Request
+from ..session import api_session
 
 __all__ = (
     'Agent',
     'AgentWatcher',
 )
 
 
-class Agent:
-    '''
+class Agent(BaseFunction):
+    """
     Provides a shortcut of :func:`Admin.query()
     <ai.backend.client.admin.Admin.query>` that fetches various agent
     information.
@@ -20,10 +21,7 @@ class Agent:
 
       All methods in this function class require your API access key to
       have the *admin* privilege.
-    '''
-
-    session = None
-    '''The client session instance that this function class is bound to.'''
+    """
 
     @api_function
     @classmethod
@@ -32,7 +30,7 @@ async def list_with_limit(cls,
                               offset,
                               status: str = 'ALIVE',
                               fields: Iterable[str] = None) -> Sequence[dict]:
-        '''
+        """
         Fetches the list of agents with the given status with limit and offset for
         pagination.
 
@@ -42,7 +40,7 @@ async def list_with_limit(cls,
             status (one of ``'ALIVE'``, ``'TERMINATED'``, ``'LOST'``,
             etc.)
         :param fields: Additional per-agent query fields to fetch.
-        '''
+        """
         if fields is None:
             fields = (
                 'id',
@@ -65,7 +63,7 @@ async def list_with_limit(cls,
             'offset': offset,
             'status': status,
         }
-        rqst = Request(cls.session, 'POST', '/admin/graphql')
+        rqst = Request(api_session.get(), 'POST', '/admin/graphql')
         rqst.set_json({
             'query': q,
             'variables': variables,
@@ -81,14 +79,14 @@ async def detail(cls, agent_id: str, fields: Iterable[str] = None) -> Sequence[d
             fields = ('id', 'status', 'addr', 'region', 'first_contact',
                       'cpu_cur_pct', 'mem_cur_bytes',
                       'available_slots', 'occupied_slots')
-        query = textwrap.dedent('''\
+        query = textwrap.dedent("""\
             query($agent_id: String!) {
                 agent(agent_id: $agent_id) {$fields}
             }
-        ''')
+        """)
         query = query.replace('$fields', ' '.join(fields))
         variables = {'agent_id': agent_id}
-        rqst = Request(cls.session, 'POST', '/admin/graphql')
+        rqst = Request(api_session.get(), 'POST', '/admin/graphql')
         rqst.set_json({
             'query': query,
             'variables': variables,
@@ -98,27 +96,24 @@ async def detail(cls, agent_id: str, fields: Iterable[str] = None) -> Sequence[d
             return data['agent']
 
 
-class AgentWatcher:
-    '''
+class AgentWatcher(BaseFunction):
+    """
     Provides a shortcut of :func:`Admin.query()
     <ai.backend.client.admin.Admin.query>` that manipulate agent status.
 
     .. note::
 
       All methods in this function class require you to
       have the *superadmin* privilege.
-    '''
-
-    session = None
-    '''The client session instance that this function class is bound to.'''
+    """
 
     @api_function
     @classmethod
     async def get_status(cls, agent_id: str) -> dict:
-        '''
+        """
         Get agent and watcher status.
-        '''
-        rqst = Request(cls.session, 'GET', '/resource/watcher')
+        """
+        rqst = Request(api_session.get(), 'GET', '/resource/watcher')
         rqst.set_json({'agent_id': agent_id})
         async with rqst.fetch() as resp:
             data = await resp.json()
@@ -130,10 +125,10 @@ async def get_status(cls, agent_id: str) -> dict:
     @api_function
     @classmethod
     async def agent_start(cls, agent_id: str) -> dict:
-        '''
+        """
         Start agent.
-        '''
-        rqst = Request(cls.session, 'POST', '/resource/watcher/agent/start')
+        """
+        rqst = Request(api_session.get(), 'POST', '/resource/watcher/agent/start')
         rqst.set_json({'agent_id': agent_id})
         async with rqst.fetch() as resp:
             data = await resp.json()
@@ -145,10 +140,10 @@ async def agent_start(cls, agent_id: str) -> dict:
     @api_function
     @classmethod
     async def agent_stop(cls, agent_id: str) -> dict:
-        '''
+        """
         Stop agent.
-        '''
-        rqst = Request(cls.session, 'POST', '/resource/watcher/agent/stop')
+        """
+        rqst = Request(api_session.get(), 'POST', '/resource/watcher/agent/stop')
         rqst.set_json({'agent_id': agent_id})
         async with rqst.fetch() as resp:
             data = await resp.json()
@@ -160,10 +155,10 @@ async def agent_stop(cls, agent_id: str) -> dict:
     @api_function
     @classmethod
     async def agent_restart(cls, agent_id: str) -> dict:
-        '''
+        """
         Restart agent.
-        '''
-        rqst = Request(cls.session, 'POST', '/resource/watcher/agent/restart')
+        """
+        rqst = Request(api_session.get(), 'POST', '/resource/watcher/agent/restart')
         rqst.set_json({'agent_id': agent_id})
         async with rqst.fetch() as resp:
             data = await resp.json()

src/ai/backend/client/func/agent.py

@@ -1,26 +1,27 @@
-from .base import api_function
+from .base import api_function, BaseFunction
 from ..request import Request
+from ..session import api_session
 
 __all__ = (
     'Auth',
 )
 
 
-class Auth:
-    '''
+class Auth(BaseFunction):
+    """
     Provides the function interface for login session management and authorization.
-    '''
+    """
 
     @api_function
     @classmethod
     async def login(cls, user_id: str, password: str) -> dict:
-        '''
+        """
         Log-in into the endpoint with the given user ID and password.
         It creates a server-side web session and return
         a dictionary with ``"authenticated"`` boolean field and
         JSON-encoded raw cookie data.
-        '''
-        rqst = Request(cls.session, 'POST', '/server/login')
+        """
+        rqst = Request(api_session.get(), 'POST', '/server/login')
         rqst.set_json({
             'username': user_id,
             'password': password,
@@ -36,11 +37,11 @@ async def login(cls, user_id: str, password: str) -> dict:
     @api_function
     @classmethod
     async def logout(cls) -> None:
-        '''
+        """
         Log-out from the endpoint.
         It clears the server-side web session.
-        '''
-        rqst = Request(cls.session, 'POST', '/server/logout')
+        """
+        rqst = Request(api_session.get(), 'POST', '/server/logout')
         async with rqst.fetch() as resp:
             resp.raw_response.raise_for_status()
 
@@ -50,7 +51,7 @@ async def update_password(cls, old_password: str, new_password: str, new_passwor
         """
         Update user's password. This API works only for account owner.
         """
-        rqst = Request(cls.session, 'POST', '/auth/update-password')
+        rqst = Request(api_session.get(), 'POST', '/auth/update-password')
         rqst.set_json({
             'old_password': old_password,
             'new_password': new_password,