Add support for kernel launch parameters

jupyter_kernel_mgmt/discovery.py

@@ -25,7 +25,7 @@ def find_kernels(self):
         pass
 
     @abstractmethod
-    def launch(self, name, cwd=None):
+    def launch(self, name, cwd=None, launch_params=None):
         """Launch a kernel, return (connection_info, kernel_manager).
 
         name will be one of the kernel names produced by find_kernels()
@@ -34,7 +34,7 @@ def launch(self, name, cwd=None):
         """
         pass
 
-    def launch_async(self, name, cwd=None):
+    def launch_async(self, name, cwd=None, launch_params=None):
         """Launch a kernel asynchronously using asyncio.
 
         name will be one of the kernel names produced by find_kernels()
@@ -80,16 +80,17 @@ def find_kernels(self):
                 'metadata': spec.metadata,
             }
 
-    def launch(self, name, cwd=None):
+    def launch(self, name, cwd=None, launch_params=None):
         spec = self.ksm.get_kernel_spec(name)
-        launcher = SubprocessKernelLauncher(kernel_cmd=spec.argv, extra_env=spec.env, cwd=cwd)
+        launcher = SubprocessKernelLauncher(kernel_cmd=spec.argv, extra_env=spec.env, cwd=cwd,
+                                            launch_params=launch_params)
         return launcher.launch()
 
-    def launch_async(self, name, cwd=None):
+    def launch_async(self, name, cwd=None, launch_params=None):
         from .subproc.async_manager import AsyncSubprocessKernelLauncher
         spec = self.ksm.get_kernel_spec(name)
         return AsyncSubprocessKernelLauncher(
-            kernel_cmd=spec.argv, extra_env=spec.env, cwd=cwd).launch()
+            kernel_cmd=spec.argv, extra_env=spec.env, cwd=cwd, launch_params=launch_params).launch()
 
 
 class IPykernelProvider(KernelProviderBase):
@@ -123,22 +124,22 @@ def find_kernels(self):
                 'resource_dir': info['resource_dir'],
             }
 
-    def launch(self, name, cwd=None):
+    def launch(self, name, cwd=None, launch_params=None):
         info = self._check_for_kernel()
         if info is None:
             raise Exception("ipykernel is not importable")
 
         launcher = SubprocessKernelLauncher(kernel_cmd=info['spec']['argv'],
-                                            extra_env={}, cwd=cwd)
+                                            extra_env={}, cwd=cwd, launch_params=launch_params)
         return launcher.launch()
 
-    def launch_async(self, name, cwd=None):
+    def launch_async(self, name, cwd=None, launch_params=None):
         from .subproc.async_manager import AsyncSubprocessKernelLauncher
         info = self._check_for_kernel()
         if info is None:
             raise Exception("ipykernel is not importable")
         return AsyncSubprocessKernelLauncher(
-            kernel_cmd=info['spec']['argv'], extra_env={}, cwd=cwd).launch()
+            kernel_cmd=info['spec']['argv'], extra_env={}, cwd=cwd, launch_params=launch_params).launch()
 
 
 class KernelFinder(object):
@@ -188,22 +189,22 @@ def find_kernels(self):
                 kernel_type = provider.id + '/' + kernel_name
                 yield kernel_type, attributes
 
-    def launch(self, name, cwd=None):
+    def launch(self, name, cwd=None, launch_params=None):
         """Launch a kernel of a given kernel type.
         """
         provider_id, kernel_id = name.split('/', 1)
         for provider in self.providers:
             if provider_id == provider.id:
-                return provider.launch(kernel_id, cwd)
+                return provider.launch(kernel_id, cwd=cwd, launch_params=launch_params)
         raise KeyError(provider_id)
 
-    def launch_async(self, name, cwd=None):
+    def launch_async(self, name, cwd=None, launch_params=None):
         """Launch a kernel of a given kernel type, using asyncio.
         """
         provider_id, kernel_id = name.split('/', 1)
         for provider in self.providers:
             if provider_id == provider.id:
-                return provider.launch_async(kernel_id, cwd)
+                return provider.launch_async(kernel_id, cwd=cwd, launch_params=launch_params)
         raise KeyError(provider_id)
 
 

jupyter_kernel_mgmt/subproc/launcher.py

@@ -24,6 +24,7 @@
 port_names = ['shell_port', 'iopub_port', 'stdin_port', 'control_port',
               'hb_port']
 
+
 class SubprocessKernelLauncher:
     """Launch kernels in a subprocess.
 
