From 433b519b8d6f60149276dc6a24a0768053586b72 Mon Sep 17 00:00:00 2001 From: gongruan Date: Thu, 10 Dec 2020 16:08:54 +0800 Subject: [PATCH] Fix the bug that reconnectsparkmagic cause 'Error Starting Kernel' in jupyterlab, add the parameter 'name' in create new session. --- .../sparkmagic/serverextension/handlers.py | 17 +++++++++++++---- sparkmagic/sparkmagic/tests/test_handlers.py | 13 +++++++------ 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/sparkmagic/sparkmagic/serverextension/handlers.py b/sparkmagic/sparkmagic/serverextension/handlers.py index c94a258ca..b228b65e1 100644 --- a/sparkmagic/sparkmagic/serverextension/handlers.py +++ b/sparkmagic/sparkmagic/serverextension/handlers.py @@ -1,4 +1,5 @@ import json +import os from notebook.utils import url_path_join from notebook.base.handlers import IPythonHandler from tornado import web @@ -99,20 +100,25 @@ def _get_kernel_manager(self, path, kernel_name): sessions = self.session_manager.list_sessions() kernel_id = None + name = None for session in sessions: if session['notebook']['path'] == path: session_id = session['id'] kernel_id = session['kernel']['id'] existing_kernel_name = session['kernel']['name'] + if 'name' in session: + name = session['name'] + break if kernel_id is None: self.logger.debug(u"Kernel not found. Starting a new kernel.") - k_m = yield self._get_kernel_manager_new_session(path, kernel_name) + name = os.path.basename(path) + k_m = yield self._get_kernel_manager_new_session(path, kernel_name, name) elif existing_kernel_name != kernel_name: self.logger.debug(u"Existing kernel name '{}' does not match requested '{}'. Starting a new kernel.".format(existing_kernel_name, kernel_name)) self._delete_session(session_id) - k_m = yield self._get_kernel_manager_new_session(path, kernel_name) + k_m = yield self._get_kernel_manager_new_session(path, kernel_name, name) else: self.logger.debug(u"Kernel found. Restarting kernel.") k_m = self.kernel_manager.get_kernel(kernel_id) @@ -121,8 +127,11 @@ def _get_kernel_manager(self, path, kernel_name): raise gen.Return(k_m) @gen.coroutine - def _get_kernel_manager_new_session(self, path, kernel_name): - model_future = self.session_manager.create_session(kernel_name=kernel_name, path=path, type="notebook") + def _get_kernel_manager_new_session(self, path, kernel_name, name=None): + if name is None: + name = os.path.basename(path) + + model_future = self.session_manager.create_session(kernel_name=kernel_name, path=path, name=name, type="notebook") model = yield model_future kernel_id = model["kernel"]["id"] self.logger.debug("Kernel created with id {}".format(str(kernel_id))) diff --git a/sparkmagic/sparkmagic/tests/test_handlers.py b/sparkmagic/sparkmagic/tests/test_handlers.py index 4d01734fd..e4601f0ac 100644 --- a/sparkmagic/sparkmagic/tests/test_handlers.py +++ b/sparkmagic/sparkmagic/tests/test_handlers.py @@ -27,6 +27,7 @@ class TestSparkMagicHandler(AsyncTestCase): kernel_id = '1' kernel_name = 'pysparkkernel' session_id = '1' + session_name = 'some_path.ipynb' username = 'username' password = 'password' endpoint = 'http://endpoint.com' @@ -36,8 +37,8 @@ class TestSparkMagicHandler(AsyncTestCase): bad_msg = dict(content=dict(status='error', ename='SyntaxError', evalue='oh no!')) request = None - def create_session_dict(self, path, kernel_id): - return dict(notebook=dict(path=path), kernel=dict(id=kernel_id, name=self.kernel_name), id=self.session_id) + def create_session_dict(self, path, kernel_id, name): + return dict(notebook=dict(path=path), kernel=dict(id=kernel_id, name=self.kernel_name), id=self.session_id, name=name) def get_argument(self, key): return dict(username=self.username, password=self.password, endpoint=self.endpoint, path=self.path)[key] @@ -53,10 +54,10 @@ def setUp(self): self.kernel_manager.get_kernel = MagicMock(return_value=self.individual_kernel_manager) # Mock session manager - self.session_list = [self.create_session_dict(self.path, self.kernel_id)] + self.session_list = [self.create_session_dict(self.path, self.kernel_id, self.session_name)] self.session_manager = MagicMock() self.session_manager.list_sessions = MagicMock(return_value=self.session_list) - self.session_manager.create_session = MagicMock(return_value=self.create_session_dict(self.path, self.kernel_id)) + self.session_manager.create_session = MagicMock(return_value=self.create_session_dict(self.path, self.kernel_id, self.session_name)) # Mock spark events self.spark_events = MagicMock() @@ -195,7 +196,7 @@ def test_get_kernel_manager_no_existing_kernel(self, _get_kernel_manager_new_ses assert_equals(self.individual_kernel_manager, km) self.individual_kernel_manager.restart_kernel.assert_not_called() self.kernel_manager.get_kernel.assert_not_called() - _get_kernel_manager_new_session.assert_called_once_with(different_path, self.kernel_name) + _get_kernel_manager_new_session.assert_called_once_with(different_path, self.kernel_name, different_path) @patch('sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager_new_session') @gen_test @@ -219,5 +220,5 @@ def test_get_kernel_manager_different_kernel_type(self, _get_kernel_manager_new_ assert_equals(self.individual_kernel_manager, km) self.individual_kernel_manager.restart_kernel.assert_not_called() self.kernel_manager.get_kernel.assert_not_called() - _get_kernel_manager_new_session.assert_called_once_with(self.path, different_kernel) + _get_kernel_manager_new_session.assert_called_once_with(self.path, different_kernel, self.session_name) self.session_manager.delete_session.assert_called_once_with(self.session_id)