1"""Kernel Provisioner Classes"""
2
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
11
12import zmq
13
14from ..connect import KernelConnectionInfo, LocalPortCache
15from ..launcher import launch_kernel
16from ..localinterfaces import is_local_ip, local_ips
17from .provisioner_base import KernelProvisionerBase
18
19
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.
27
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 """
31
32 process = None
33 _exit_future = None
34 pid = None
35 pgid = None
36 ip = None
37 ports_cached = False
38 cwd = None
39
40 @property
41 def has_process(self) -> bool:
42 return self.process is not None
43
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
50
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)
61
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
71
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).
76
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
84
85 send_interrupt(self.process.win32_interrupt_event)
86 return
87
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
95
96 # If we're here, send the signal to the process and let caller handle exceptions
97 self.process.send_signal(signum)
98 return
99
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)
114
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)
129
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)
138
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
143
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)
149
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)
166
167 async def pre_launch(self, **kwargs: Any) -> dict[str, Any]:
168 """Perform any steps in preparation for kernel process launch.
169
170 This includes applying additional substitutions to the kernel launch command and env.
171 It also includes preparation of launch parameters.
172
173 Returns the updated kwargs.
174 """
175
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 # CurveZMQ encryption is supported over the tcp and ipc transports.
190 # inproc is in-process (no wire handshake) and is never encrypted.
191 encryption_supported_transport = km.transport in {"tcp", "ipc"}
192 curve_publickey: bytes | None = None
193 curve_secretkey: bytes | None = None
194 if encryption_required and not encryption_supported_transport:
195 msg = (
196 "transport_encryption='required' is only supported when "
197 "transport is 'tcp' or 'ipc'."
198 )
199 raise RuntimeError(msg)
200 if km.transport == "tcp" and not is_local_ip(km.ip):
201 msg = (
202 "Can only launch a kernel on a local interface. "
203 f"This one is not: {km.ip}."
204 "Make sure that the '*_address' attributes are "
205 "configured properly. "
206 f"Currently valid addresses are: {local_ips()}"
207 )
208 raise RuntimeError(msg)
209 # build the Popen cmd
210 extra_arguments = kwargs.pop("extra_arguments", [])
211
212 # write connection file / get default ports
213 # TODO - change when handshake pattern is adopted
214 if km.cache_ports and not self.ports_cached:
215 lpc = LocalPortCache.instance()
216 km.shell_port = lpc.find_available_port(km.ip)
217 km.iopub_port = lpc.find_available_port(km.ip)
218 km.stdin_port = lpc.find_available_port(km.ip)
219 km.hb_port = lpc.find_available_port(km.ip)
220 km.control_port = lpc.find_available_port(km.ip)
221 self.ports_cached = True
222
223 if encryption_enabled and encryption_supported_transport:
224 kernel_curve_ok = encryption_required or (
225 hasattr(km, "_kernel_supports_curve_encryption")
226 and km._kernel_supports_curve_encryption()
227 )
228 if kernel_curve_ok:
229 if km.curve_publickey is None:
230 curve_publickey, curve_secretkey = zmq.curve_keypair()
231 km.curve_publickey = curve_publickey
232 km.curve_secretkey = curve_secretkey
233 else:
234 # Reuse existing keys across restart (same as session.key).
235 # The connection file is preserved on restart, so the kernel
236 # process will read the same keys the manager already holds.
237 curve_publickey = km.curve_publickey
238 curve_secretkey = km.curve_secretkey
239 if "env" in kwargs:
240 jupyter_session = kwargs["env"].get("JPY_SESSION_NAME", "")
241 km.write_connection_file(jupyter_session=jupyter_session)
242 else:
243 km.write_connection_file()
244 self.connection_info = km.get_connection_info()
245
246 kernel_cmd = km.format_kernel_cmd(
247 extra_arguments=extra_arguments
248 ) # This needs to remain here for b/c
249 else:
250 extra_arguments = kwargs.pop("extra_arguments", [])
251 kernel_cmd = self.kernel_spec.argv + extra_arguments
252
253 return await super().pre_launch(cmd=kernel_cmd, **kwargs)
254
255 async def launch_kernel(self, cmd: list[str], **kwargs: Any) -> KernelConnectionInfo:
256 """Launch a kernel with a command."""
257
258 scrubbed_kwargs = LocalProvisioner._scrub_kwargs(kwargs)
259 self.process = launch_kernel(cmd, **scrubbed_kwargs)
260 pgid = None
261 if hasattr(os, "getpgid"):
262 try:
263 pgid = os.getpgid(self.process.pid)
264 except OSError:
265 pass
266
267 self.pid = self.process.pid
268 self.pgid = pgid
269 self.cwd = kwargs.get("cwd", pathlib.Path.cwd())
270 return self.connection_info
271
272 def resolve_path(self, path_str: str) -> str | None:
273 """Resolve path to given file."""
274 path = pathlib.Path(path_str).expanduser()
275 if not path.is_absolute() and self.cwd:
276 path = (pathlib.Path(self.cwd) / path).resolve()
277 if path.exists():
278 return path.as_posix()
279 return None
280
281 @staticmethod
282 def _scrub_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
283 """Remove any keyword arguments that Popen does not tolerate."""
284 keywords_to_scrub: list[str] = ["extra_arguments", "kernel_id"]
285 scrubbed_kwargs = kwargs.copy()
286 for kw in keywords_to_scrub:
287 scrubbed_kwargs.pop(kw, None)
288 return scrubbed_kwargs
289
290 async def get_provisioner_info(self) -> dict:
291 """Captures the base information necessary for persistence relative to this instance."""
292 provisioner_info = await super().get_provisioner_info()
293 provisioner_info.update({"pid": self.pid, "pgid": self.pgid, "ip": self.ip})
294 return provisioner_info
295
296 async def load_provisioner_info(self, provisioner_info: dict) -> None:
297 """Loads the base information necessary for persistence relative to this instance."""
298 await super().load_provisioner_info(provisioner_info)
299 self.pid = provisioner_info["pid"]
300 self.pgid = provisioner_info["pgid"]
301 self.ip = provisioner_info["ip"]