@@ -44,13 +45,14 @@ class SubprocessKernelLauncher:
     """
     transport = 'tcp'
 
-    def __init__(self, kernel_cmd, cwd, extra_env=None, ip=None):
+    def __init__(self, kernel_cmd, cwd, extra_env=None, ip=None, launch_params=None):
         self.kernel_cmd = kernel_cmd
         self.cwd = cwd
         self.extra_env = extra_env
         if ip is None:
             ip = localhost()
         self.ip = ip
+        self.launch_params = launch_params
         self.log = get_app_logger()
 
         if self.transport == 'tcp' and not is_local_ip(ip):
@@ -138,18 +140,29 @@ def format_kernel_cmd(self, connection_file, kernel_resource_dir=None):
             # but it should be.
             cmd[0] = sys.executable
 
-        ns = dict(connection_file=connection_file,
+        # Preserve system-owned substitutions by starting with launch params
+        ns = dict()
+        if isinstance(self.launch_params, dict):
+            ns.update(self.launch_params)
+
+        # Add system-owned substitutions
+        ns.update(dict(connection_file=connection_file,
                   prefix=sys.prefix,
-                  )
+                  ))
 
         if kernel_resource_dir:
             ns["resource_dir"] = kernel_resource_dir
 
         pat = re.compile(r'{([A-Za-z0-9_]+)}')
 
         def from_ns(match):
-            """Get the key out of ns if it's there, otherwise no change."""
-            return ns.get(match.group(1), match.group())
+            """Get the key out of ns if it's there, otherwise no change.
+               Return as string since that's what is required by pattern
+               matching.  We know this should be safe currently, because
+               only 'connection_file', 'sys.prefix' and 'resource_dir' are
+               candidates - all of which are strings.
+            """
+            return str(ns.get(match.group(1), match.group()))
 
         return [pat.sub(from_ns, arg) for arg in cmd]
 
@@ -310,12 +323,12 @@ def prepare_interrupt_event(env, interrupt_event=None):
         env["IPY_INTERRUPT_EVENT"] = env["JPY_INTERRUPT_EVENT"]
         return interrupt_event
 
-def start_new_kernel(kernel_cmd, startup_timeout=60, cwd=None):
+def start_new_kernel(kernel_cmd, startup_timeout=60, cwd=None, launch_params=None):
     """Start a new kernel, and return its Manager and a blocking client"""
     from ..client import BlockingKernelClient
     cwd = cwd or os.getcwd()
 
-    launcher = SubprocessKernelLauncher(kernel_cmd, cwd=cwd)
+    launcher = SubprocessKernelLauncher(kernel_cmd, cwd=cwd, launch_params=launch_params)
     connection_info, km = launcher.launch()
     kc = BlockingKernelClient(connection_info, manager=km)
     try:

jupyter_kernel_mgmt/tests/test_discovery.py

@@ -19,10 +19,10 @@ class DummyKernelProvider(discovery.KernelProviderBase):
     def find_kernels(self):
         yield 'sample', {'argv': ['dummy_kernel']}
 
-    def launch(self, name, cwd=None):
+    def launch(self, name, cwd=None, launch_params=None):
         return {}, DummyKernelManager()
 
-    def launch_async(self, name, cwd=None):
+    def launch_async(self, name, cwd=None, launch_params=None):
         pass
 
 
