diff --git a/jupyter_kernel_mgmt/discovery.py b/jupyter_kernel_mgmt/discovery.py index 19e95c4..e37fe93 100644 --- a/jupyter_kernel_mgmt/discovery.py +++ b/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) diff --git a/jupyter_kernel_mgmt/subproc/launcher.py b/jupyter_kernel_mgmt/subproc/launcher.py index 3688c29..2bec592 100644 --- a/jupyter_kernel_mgmt/subproc/launcher.py +++ b/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,9 +140,15 @@ 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 @@ -148,8 +156,13 @@ def format_kernel_cmd(self, connection_file, kernel_resource_dir=None): 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: diff --git a/jupyter_kernel_mgmt/tests/test_discovery.py b/jupyter_kernel_mgmt/tests/test_discovery.py index bd0443b..69b3baa 100644 --- a/jupyter_kernel_mgmt/tests/test_discovery.py +++ b/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 diff --git a/jupyter_kernel_mgmt/tests/test_kernelspec.py b/jupyter_kernel_mgmt/tests/test_kernelspec.py index 4b08dad..6d9b7a1 100644 --- a/jupyter_kernel_mgmt/tests/test_kernelspec.py +++ b/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