diff --git a/jupyter_server/services/kernels/kernelmanager.py b/jupyter_server/services/kernels/kernelmanager.py index 8f4e8277f9..ee80799c21 100644 --- a/jupyter_server/services/kernels/kernelmanager.py +++ b/jupyter_server/services/kernels/kernelmanager.py @@ -18,6 +18,7 @@ from datetime import datetime, timedelta from functools import partial, wraps +from jupyter_client.asynchronous.client import AsyncKernelClient from jupyter_client.ioloop.manager import AsyncIOLoopKernelManager from jupyter_client.multikernelmanager import AsyncMultiKernelManager, MultiKernelManager from jupyter_client.session import Session @@ -37,11 +38,13 @@ Any, Bool, Dict, + DottedObjectName, Float, Instance, Integer, List, TraitError, + Type, Unicode, default, validate, @@ -845,6 +848,30 @@ async def wrapped_method(self, *args, **kwargs): class ServerKernelManager(AsyncIOLoopKernelManager): """A server-specific kernel manager.""" + # Override parent traits to make them configurable + client_class = DottedObjectName( + "jupyter_client.asynchronous.AsyncKernelClient", + config=True, + help="""The kernel client class to use for communicating with kernels. + + This should be a subclass of KernelClient, and it should accept the + following arguments: + - kernel_manager + - blocking + - loop + """, + ) + + client_factory = Type( + default_value=AsyncKernelClient, + klass="jupyter_client.client.KernelClient", + config=True, + help="""The kernel client factory class to use for creating client instances. + + This should be a subclass of KernelClient. + """, + ) + # Define activity-related attributes: execution_state = Unicode( None, allow_none=True, help="The current execution state of the kernel" diff --git a/tests/services/kernels/test_config.py b/tests/services/kernels/test_config.py index 1db2e11b1f..657f8c8ea1 100644 --- a/tests/services/kernels/test_config.py +++ b/tests/services/kernels/test_config.py @@ -1,7 +1,14 @@ import pytest +from jupyter_client.asynchronous.client import AsyncKernelClient +from jupyter_client.blocking.client import BlockingKernelClient +from jupyter_client.client import KernelClient from traitlets.config import Config -from jupyter_server.services.kernels.kernelmanager import AsyncMappingKernelManager +from jupyter_server.services.kernels.kernelmanager import ( + AsyncMappingKernelManager, + ServerKernelManager, +) +from jupyter_server.utils import import_item @pytest.fixture @@ -29,3 +36,31 @@ def test_not_server_kernel_manager(jp_configurable_serverapp): ] with pytest.warns(FutureWarning, match="is not a subclass of 'ServerKernelManager'"): jp_configurable_serverapp(argv=argv) + + +def test_server_kernel_manager_client_traits_via_config(): + """Test that ServerKernelManager client traits can be configured via Config object.""" + config = Config() + config.ServerKernelManager.client_class = "jupyter_client.blocking.client.BlockingKernelClient" + config.ServerKernelManager.client_factory = BlockingKernelClient + + km = ServerKernelManager(config=config) + assert km.client_class == "jupyter_client.blocking.client.BlockingKernelClient" + assert km.client_factory == BlockingKernelClient + + +def test_server_kernel_manager_client_traits_default_values(): + """Test that ServerKernelManager client traits have correct default values.""" + km = ServerKernelManager() + assert km.client_class == "jupyter_client.asynchronous.AsyncKernelClient" + # Default client_factory should be the AsyncKernelClient class + assert km.client_factory == AsyncKernelClient + + +def test_server_kernel_manager_client_class_string_configuration(): + """Test that client_class can be configured with different string values.""" + config = Config() + config.ServerKernelManager.client_class = "jupyter_client.client.KernelClient" + + km = ServerKernelManager(config=config) + assert km.client_class == "jupyter_client.client.KernelClient"