@@ -33,10 +33,18 @@ class DummyKernelSpecProvider(discovery.KernelSpecProvider):
 
     # find_kernels() is inherited from KernelsSpecProvider
 
-    def launch(self, name, cwd=None):
+    def launch(self, name, cwd=None, launch_params=None):
         return {}, DummyKernelManager()
 
 
+class LaunchParamsKernelProvider(discovery.KernelSpecProvider):
+    """A dummy kernelspec provider subclass for testing KernelFinder and KernelSpecProvider subclasses"""
+    id = 'params_kspec'
+    kernel_file = 'params_kspec.json'
+
+    # find_kernels() and launch() are inherited from KernelsSpecProvider
+
+
 class DummyKernelManager(KernelManagerABC):
     _alive = True
 
@@ -60,10 +68,6 @@ def interrupt(self):
     def kill(self):
         self._alive = False
 
-    def get_connection_info(self):
-        """Return a dictionary of connection information"""
-        return {}
-
 
 class ProviderApplication(Application):
     name = 'ProviderApplication'
@@ -99,12 +103,38 @@ class KernelDiscoveryTests(unittest.TestCase):
     def setUp(self):
         self.env_patch = test_env()
         self.env_patch.start()
-        self.sample_kernel_dir = install_sample_kernel(
-            pjoin(paths.jupyter_data_dir(), 'kernels'))
-        self.prov_sample1_kernel_dir = install_sample_kernel(
-            pjoin(paths.jupyter_data_dir(), 'kernels'), 'dummy_kspec1', 'dummy_kspec.json')
-        self.prov_sample2_kernel_dir = install_sample_kernel(
-            pjoin(paths.jupyter_data_dir(), 'kernels'), 'dummy_kspec2', 'dummy_kspec.json')
+        install_sample_kernel(pjoin(paths.jupyter_data_dir(), 'kernels'))
+        install_sample_kernel(pjoin(paths.jupyter_data_dir(), 'kernels'), 'dummy_kspec1', 'dummy_kspec.json')
+        install_sample_kernel(pjoin(paths.jupyter_data_dir(), 'kernels'), 'dummy_kspec2', 'dummy_kspec.json')
+
+        # This provides an example of what a kernel provider might do for describing the launch parameters
+        # it supports.  By creating the metadata in the form of JSON schema, applications can easily build
+        # forms that gather the values.
+        # Note that not all parameters are fed to `argv`.  Some may be used by the provider
+        # to configure an environment (e.g., a kubernetes pod) in which the kernel will run.  The idea
+        # is that the front-end will get the parameter metadata, consume and prompt for values, and return
+        # the launch_parameters (name, value pairs) in the kernel startup POST json body, which then
+        # gets passed into the kernel provider's launch method.
+        #
+        # See test_kernel_launch_params() for usage.
+
+        params_json = {'argv': ['tail', '{follow}', '-n {line_count}', '{connection_file}'],
+                       'display_name': 'Test kernel',
+                       'metadata': {
+                           'launch_parameter_schema': {
+                             "title": "Params_kspec Kernel Provider Launch Parameter Schema",
+                             "properties": {
+                               "line_count": {"type": "integer", "minimum": 1, "default": 20, "description": "The number of lines to tail"},
+                               "follow": {"type": "string", "enum": ["-f", "-F"], "default": "-f", "description": "The follow option to tail"},
+                               "cpus": {"type": "number", "minimum": 0.5, "maximum": 8.0, "default": 4.0, "description": "The number of CPUs to use for this kernel"},
+                               "memory": {"type": "integer", "minimum": 2, "maximum": 1024, "default": 8, "description": "The number of GB to reserve for memory for this kernel"}
+                             },
+                             "required": ["line_count", "follow"]
+                           }
+                        }
+                      }
+        install_sample_kernel(pjoin(paths.jupyter_data_dir(), 'kernels'), 'params_kspec', 'params_kspec.json',
+                              kernel_json=params_json)
 
     def tearDown(self):
         self.env_patch.stop()
