Skip to content

Commit b5f73cf

Browse files
committed
cocalc-api: jupyter tests, add endpoints to list/stop kernels
1 parent 1aed72c commit b5f73cf

File tree

11 files changed

+527
-41
lines changed

11 files changed

+527
-41
lines changed

src/packages/conat/project/api/system.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,10 @@ export const system = {
3434

3535
// jupyter stateless API
3636
jupyterExecute: true,
37+
38+
// jupyter kernel management
39+
listJupyterKernels: true,
40+
stopJupyterKernel: true,
3741
};
3842

3943
export interface System {
@@ -74,4 +78,9 @@ export interface System {
7478
}) => Promise<void>;
7579

7680
jupyterExecute: (opts: ProjectJupyterApiOptions) => Promise<object[]>;
81+
82+
listJupyterKernels: () => Promise<
83+
{ pid: number; connectionFile: string; kernel_name?: string }[]
84+
>;
85+
stopJupyterKernel: (opts: { pid: number }) => Promise<{ success: boolean }>;
7786
}

src/packages/jupyter/kernel/launch-kernel.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,11 @@ async function launchKernelSpec(
168168
} else {
169169
running_kernel = spawn(argv[0], argv.slice(1), full_spawn_options);
170170
}
171+
172+
// Store kernel info for tracking
173+
running_kernel.connectionFile = connectionFile;
174+
running_kernel.kernel_spec = kernel_spec;
175+
171176
spawned.push(running_kernel);
172177

