"""Tests for the notebook kernel and session manager.""" import asyncio import concurrent.futures import os import sys import uuid from asyncio import ensure_future from subprocess import PIPE from unittest import TestCase import pytest from jupyter_core import paths from tornado.testing import AsyncTestCase, gen_test from traitlets.config.loader import Config from jupyter_client import AsyncKernelManager, KernelManager from jupyter_client.localinterfaces import localhost from jupyter_client.multikernelmanager import AsyncMultiKernelManager, MultiKernelManager from .utils import ( AsyncKMSubclass, AsyncMKMSubclass, SyncKMSubclass, SyncMKMSubclass, install_kernel, skip_win32, ) TIMEOUT = 30 async def now(awaitable): """Use this function ensure that this awaitable happens before other awaitables defined after it. """ (out,) = await asyncio.gather(awaitable) return out class TestKernelManager(TestCase): # static so picklable for multiprocessing on Windows @staticmethod def _get_tcp_km(): c = Config() km = MultiKernelManager(config=c) return km @staticmethod def _get_tcp_km_sub(): c = Config() km = SyncMKMSubclass(config=c) return km # static so picklable for multiprocessing on Windows @staticmethod def _get_ipc_km(): c = Config() c.KernelManager.transport = "ipc" c.KernelManager.ip = "test" km = MultiKernelManager(config=c) return km # static so picklable for multiprocessing on Windows @staticmethod def _run_lifecycle(km, test_kid=None): if test_kid: kid = km.start_kernel(stdout=PIPE, stderr=PIPE, kernel_id=test_kid) assert kid == test_kid else: kid = km.start_kernel(stdout=PIPE, stderr=PIPE) assert km.is_alive(kid) assert km.get_kernel(kid).ready.done() assert kid in km assert kid in km.list_kernel_ids() assert len(km) == 1, f"{len(km)} != {1}" km.restart_kernel(kid, now=True) assert km.is_alive(kid) assert kid in km.list_kernel_ids() km.interrupt_kernel(kid) k = km.get_kernel(kid) kc = k.client() assert isinstance(k, KernelManager) km.shutdown_kernel(kid, now=True) assert kid not in km, f"{kid} not in {km}" kc.stop_channels() def _run_cinfo(self, km, transport, ip): kid = km.start_kernel(stdout=PIPE, stderr=PIPE) km.get_kernel(kid) cinfo = km.get_connection_info(kid) self.assertEqual(transport, cinfo["transport"]) self.assertEqual(ip, cinfo["ip"]) self.assertTrue("stdin_port" in cinfo) self.assertTrue("iopub_port" in cinfo) stream = km.connect_iopub(kid) stream.close() self.assertTrue("shell_port" in cinfo) stream = km.connect_shell(kid) stream.close() self.assertTrue("hb_port" in cinfo) stream = km.connect_hb(kid) stream.close() km.shutdown_kernel(kid, now=True) # static so picklable for multiprocessing on Windows @classmethod def test_tcp_lifecycle(cls): km = cls._get_tcp_km() cls._run_lifecycle(km) def test_tcp_lifecycle_with_kernel_id(self): km = self._get_tcp_km() self._run_lifecycle(km, test_kid=str(uuid.uuid4())) def test_shutdown_all(self): km = self._get_tcp_km() kid = km.start_kernel(stdout=PIPE, stderr=PIPE) self.assertIn(kid, km) km.shutdown_all() self.assertNotIn(kid, km) # shutdown again is okay, because we have no kernels km.shutdown_all() def test_tcp_cinfo(self): km = self._get_tcp_km() self._run_cinfo(km, "tcp", localhost()) @skip_win32 def test_ipc_lifecycle(self): km = self._get_ipc_km() self._run_lifecycle(km) @skip_win32 def test_ipc_cinfo(self): km = self._get_ipc_km() self._run_cinfo(km, "ipc", "test") def test_start_sequence_tcp_kernels(self): """Ensure that a sequence of kernel startups doesn't break anything.""" self._run_lifecycle(self._get_tcp_km()) self._run_lifecycle(self._get_tcp_km()) self._run_lifecycle(self._get_tcp_km()) @skip_win32 def test_start_sequence_ipc_kernels(self): """Ensure that a sequence of kernel startups doesn't break anything.""" self._run_lifecycle(self._get_ipc_km()) self._run_lifecycle(self._get_ipc_km()) self._run_lifecycle(self._get_ipc_km()) def tcp_lifecycle_with_loop(self): # Ensure each thread has an event loop loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) self.test_tcp_lifecycle() loop.close() def test_start_parallel_thread_kernels(self): self.test_tcp_lifecycle() with concurrent.futures.ThreadPoolExecutor(max_workers=2) as thread_executor: future1 = thread_executor.submit(self.tcp_lifecycle_with_loop) future2 = thread_executor.submit(self.tcp_lifecycle_with_loop) future1.result() future2.result() @pytest.mark.skipif( (sys.platform == "darwin") and (sys.version_info >= (3, 6)) and (sys.version_info < (3, 8)), reason='"Bad file descriptor" error', ) @pytest.mark.skipif( sys.platform == "linux", reason='Kernel refuses to start in process pool', ) def test_start_parallel_process_kernels(self): self.test_tcp_lifecycle() with concurrent.futures.ThreadPoolExecutor(max_workers=1) as thread_executor: future1 = thread_executor.submit(self.tcp_lifecycle_with_loop) with concurrent.futures.ProcessPoolExecutor(max_workers=1) as process_executor: # Windows tests needs this target to be picklable: future2 = process_executor.submit(self.test_tcp_lifecycle) future2.result() future1.result() def test_subclass_callables(self): km = self._get_tcp_km_sub() km.reset_counts() kid = km.start_kernel(stdout=PIPE, stderr=PIPE) assert km.call_count("start_kernel") == 1 assert isinstance(km.get_kernel(kid), SyncKMSubclass) assert km.get_kernel(kid).call_count("start_kernel") == 1 assert km.get_kernel(kid).call_count("_async_launch_kernel") == 1 assert km.is_alive(kid) assert kid in km assert kid in km.list_kernel_ids() assert len(km) == 1, f"{len(km)} != {1}" km.get_kernel(kid).reset_counts() km.reset_counts() km.restart_kernel(kid, now=True) assert km.call_count("restart_kernel") == 1 assert km.call_count("get_kernel") == 1 assert km.get_kernel(kid).call_count("restart_kernel") == 1 assert km.get_kernel(kid).call_count("_async_shutdown_kernel") == 1 assert km.get_kernel(kid).call_count("_async_interrupt_kernel") == 1 assert km.get_kernel(kid).call_count("_async_kill_kernel") == 1 assert km.get_kernel(kid).call_count("_async_cleanup_resources") == 1 assert km.get_kernel(kid).call_count("_async_launch_kernel") == 1 assert km.is_alive(kid) assert kid in km.list_kernel_ids() km.get_kernel(kid).reset_counts() km.reset_counts() km.interrupt_kernel(kid) assert km.call_count("interrupt_kernel") == 1 assert km.call_count("get_kernel") == 1 assert km.get_kernel(kid).call_count("interrupt_kernel") == 1 km.get_kernel(kid).reset_counts() km.reset_counts() k = km.get_kernel(kid) assert isinstance(k, SyncKMSubclass) assert km.call_count("get_kernel") == 1 km.get_kernel(kid).reset_counts() km.reset_counts() km.shutdown_all(now=True) assert km.call_count("remove_kernel") == 1 assert km.call_count("request_shutdown") == 0 assert km.call_count("finish_shutdown") == 0 assert km.call_count("cleanup_resources") == 0 assert kid not in km, f"{kid} not in {km}" def test_stream_on_recv(self): mkm = self._get_tcp_km() kid = mkm.start_kernel(stdout=PIPE, stderr=PIPE) stream = mkm.connect_iopub(kid) km = mkm.get_kernel(kid) client = km.client() session = km.session called = False def record_activity(msg_list): nonlocal called """Record an IOPub message arriving from a kernel""" idents, fed_msg_list = session.feed_identities(msg_list) msg = session.deserialize(fed_msg_list, content=False) msg_type = msg["header"]["msg_type"] stream.send(msg) called = True stream.on_recv(record_activity) while True: client.kernel_info() import time time.sleep(0.1) if called: break client.stop_channels() km.shutdown_kernel(now=True) class TestAsyncKernelManager(AsyncTestCase): # static so picklable for multiprocessing on Windows @staticmethod def _get_tcp_km(): c = Config() km = AsyncMultiKernelManager(config=c) return km @staticmethod def _get_tcp_km_sub(): c = Config() km = AsyncMKMSubclass(config=c) return km # static so picklable for multiprocessing on Windows @staticmethod def _get_ipc_km(): c = Config() c.KernelManager.transport = "ipc" c.KernelManager.ip = "test" km = AsyncMultiKernelManager(config=c) return km @staticmethod def _get_pending_kernels_km(): c = Config() c.AsyncMultiKernelManager.use_pending_kernels = True km = AsyncMultiKernelManager(config=c) return km # static so picklable for multiprocessing on Windows @staticmethod async def _run_lifecycle(km, test_kid=None): if test_kid: kid = await km.start_kernel(stdout=PIPE, stderr=PIPE, kernel_id=test_kid) assert kid == test_kid else: kid = await km.start_kernel(stdout=PIPE, stderr=PIPE) assert await km.is_alive(kid) assert kid in km assert kid in km.list_kernel_ids() assert len(km) == 1, f"{len(km)} != {1}" # Ensure we can interrupt during a restart. fut = km.restart_kernel(kid, now=True) await km.interrupt_kernel(kid) assert await km.is_alive(kid) await fut assert kid in km.list_kernel_ids() k = km.get_kernel(kid) assert isinstance(k, AsyncKernelManager) await km.shutdown_kernel(kid, now=True) assert kid not in km, f"{kid} not in {km}" async def _run_cinfo(self, km, transport, ip): kid = await km.start_kernel(stdout=PIPE, stderr=PIPE) km.get_kernel(kid) cinfo = km.get_connection_info(kid) self.assertEqual(transport, cinfo["transport"]) self.assertEqual(ip, cinfo["ip"]) self.assertTrue("stdin_port" in cinfo) self.assertTrue("iopub_port" in cinfo) stream = km.connect_iopub(kid) stream.close() self.assertTrue("shell_port" in cinfo) stream = km.connect_shell(kid) stream.close() self.assertTrue("hb_port" in cinfo) stream = km.connect_hb(kid) stream.close() await km.shutdown_kernel(kid, now=True) self.assertNotIn(kid, km) @gen_test async def test_tcp_lifecycle(self): await self.raw_tcp_lifecycle() @gen_test async def test_tcp_lifecycle_with_kernel_id(self): await self.raw_tcp_lifecycle(test_kid=str(uuid.uuid4())) @gen_test async def test_shutdown_all(self): km = self._get_tcp_km() kid = await km.start_kernel(stdout=PIPE, stderr=PIPE) self.assertIn(kid, km) await km.shutdown_all() self.assertNotIn(kid, km) # shutdown again is okay, because we have no kernels await km.shutdown_all() @gen_test(timeout=20) async def test_use_after_shutdown_all(self): km = self._get_tcp_km() kid = await km.start_kernel(stdout=PIPE, stderr=PIPE) self.assertIn(kid, km) await km.shutdown_all() self.assertNotIn(kid, km) # Start another kernel kid = await km.start_kernel(stdout=PIPE, stderr=PIPE) self.assertIn(kid, km) await km.shutdown_all() self.assertNotIn(kid, km) # shutdown again is okay, because we have no kernels await km.shutdown_all() @gen_test(timeout=20) async def test_shutdown_all_while_starting(self): km = self._get_tcp_km() kid_future = asyncio.ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE)) # This is relying on the ordering of the asyncio queue, not sure if guaranteed or not: kid, _ = await asyncio.gather(kid_future, km.shutdown_all()) self.assertNotIn(kid, km) # Start another kernel kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE)) self.assertIn(kid, km) self.assertEqual(len(km), 1) await km.shutdown_all() self.assertNotIn(kid, km) # shutdown again is okay, because we have no kernels await km.shutdown_all() @gen_test async def test_use_pending_kernels(self): km = self._get_pending_kernels_km() kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE)) kernel = km.get_kernel(kid) assert not kernel.ready.done() assert kid in km assert kid in km.list_kernel_ids() assert len(km) == 1, f"{len(km)} != {1}" # Wait for the kernel to start. await kernel.ready await km.restart_kernel(kid, now=True) out = await km.is_alive(kid) assert out assert kid in km.list_kernel_ids() await km.interrupt_kernel(kid) k = km.get_kernel(kid) assert isinstance(k, AsyncKernelManager) await ensure_future(km.shutdown_kernel(kid, now=True)) # Wait for the kernel to shutdown await kernel.ready assert kid not in km, f"{kid} not in {km}" @gen_test async def test_use_pending_kernels_early_restart(self): km = self._get_pending_kernels_km() kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE)) kernel = km.get_kernel(kid) assert not kernel.ready.done() with pytest.raises(RuntimeError): await km.restart_kernel(kid, now=True) await kernel.ready await ensure_future(km.shutdown_kernel(kid, now=True)) # Wait for the kernel to shutdown await kernel.ready assert kid not in km, f"{kid} not in {km}" @gen_test async def test_use_pending_kernels_early_shutdown(self): km = self._get_pending_kernels_km() kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE)) kernel = km.get_kernel(kid) assert not kernel.ready.done() # Try shutting down while the kernel is pending await ensure_future(km.shutdown_kernel(kid, now=True)) # Wait for the kernel to shutdown await kernel.ready assert kid not in km, f"{kid} not in {km}" @gen_test async def test_use_pending_kernels_early_interrupt(self): km = self._get_pending_kernels_km() kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE)) kernel = km.get_kernel(kid) assert not kernel.ready.done() with pytest.raises(RuntimeError): await km.interrupt_kernel(kid) # Now wait for the kernel to be ready. await kernel.ready await ensure_future(km.shutdown_kernel(kid, now=True)) # Wait for the kernel to shutdown await kernel.ready assert kid not in km, f"{kid} not in {km}" @gen_test async def test_tcp_cinfo(self): km = self._get_tcp_km() await self._run_cinfo(km, "tcp", localhost()) @skip_win32 @gen_test async def test_ipc_lifecycle(self): km = self._get_ipc_km() await self._run_lifecycle(km) @skip_win32 @gen_test async def test_ipc_cinfo(self): km = self._get_ipc_km() await self._run_cinfo(km, "ipc", "test") @gen_test async def test_start_sequence_tcp_kernels(self): """Ensure that a sequence of kernel startups doesn't break anything.""" await self._run_lifecycle(self._get_tcp_km()) await self._run_lifecycle(self._get_tcp_km()) await self._run_lifecycle(self._get_tcp_km()) @skip_win32 @gen_test async def test_start_sequence_ipc_kernels(self): """Ensure that a sequence of kernel startups doesn't break anything.""" await self._run_lifecycle(self._get_ipc_km()) await self._run_lifecycle(self._get_ipc_km()) await self._run_lifecycle(self._get_ipc_km()) def tcp_lifecycle_with_loop(self): # Ensure each thread has an event loop loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(self.raw_tcp_lifecycle()) loop.close() # static so picklable for multiprocessing on Windows @classmethod async def raw_tcp_lifecycle(cls, test_kid=None): # Since @gen_test creates an event loop, we need a raw form of # test_tcp_lifecycle that assumes the loop already exists. km = cls._get_tcp_km() await cls._run_lifecycle(km, test_kid=test_kid) # static so picklable for multiprocessing on Windows @classmethod def raw_tcp_lifecycle_sync(cls, test_kid=None): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(cls.raw_tcp_lifecycle(test_kid=test_kid)) loop.close() @gen_test async def test_start_parallel_thread_kernels(self): await self.raw_tcp_lifecycle() with concurrent.futures.ThreadPoolExecutor(max_workers=2) as thread_executor: future1 = thread_executor.submit(self.tcp_lifecycle_with_loop) future2 = thread_executor.submit(self.tcp_lifecycle_with_loop) future1.result() future2.result() @gen_test async def test_start_parallel_process_kernels(self): await self.raw_tcp_lifecycle() with concurrent.futures.ThreadPoolExecutor(max_workers=1) as thread_executor: future1 = thread_executor.submit(self.tcp_lifecycle_with_loop) with concurrent.futures.ProcessPoolExecutor(max_workers=1) as process_executor: # Windows tests needs this target to be picklable: future2 = process_executor.submit(self.raw_tcp_lifecycle_sync) future2.result() future1.result() @gen_test async def test_subclass_callables(self): mkm = self._get_tcp_km_sub() mkm.reset_counts() kid = await mkm.start_kernel(stdout=PIPE, stderr=PIPE) assert mkm.call_count("start_kernel") == 1 assert isinstance(mkm.get_kernel(kid), AsyncKMSubclass) assert mkm.get_kernel(kid).call_count("start_kernel") == 1 assert mkm.get_kernel(kid).call_count("_async_launch_kernel") == 1 assert await mkm.is_alive(kid) assert kid in mkm assert kid in mkm.list_kernel_ids() assert len(mkm) == 1, f"{len(mkm)} != {1}" mkm.get_kernel(kid).reset_counts() mkm.reset_counts() await mkm.restart_kernel(kid, now=True) assert mkm.call_count("restart_kernel") == 1 assert mkm.call_count("get_kernel") == 1 assert mkm.get_kernel(kid).call_count("restart_kernel") == 1 assert mkm.get_kernel(kid).call_count("_async_interrupt_kernel") == 1 assert mkm.get_kernel(kid).call_count("_async_kill_kernel") == 1 assert mkm.get_kernel(kid).call_count("_async_cleanup_resources") == 1 assert mkm.get_kernel(kid).call_count("_async_launch_kernel") == 1 assert await mkm.is_alive(kid) assert kid in mkm.list_kernel_ids() mkm.get_kernel(kid).reset_counts() mkm.reset_counts() await mkm.interrupt_kernel(kid) assert mkm.call_count("interrupt_kernel") == 1 assert mkm.call_count("get_kernel") == 1 assert mkm.get_kernel(kid).call_count("interrupt_kernel") == 1 mkm.get_kernel(kid).reset_counts() mkm.reset_counts() k = mkm.get_kernel(kid) assert isinstance(k, AsyncKMSubclass) assert mkm.call_count("get_kernel") == 1 mkm.get_kernel(kid).reset_counts() mkm.reset_counts() await mkm.shutdown_all(now=True) assert mkm.call_count("remove_kernel") == 1 assert mkm.call_count("_async_request_shutdown") == 0 assert mkm.call_count("_async_finish_shutdown") == 0 assert mkm.call_count("_async_cleanup_resources") == 0 assert kid not in mkm, f"{kid} not in {mkm}" @gen_test async def test_bad_kernelspec(self): km = self._get_tcp_km() install_kernel( os.path.join(paths.jupyter_data_dir(), "kernels"), argv=["non_existent_executable"], name="bad", ) with pytest.raises(FileNotFoundError): await ensure_future(km.start_kernel(kernel_name="bad", stdout=PIPE, stderr=PIPE)) @gen_test async def test_bad_kernelspec_pending(self): km = self._get_pending_kernels_km() install_kernel( os.path.join(paths.jupyter_data_dir(), "kernels"), argv=["non_existent_executable"], name="bad", ) kernel_id = await ensure_future( km.start_kernel(kernel_name="bad", stdout=PIPE, stderr=PIPE) ) with pytest.raises(FileNotFoundError): await km.get_kernel(kernel_id).ready assert kernel_id in km.list_kernel_ids() await ensure_future(km.shutdown_kernel(kernel_id)) assert kernel_id not in km.list_kernel_ids() @gen_test async def test_stream_on_recv(self): mkm = self._get_tcp_km() kid = await mkm.start_kernel(stdout=PIPE, stderr=PIPE) stream = mkm.connect_iopub(kid) km = mkm.get_kernel(kid) client = km.client() session = km.session called = False def record_activity(msg_list): nonlocal called """Record an IOPub message arriving from a kernel""" idents, fed_msg_list = session.feed_identities(msg_list) msg = session.deserialize(fed_msg_list, content=False) msg_type = msg["header"]["msg_type"] stream.send(msg) called = True stream.on_recv(record_activity) while True: await client.kernel_info(reply=True) if called: break await asyncio.sleep(0.1) client.stop_channels() await km.shutdown_kernel(now=True)