Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Kernel Provisioner Classes"""
3# Copyright (c) Jupyter Development Team.
4# Distributed under the terms of the Modified BSD License.
5import asyncio
6import os
7import pathlib
8import signal
9import sys
10from typing import TYPE_CHECKING, Any
12import zmq
14from ..connect import KernelConnectionInfo, LocalPortCache
15from ..launcher import launch_kernel
16from ..localinterfaces import is_local_ip, local_ips
17from .provisioner_base import KernelProvisionerBase
20class LocalProvisioner(KernelProvisionerBase):
21 """
22 :class:`LocalProvisioner` is a concrete class of ABC :py:class:`KernelProvisionerBase`
23 and is the out-of-box default implementation used when no kernel provisioner is
24 specified in the kernel specification (``kernel.json``). It provides functional
25 parity to existing applications by launching the kernel locally and using
26 :class:`subprocess.Popen` to manage its lifecycle.
28 This class is intended to be subclassed for customizing local kernel environments
29 and serve as a reference implementation for other custom provisioners.
30 """
32 process = None
33 _exit_future = None
34 pid = None
35 pgid = None
36 ip = None
37 ports_cached = False
38 cwd = None
40 @property
41 def has_process(self) -> bool:
42 return self.process is not None
44 async def poll(self) -> int | None:
45 """Poll the provisioner."""
46 ret = 0
47 if self.process:
48 ret = self.process.poll() # type:ignore[unreachable]
49 return ret
51 async def wait(self) -> int | None:
52 """Wait for the provisioner process."""
53 ret = 0
54 if self.process:
55 # Use busy loop at 100ms intervals, polling until the process is
56 # not alive. If we find the process is no longer alive, complete
57 # its cleanup via the blocking wait(). Callers are responsible for
58 # issuing calls to wait() using a timeout (see kill()).
59 while await self.poll() is None: # type:ignore[unreachable]
60 await asyncio.sleep(0.1)
62 # Process is no longer alive, wait and clear
63 ret = self.process.wait()
64 # Make sure all the fds get closed.
65 for attr in ["stdout", "stderr", "stdin"]:
66 fid = getattr(self.process, attr)
67 if fid:
68 fid.close()
69 self.process = None # allow has_process to now return False
70 return ret
72 async def send_signal(self, signum: int) -> None:
73 """Sends a signal to the process group of the kernel (this
74 usually includes the kernel and any subprocesses spawned by
75 the kernel).
77 Note that since only SIGTERM is supported on Windows, we will
78 check if the desired signal is for interrupt and apply the
79 applicable code on Windows in that case.
80 """
81 if self.process:
82 if signum == signal.SIGINT and sys.platform == "win32": # type:ignore[unreachable]
83 from ..win_interrupt import send_interrupt
85 send_interrupt(self.process.win32_interrupt_event)
86 return
88 # Prefer process-group over process
89 if self.pgid and hasattr(os, "killpg"):
90 try:
91 os.killpg(self.pgid, signum)
92 return
93 except OSError:
94 pass # We'll retry sending the signal to only the process below
96 # If we're here, send the signal to the process and let caller handle exceptions
97 self.process.send_signal(signum)
98 return
100 async def kill(self, restart: bool = False) -> None:
101 """Kill the provisioner and optionally restart."""
102 if self.process:
103 if hasattr(signal, "SIGKILL"): # type:ignore[unreachable]
104 # If available, give preference to signalling the process-group over `kill()`.
105 try:
106 await self.send_signal(signal.SIGKILL)
107 return
108 except OSError:
109 pass
110 try:
111 self.process.kill()
112 except OSError as e:
113 LocalProvisioner._tolerate_no_process(e)
115 async def terminate(self, restart: bool = False) -> None:
116 """Terminate the provisioner and optionally restart."""
117 if self.process:
118 if hasattr(signal, "SIGTERM"): # type:ignore[unreachable]
119 # If available, give preference to signalling the process group over `terminate()`.
120 try:
121 await self.send_signal(signal.SIGTERM)
122 return
123 except OSError:
124 pass
125 try:
126 self.process.terminate()
127 except OSError as e:
128 LocalProvisioner._tolerate_no_process(e)
130 @staticmethod
131 def _tolerate_no_process(os_error: OSError) -> None:
132 # In Windows, we will get an Access Denied error if the process
133 # has already terminated. Ignore it.
134 if sys.platform == "win32":
135 if os_error.winerror != 5:
136 err_message = f"Invalid Error, expecting error number to be 5, got {os_error}"
137 raise ValueError(err_message)
139 # On Unix, we may get an ESRCH error (or ProcessLookupError instance) if
140 # the process has already terminated. Ignore it.
141 else:
142 from errno import ESRCH
144 if not isinstance(os_error, ProcessLookupError) or os_error.errno != ESRCH:
145 err_message = (
146 f"Invalid Error, expecting ProcessLookupError or ESRCH, got {os_error}"
147 )
148 raise ValueError(err_message)
150 async def cleanup(self, restart: bool = False) -> None:
151 """Clean up the resources used by the provisioner and optionally restart."""
152 if self.ports_cached and not restart:
153 # provisioner is about to be destroyed, return cached ports
154 lpc = LocalPortCache.instance()
155 ports = (
156 self.connection_info["shell_port"],
157 self.connection_info["iopub_port"],
158 self.connection_info["stdin_port"],
159 self.connection_info["hb_port"],
160 self.connection_info["control_port"],
161 )
162 for port in ports:
163 if TYPE_CHECKING:
164 assert isinstance(port, int)
165 lpc.return_port(port)
167 async def pre_launch(self, **kwargs: Any) -> dict[str, Any]:
168 """Perform any steps in preparation for kernel process launch.
170 This includes applying additional substitutions to the kernel launch command and env.
171 It also includes preparation of launch parameters.
173 Returns the updated kwargs.
174 """
176 # This should be considered temporary until a better division of labor can be defined.
177 km = self.parent
178 if km:
179 transport_encryption = kwargs.pop(
180 "transport_encryption", getattr(km, "transport_encryption", "disabled")
181 )
182 transport_encryption_policy = (
183 km._transport_encryption_policy(transport_encryption)
184 if hasattr(km, "_transport_encryption_policy")
185 else ("auto" if bool(transport_encryption) else "disabled")
186 )
187 encryption_required = transport_encryption_policy == "required"
188 encryption_enabled = transport_encryption_policy in {"auto", "required"}
189 curve_publickey: bytes | None = None
190 curve_secretkey: bytes | None = None
191 if encryption_required and km.transport != "tcp":
192 msg = "transport_encryption='required' is only supported when transport='tcp'."
193 raise RuntimeError(msg)
194 if km.transport == "tcp" and not is_local_ip(km.ip):
195 msg = (
196 "Can only launch a kernel on a local interface. "
197 f"This one is not: {km.ip}."
198 "Make sure that the '*_address' attributes are "
199 "configured properly. "
200 f"Currently valid addresses are: {local_ips()}"
201 )
202 raise RuntimeError(msg)
203 # build the Popen cmd
204 extra_arguments = kwargs.pop("extra_arguments", [])
206 # write connection file / get default ports
207 # TODO - change when handshake pattern is adopted
208 if km.cache_ports and not self.ports_cached:
209 lpc = LocalPortCache.instance()
210 km.shell_port = lpc.find_available_port(km.ip)
211 km.iopub_port = lpc.find_available_port(km.ip)
212 km.stdin_port = lpc.find_available_port(km.ip)
213 km.hb_port = lpc.find_available_port(km.ip)
214 km.control_port = lpc.find_available_port(km.ip)
215 self.ports_cached = True
217 if encryption_enabled and km.transport == "tcp":
218 kernel_curve_ok = encryption_required or (
219 hasattr(km, "_kernel_supports_curve_encryption")
220 and km._kernel_supports_curve_encryption()
221 )
222 if kernel_curve_ok:
223 if km.curve_publickey is None:
224 curve_publickey, curve_secretkey = zmq.curve_keypair()
225 km.curve_publickey = curve_publickey
226 km.curve_secretkey = curve_secretkey
227 else:
228 # Reuse existing keys across restart (same as session.key).
229 # The connection file is preserved on restart, so the kernel
230 # process will read the same keys the manager already holds.
231 curve_publickey = km.curve_publickey
232 curve_secretkey = km.curve_secretkey
233 if "env" in kwargs:
234 jupyter_session = kwargs["env"].get("JPY_SESSION_NAME", "")
235 km.write_connection_file(jupyter_session=jupyter_session)
236 else:
237 km.write_connection_file()
238 self.connection_info = km.get_connection_info()
240 kernel_cmd = km.format_kernel_cmd(
241 extra_arguments=extra_arguments
242 ) # This needs to remain here for b/c
243 else:
244 extra_arguments = kwargs.pop("extra_arguments", [])
245 kernel_cmd = self.kernel_spec.argv + extra_arguments
247 return await super().pre_launch(cmd=kernel_cmd, **kwargs)
249 async def launch_kernel(self, cmd: list[str], **kwargs: Any) -> KernelConnectionInfo:
250 """Launch a kernel with a command."""
252 scrubbed_kwargs = LocalProvisioner._scrub_kwargs(kwargs)
253 self.process = launch_kernel(cmd, **scrubbed_kwargs)
254 pgid = None
255 if hasattr(os, "getpgid"):
256 try:
257 pgid = os.getpgid(self.process.pid)
258 except OSError:
259 pass
261 self.pid = self.process.pid
262 self.pgid = pgid
263 self.cwd = kwargs.get("cwd", pathlib.Path.cwd())
264 return self.connection_info
266 def resolve_path(self, path_str: str) -> str | None:
267 """Resolve path to given file."""
268 path = pathlib.Path(path_str).expanduser()
269 if not path.is_absolute() and self.cwd:
270 path = (pathlib.Path(self.cwd) / path).resolve()
271 if path.exists():
272 return path.as_posix()
273 return None
275 @staticmethod
276 def _scrub_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
277 """Remove any keyword arguments that Popen does not tolerate."""
278 keywords_to_scrub: list[str] = ["extra_arguments", "kernel_id"]
279 scrubbed_kwargs = kwargs.copy()
280 for kw in keywords_to_scrub:
281 scrubbed_kwargs.pop(kw, None)
282 return scrubbed_kwargs
284 async def get_provisioner_info(self) -> dict:
285 """Captures the base information necessary for persistence relative to this instance."""
286 provisioner_info = await super().get_provisioner_info()
287 provisioner_info.update({"pid": self.pid, "pgid": self.pgid, "ip": self.ip})
288 return provisioner_info
290 async def load_provisioner_info(self, provisioner_info: dict) -> None:
291 """Loads the base information necessary for persistence relative to this instance."""
292 await super().load_provisioner_info(provisioner_info)
293 self.pid = provisioner_info["pid"]
294 self.pgid = provisioner_info["pgid"]
295 self.ip = provisioner_info["ip"]