173178
running_kernel.on("error", (code, signal) => {
@@ -221,6 +226,48 @@ async function ensureDirectoryExists(path: string) {
221226

222227
// Clean up after any children created here
223228
const spawned: any[] = [];
229+
230+
export interface RunningKernel {
231+
pid: number;
232+
connectionFile: string;
233+
kernel_name?: string;
234+
}
235+
236+
export function listRunningKernels(): RunningKernel[] {
237+
return spawned
238+
.filter((child) => child.pid)
239+
.map((child) => ({
240+
pid: child.pid,
241+
connectionFile: child.connectionFile || "unknown",
242+
kernel_name: child.kernel_spec?.name,
243+
}));
244+
}
245+
246+
export function stopKernel(pid: number): boolean {
247+
const index = spawned.findIndex((child) => child.pid === pid);
248+
if (index === -1) {
249+
return false;
250+
}
251+
252+
const child = spawned[index];
253+
try {
254+
// Try to kill the process group first (negative PID)
255+
process.kill(-child.pid, "SIGKILL");
256+
} catch (err) {
257+
// If that fails, try killing the process directly
258+
try {
259+
child.kill("SIGKILL");
260+
} catch (err2) {
261+
logger.debug(`stopKernel: failed to kill ${child.pid}: ${err2}`);
262+
return false;
263+
}
264+
}
265+
266+
// Remove from spawned array
267+
spawned.splice(index, 1);
268+
return true;
269+
}
270+
224271
export function closeAll() {
225272
for (const child of spawned) {
226273
if (child.pid) {

src/packages/next/pages/api/conat/project.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ export default async function handle(req, res) {
4040
throw Error("must specify project_id or use project-specific api key");
4141
}
4242
if (project_id0) {
43-
// auth via project_id
43+
// auth via project-specific API key
4444
if (project_id0 != project_id) {
4545
throw Error("project specific api key must match requested project");
4646
}

src/packages/project/conat/api/system.ts

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ export async function renameFile({ src, dest }: { src: string; dest: string }) {
5050
import { get_configuration } from "@cocalc/project/configuration";
5151
export { get_configuration as configuration };
5252

53-
import { canonical_paths } from "../../browser-websocket/canonical-path";
53+
import { canonical_paths } from "@cocalc/project/browser-websocket/canonical-path";
5454
export { canonical_paths as canonicalPaths };
5555

5656
import ensureContainingDirectoryExists from "@cocalc/backend/misc/ensure-containing-directory-exists";
@@ -107,3 +107,17 @@ export async function signal({
107107

108108
import jupyterExecute from "@cocalc/jupyter/stateless-api/execute";
109109
export { jupyterExecute };
110+
111+
import {
112+
listRunningKernels,
113+
stopKernel,
114+
} from "@cocalc/jupyter/kernel/launch-kernel";
115+
116+
export async function listJupyterKernels() {
117+
return listRunningKernels();
118+
}
119+
120+
export async function stopJupyterKernel({ pid }: { pid: number }) {
121+
const success = stopKernel(pid);
122+
return { success };
123+
}

src/packages/server/projects/control/single-user.ts

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -153,20 +153,33 @@ class Project extends BaseProject {
153153
this.stateChanging = { state: "stopping" };
154154
await this.saveStateToDatabase(this.stateChanging);
155155
const pid = await getProjectPID(this.HOME);
156-
const killProject = () => {
156+
157+
// First attempt: graceful shutdown with SIGTERM
158+
// This allows the process to clean up child processes (e.g., Jupyter kernels)
159+
let usedSigterm = false;
160+
const killProject = (signal: NodeJS.Signals = "SIGTERM") => {
157161
try {
158-
logger.debug(`stop: sending kill -${pid}`);
159-
kill(-pid, "SIGKILL");
162+
logger.debug(`stop: sending kill -${pid} with ${signal}`);
163+
kill(-pid, signal);
160164
} catch (err) {
161165
// expected exception if no pid
162166
logger.debug(`stop: kill err ${err}`);
163167
}
164168
};
165-
killProject();
169+
170+
// Try SIGTERM first for graceful shutdown
171+
killProject("SIGTERM");
172+
usedSigterm = true;
173+
166174
await this.wait({
167175
until: async () => {
168176
if (await isProjectRunning(this.HOME)) {
169-
killProject();
177+
// After 5 seconds, escalate to SIGKILL
178+
if (usedSigterm) {
179+
logger.debug("stop: escalating to SIGKILL");
180+
killProject("SIGKILL");
181+
usedSigterm = false;
182+
}
170183
return false;
171184
} else {
172185
return true;

src/python/cocalc-api/src/cocalc_api/hub.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ def __init__(self, parent: "Hub"):
289289
self._parent = parent
290290

291291
@api_method("jupyter.kernels")
292-
def kernels(self, project_id: Optional[str] = None) -> dict[str, Any]:
292+
def kernels(self, project_id: Optional[str] = None) -> list[dict[str, Any]]:
293293
"""
294294
Get specifications of available Jupyter kernels.
295295
@@ -298,7 +298,18 @@ def kernels(self, project_id: Optional[str] = None) -> dict[str, Any]:
298298
If not given, a global anonymous project may be used.
299299
300300
Returns:
301-
dict[str, Any]: JSON response containing kernel specs.
301+
list[dict[str, Any]]: List of kernel specification objects. Each kernel object
302+
contains information like 'name', 'display_name', 'language', etc.
303+
304+
Examples:
305+
Get available kernels for a project:
306+
307+
>>> import cocalc_api; hub = cocalc_api.Hub(api_key="sk-...")
308+
>>> kernels = hub.jupyter.kernels(project_id='6e75dbf1-0342-4249-9dce-6b21648656e9')
309+
>>> # Extract kernel names
310+
>>> kernel_names = [k['name'] for k in kernels]
311+
>>> 'python3' in kernel_names
312+
True
302313
"""
303314
...
304315

src/python/cocalc-api/src/cocalc_api/project.py

Lines changed: 66 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -116,26 +116,84 @@ def jupyter_execute(
116116
kernel: str,
117117
history: Optional[list[str]] = None,
118118
path: Optional[str] = None,
119-
) -> dict[str, Any]: # type: ignore[empty-body]
119+
) -> list[dict[str, Any]]: # type: ignore[empty-body]
120120
"""
121121
Execute code using a Jupyter kernel.
122122
123123
Args:
124124
input (str): Code to execute.
125-
kernel (str): Name of kernel to use. Get options using jupyter.kernels().
125+
kernel (str): Name of kernel to use. Get options using hub.jupyter.kernels().
126126
history (Optional[list[str]]): Array of previous inputs (they get evaluated every time, but without output being captured).
127127
path (Optional[str]): File path context for execution.
128128
129129
Returns:
130-
dict[str, Any]: JSON response containing execution results.
130+
list[dict[str, Any]]: List of output items. Each output item contains
131+
execution results with 'data' field containing output by MIME type
132+
(e.g., 'text/plain' for text output) or 'name'/'text' fields for
133+
stream output (stdout/stderr).
131134
132135
Examples:
133136
Execute a simple sum using a Jupyter kernel:
134137
135-
>>> import cocalc_api; project = cocalc_api.Project(api_key="sk-...")
136-
>>> project.jupyter.execute(history=['a=100;print(a)'],
137-
input='sum(range(a+1))',
138-
kernel='python3')
139-
{'output': [{'data': {'text/plain': '5050'}}], ...}
138+
>>> import cocalc_api; project = cocalc_api.Project(api_key="sk-...", project_id='...')
139+
>>> result = project.system.jupyter_execute(input='sum(range(100))', kernel='python3')
140+
>>> result
141+
[{'data': {'text/plain': '4950'}}]
142+
143+
Execute with history context:
144+
145+
>>> result = project.system.jupyter_execute(
146+
... history=['a = 100'],
147+
... input='sum(range(a + 1))',
148+
... kernel='python3')
149+
>>> result
150+
[{'data': {'text/plain': '5050'}}]
151+
152+
Print statements produce stream output:
153+
154+
>>> result = project.system.jupyter_execute(input='print("Hello")', kernel='python3')
155+
>>> result
156+
[{'name': 'stdout', 'text': 'Hello\\n'}]
157+
"""
158+
...
159+
160+
@api_method("system.listJupyterKernels")
161+
def list_jupyter_kernels(self) -> list[dict[str, Any]]: # type: ignore[empty-body]
162+
"""
163+
List all running Jupyter kernels in the project.
164+
165+
Returns:
166+
list[dict[str, Any]]: List of running kernels. Each kernel has:
167+
- pid (int): Process ID of the kernel
168+
- connectionFile (str): Path to the kernel connection file
169+
- kernel_name (str, optional): Name of the kernel (e.g., 'python3')
170+
171+
Examples:
172+
List all running kernels:
173+
174+
>>> import cocalc_api; project = cocalc_api.Project(api_key="sk-...", project_id='...')
175+
>>> kernels = project.system.list_jupyter_kernels()
176+
>>> kernels
177+
[{'pid': 12345, 'connectionFile': '/run/user/1000/jupyter/kernel-abc123.json', 'kernel_name': 'python3'}]
178+
"""
179+
...
180+
181+
@api_method("system.stopJupyterKernel")
182+
def stop_jupyter_kernel(self, pid: int) -> dict[str, bool]: # type: ignore[empty-body]
183+
"""
184+
Stop a specific Jupyter kernel by process ID.
185+
186+
Args:
187+
pid (int): Process ID of the kernel to stop
188+
189+
Returns:
190+
dict[str, bool]: Dictionary with 'success' key indicating if the kernel was stopped
191+
192+
Examples:
193+
Stop a kernel by PID:
194+
195+
>>> import cocalc_api; project = cocalc_api.Project(api_key="sk-...", project_id='...')
196+
>>> project.system.stop_jupyter_kernel(pid=12345)
197+
{'success': True}
140198
"""
141199
...

0 commit comments

Comments
 (0)