add load-balanced spawner

* Custom Jupyter spawner for starting Jupyter Servers on remote
  platforms over SSH.

Author: oliver-sanders

cylc/uiserver/config_defaults.py

@@ -22,7 +22,7 @@
 
 
 # the command the hub should spawn (i.e. the cylc uiserver itself)
-c.Spawner.cmd = ['cylc', 'uiserver']
+c.Spawner.cmd = ['cylc', 'uis']
 
 # the spawner to invoke this command
 c.JupyterHub.spawner_class = 'jupyterhub.spawner.LocalProcessSpawner'

cylc/uiserver/distributed_spawner.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+# Copyright (C) NIWA & British Crown (Met Office) & Contributors.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+import logging
+from subprocess import Popen, PIPE, DEVNULL
+import sys
+from textwrap import dedent
+from typing import List, Tuple
+
+from traitlets import (
+    DottedObjectName,
+    List as TList,
+    Unicode,
+    default,
+)
+from jupyterhub.spawner import Spawner
+
+from cylc.flow import __version__ as CYLC_VERSION
+from cylc.flow.host_select import select_host
+
+
+logger = logging.getLogger(__name__)
+
+
+class DottedObject(DottedObjectName):
+    """Like DottedObjectName, only it actually imports the thing."""
+
+    def validate(self, obj, value):
+        """Import and return bar given the string foo.bar."""
+        package = '.'.join(value.split('.')[0:-1])
+        obj = value.split('.')[-1]
+        try:
+            if package:
+                module = __import__(package, fromlist=[obj])
+                return module.__dict__[obj]
+            else:
+                return __import__(obj)
+        except ImportError:
+            self.error(obj, value)
+
+
+class DistributedSpawner(Spawner):
+    """A simple SSH Spawner with load balancing capability.
+
+    Runs as the user, no elevated privileges required.
+
+    Requires both passphraseless SSH and a shared filesystem between the
+    hub server and all configured hosts.
+    """
+
+    hosts = TList(
+        trait=Unicode(),
+        config=True,
+        help='''
+            List of host names to choose from.
+        '''
+    )
+
+    ranking = Unicode(
+        config=True,
+        help='''
+            Ranking to use for load balancing purposes.
+
+            If unspecified a host is chosen at random.
+
+            These rankings can be used to pick the host with the most available
+            memory or filter out hosts with high server load.
+
+            These rankings are provided in the same format as
+            :cylc:conf`global.cylc[scheduler][run hosts]ranking`.
+        '''
+    )
+
+    ssh_cmd = TList(
+        trait=Unicode(),
+        config=True,
+        help='''
+            The SSH command to use for connecting to the remote hosts.
+
+            E.G: ``['ssh']`` (default)
+        '''
+    )
+
+    get_ip_from_hostname = DottedObject(
+        config=True,
+        help='''
+            Function for obtaining the IP address from a hostname.
+
+            E.G: ``socket.gethostbyname`` (default)
+        '''
+    )
+
+    @default('get_ip_from_hostname')
+    def default_ip_from_hostname_command(self):
+        return 'socket.gethostbyname'
+
+    @default('ssh_cmd')
+    def default_ssh_command(self):
+        return ['ssh']
+
+    def __init__(self, *args, **kwargs):
+        Spawner.__init__(self, *args, **kwargs)
+        self._proc = None
+
+    def choose_host(self):
+        return select_host(self.hosts, self.ranking)[1]
+
+    def get_env(self):
+        return {
+            **Spawner.get_env(self),
+            'CYLC_VERSION': CYLC_VERSION,
+            'JUPYTERHUB_SERVICE_PREFIX': '/user/osanders/'
+        }
+
+    def get_env_cmd(self) -> List[str]:
+        """Return the spawner environment as an ``env`` command.
+
+        Example output: ``['env', 'FOO=bar']``
+        """
+        env = self.get_env()
+        if not env:
+            return []
+        return [
+            'env'
+        ] + [
+            f'{key}={value}'
+            for key, value in self.get_env().items()
+        ]
+
+    def get_remote_port(self) -> int:
+        """Find an open port to spawn the app onto.
+
+        Invokes Python over SSH to call a JupyterHub utility function on the
+        remote host.
+        """
+        cmd = [
+            *self.ssh_cmd,
+            self._host,
+            sys.executable,
+        ]
+        logger.debug('$ ' + ' '.join(cmd))
+        proc = Popen(
+            cmd,
+            stdout=PIPE,
+            stdin=PIPE,
+            text=True
+        )
+        proc.communicate(dedent('''
+            from jupyterhub.utils import random_port
+            print(random_port())
+        '''))
+        if proc.returncode:
+            raise Exception('remote proc failed')
+        stdout, _ = proc.communicate()
+        try:
+            port = int(stdout)
+        except Exception:
+            raise Exception(f'invalid stdout: {stdout}')
+        return port
+
+    async def start(self) -> Tuple[str, str]:
+        self._host = self.choose_host()
+        port = self.get_remote_port()
+        cmd = [
+            *self.ssh_cmd,
+            self._host,
+            *self.get_env_cmd(),
+            *self.cmd,
+            *self.get_args(),
+            # NOTE: selg.get_args may set --port, however, we override it
+            f'--port={port}',
+        ]
+        logger.info('$ ' + ' '.join(cmd))
+        self._proc = Popen(
+            cmd,
+            stderr=PIPE,
+            stdin=DEVNULL,
+            text=True
+        )
+
+        ip = self.get_ip_from_hostname(self._host)
+        return (ip, port)
+
+    async def stop(self, now=False):
+        if self._proc:
+            self._proc.kill()
+            self._proc = None
+
+    async def poll(self):
+        if self._proc:
+            return self._proc.poll()

cylc/uiserver/main.py

@@ -23,6 +23,7 @@
 from functools import partial
 from logging.config import dictConfig
 from pathlib import Path, PurePath
+from socket import gethostname
 import sys
 from typing import Any, Tuple, Type, List
 
@@ -420,11 +421,12 @@ def _make_app(self, debug: bool):
 
     def start(self, debug: bool):
         logger.info("Starting Cylc UI Server")
+        logger.info(f"Running on: {gethostname()}")
+        logger.info(f"Listening on port: {self._port}")
+        logger.info(f'Serving UI from: {self.ui_path}')
         logger.info(
             f"JupyterHub Service Prefix: {self._jupyter_hub_service_prefix}"
         )
-        logger.info(f"Listening on port: {self._port}")
-        logger.info(f'Serving UI from: {self.ui_path}')
         app = self._make_app(debug)
         signal.signal(signal.SIGINT, app.signal_handler)
         app.listen(self._port)