@@ -165,6 +195,58 @@ def test_kernel_spec_provider_subclass():
 
         conn_info, manager = kf.launch('dummy_kspec/dummy_kspec1')
         assert isinstance(manager, DummyKernelManager)
+        manager.kill()  # no process was started, so this is only for completeness
+
+    @staticmethod
+    def test_kernel_launch_params():
+        kf = discovery.KernelFinder(providers=[LaunchParamsKernelProvider()])
+
+        kspecs = list(kf.find_kernels())
+
+        count = 0
+        param_spec = None
+        for name, spec in kspecs:
+            if name == 'params_kspec/params_kspec':
+                param_spec = spec
+                count += 1
+
+        assert count == 1
+        assert param_spec['argv'] == ['tail', '{follow}', '-n {line_count}', '{connection_file}']
+
+        # application gathers launch parameters here... Since this is full schema, application will likely
+        # just access: param_spec['metadata']['launch_parameter_schema']
+        #
+        line_count_schema = param_spec['metadata']['launch_parameter_schema']['properties']['line_count']
+        follow_schema = param_spec['metadata']['launch_parameter_schema']['properties']['follow']
+        cpus_schema = param_spec['metadata']['launch_parameter_schema']['properties']['cpus']
+        memory_schema = param_spec['metadata']['launch_parameter_schema']['properties']['memory']
+
+        # validate we have our metadata
+        assert line_count_schema['minimum'] == 1
+        assert follow_schema['default'] == '-f'
+        assert cpus_schema['maximum'] == 8.0
+        assert memory_schema['description'] == "The number of GB to reserve for memory for this kernel"
+
+        # Kernel provider would be responsible for validating values against the schema upon return from client.
+        # This includes setting any default values for parameters that were not included, etc.  The following
+        # simulates the parameter gathering...
+        launch_params = dict()
+        launch_params['follow'] = follow_schema['enum'][0]
+        launch_params['line_count'] = 8
+        launch_params['cpus'] = cpus_schema['default']
+        # add a "system-owned" parameter - connection_file - ensure this value is NOT substituted.
+        launch_params['connection_file'] = 'bad_param'
+
+        conn_info, manager = kf.launch('params_kspec/params_kspec', launch_params=launch_params)
+        assert isinstance(manager, KernelManager)
+
+        # confirm argv substitutions
+        assert manager.kernel.args[1] == '-f'
+        assert manager.kernel.args[2] == '-n 8'
+        assert manager.kernel.args[3] != 'bad_param'
+
+        # this actually starts a tail -f command, so let's make sure its terminated
+        manager.kill()
 
     def test_load_config(self):
         # create fake application

jupyter_kernel_mgmt/tests/test_kernelspec.py

@@ -5,6 +5,7 @@
 # Distributed under the terms of the Modified BSD License.
 
 import io
+import copy
 import json
 from logging import StreamHandler
 import os
@@ -27,16 +28,17 @@
 
 sample_kernel_json = {'argv':['cat', '{connection_file}'],
                       'display_name':'Test kernel',
+                      'metadata': {}
                      }
 
 
-def install_sample_kernel(kernels_dir, kernel_name='sample', kernel_file='kernel.json'):
+def install_sample_kernel(kernels_dir, kernel_name='sample', kernel_file='kernel.json', kernel_json=sample_kernel_json):
     """install a sample kernel in a kernels directory"""
     sample_kernel_dir = pjoin(kernels_dir, kernel_name)
     os.makedirs(sample_kernel_dir)
     json_file = pjoin(sample_kernel_dir, kernel_file)
     with open(json_file, 'w') as f:
-        json.dump(sample_kernel_json, f)
+        json.dump(kernel_json, f)
     return sample_kernel_dir