Skip to content

Commit

Permalink
Set JPY_SESSION_NAME to full notebook path.
Browse files Browse the repository at this point in the history
This also add some typing here and there, and extend one of the console
warning to log an exception when there is an error.

My main concern is that get_kernel_env need to become async.

Co-authored-by: Kevin Bates <kbates4@gmail.com>
  • Loading branch information
Carreau and kevin-bates committed Jan 4, 2023
1 parent 8e4b7ed commit a66e8df
Show file tree
Hide file tree
Showing 9 changed files with 119 additions and 42 deletions.
2 changes: 1 addition & 1 deletion jupyter_server/base/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -735,7 +735,7 @@ def write_error(self, status_code, **kwargs):
reply["message"] = "Unhandled error"
reply["reason"] = None
reply["traceback"] = "".join(traceback.format_exception(*exc_info))
self.log.warning("wrote error: %r", reply["message"])
self.log.warning("wrote error: %r", reply["message"], exc_info=True)
self.finish(json.dumps(reply))

def get_login_url(self):
Expand Down
4 changes: 2 additions & 2 deletions jupyter_server/gateway/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def remove_kernel(self, kernel_id):
except KeyError:
pass

async def start_kernel(self, kernel_id=None, path=None, **kwargs):
async def start_kernel(self, *, kernel_id=None, path=None, **kwargs):
"""Start a kernel for a session and return its kernel_id.
Parameters
Expand Down Expand Up @@ -323,7 +323,7 @@ class GatewaySessionManager(SessionManager):

kernel_manager = Instance("jupyter_server.gateway.managers.GatewayMappingKernelManager")

async def kernel_culled(self, kernel_id: str) -> bool:
async def kernel_culled(self, kernel_id: str) -> bool: # typing: ignore
"""Checks if the kernel is still considered alive and returns true if it's not found."""
km: Optional[GatewayKernelManager] = None
try:
Expand Down
4 changes: 2 additions & 2 deletions jupyter_server/services/contents/fileio.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from traitlets import Bool
from traitlets.config import Configurable

from jupyter_server.utils import to_api_path, to_os_path
from jupyter_server.utils import ApiPath, to_api_path, to_os_path


def replace_file(src, dst):
Expand Down Expand Up @@ -257,7 +257,7 @@ def _get_os_path(self, path):
# to_os_path is not safe if path starts with a drive, since os.path.join discards first part
if os.path.splitdrive(path)[0]:
raise HTTPError(404, "%s is not a relative API path" % path)
os_path = to_os_path(path, root)
os_path = to_os_path(ApiPath(path), root)
if not (os.path.abspath(os_path) + os.path.sep).startswith(root):
raise HTTPError(404, "%s is outside root contents directory" % path)
return os_path
Expand Down
20 changes: 15 additions & 5 deletions jupyter_server/services/kernels/kernelmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,13 @@
# Distributed under the terms of the Modified BSD License.
import asyncio
import os
import typing as t
import warnings
from collections import defaultdict
from datetime import datetime, timedelta
from functools import partial
from typing import Dict as DictType
from typing import Optional

from jupyter_client.ioloop.manager import AsyncIOLoopKernelManager
from jupyter_client.multikernelmanager import AsyncMultiKernelManager, MultiKernelManager
Expand All @@ -36,7 +39,7 @@

from jupyter_server._tz import isoformat, utcnow
from jupyter_server.prometheus.metrics import KERNEL_CURRENTLY_RUNNING_TOTAL
from jupyter_server.utils import import_item, to_os_path
from jupyter_server.utils import ApiPath, import_item, to_os_path


class MappingKernelManager(MultiKernelManager):
Expand All @@ -56,7 +59,7 @@ def _default_kernel_manager_class(self):

_kernel_connections = Dict()

_kernel_ports = Dict()
_kernel_ports: DictType[str, t.List[int]] = Dict() # type: ignore

_culler_callback = None

Expand Down Expand Up @@ -196,12 +199,16 @@ async def _remove_kernel_when_ready(self, kernel_id, kernel_awaitable):
self._kernel_connections.pop(kernel_id, None)
self._kernel_ports.pop(kernel_id, None)

async def _async_start_kernel(self, kernel_id=None, path=None, **kwargs):
# TODO DEC 2022: Revise the type-ignore once the signatures have been changed upstream
# https://github.com/jupyter/jupyter_client/pull/905
async def _async_start_kernel( # type:ignore[override]
self, *, kernel_id: Optional[str] = None, path: Optional[ApiPath] = None, **kwargs: str
) -> str:
"""Start a kernel for a session and return its kernel_id.
Parameters
----------
kernel_id : uuid
kernel_id : uuid (str)
The uuid to associate the new kernel with. If this
is not None, this kernel will be persistent whenever it is
requested.
Expand All @@ -216,6 +223,7 @@ async def _async_start_kernel(self, kernel_id=None, path=None, **kwargs):
if path is not None:
kwargs["cwd"] = self.cwd_for_path(path, env=kwargs.get("env", {}))
if kernel_id is not None:
assert kernel_id is not None, "Never Fail, but necessary for mypy "
kwargs["kernel_id"] = kernel_id
kernel_id = await self.pinned_superclass._async_start_kernel(self, **kwargs)
self._kernel_connections[kernel_id] = 0
Expand All @@ -242,7 +250,7 @@ async def _async_start_kernel(self, kernel_id=None, path=None, **kwargs):
# Initialize culling if not already
if not self._initialized_culler:
self.initialize_culler()

assert kernel_id is not None
return kernel_id

start_kernel = _async_start_kernel
Expand Down Expand Up @@ -299,6 +307,8 @@ def _get_changed_ports(self, kernel_id):
"""
# Get current ports and return comparison with ports captured at startup.
km = self.get_kernel(kernel_id)
assert isinstance(km.ports, list)
assert isinstance(self._kernel_ports[kernel_id], list)
if km.ports != self._kernel_ports[kernel_id]:
return km.ports
return None
Expand Down
9 changes: 7 additions & 2 deletions jupyter_server/services/sessions/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,12 +52,17 @@ async def post(self):
if model is None:
raise web.HTTPError(400, "No JSON data provided")

if "notebook" in model and "path" in model["notebook"]:
if "notebook" in model:
self.log.warning("Sessions API changed, see updated swagger docs")
model["path"] = model["notebook"]["path"]
model["type"] = "notebook"
if "name" in model["notebook"]:
model["path"] = model["notebook"]["name"]
elif "path" in model["notebook"]:
model["path"] = model["notebook"]["path"]

try:
# There is a high chance here that `path` is not a path but
# a unique session id
path = model["path"]
except KeyError as e:
raise web.HTTPError(400, "Missing field in JSON data: path") from e
Expand Down
87 changes: 71 additions & 16 deletions jupyter_server/services/sessions/sessionmanager.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,14 @@
"""A base class session manager."""

# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
import os
import pathlib
import uuid
from typing import Any, Dict, List, NewType, Optional, Union

KernelName = NewType("KernelName", str)
ModelName = NewType("ModelName", str)

try:
import sqlite3
Expand All @@ -12,7 +17,6 @@
from pysqlite2 import dbapi2 as sqlite3 # type:ignore[no-redef]

from dataclasses import dataclass, fields
from typing import Union

from jupyter_core.utils import ensure_async
from tornado import web
Expand All @@ -39,8 +43,8 @@ class KernelSessionRecord:
associated with them.
"""

session_id: Union[None, str] = None
kernel_id: Union[None, str] = None
session_id: Optional[str] = None
kernel_id: Optional[str] = None

def __eq__(self, other: object) -> bool:
"""Whether a record equals another."""
Expand Down Expand Up @@ -98,7 +102,9 @@ class KernelSessionRecordList:
it will be appended.
"""

def __init__(self, *records):
_records: List[KernelSessionRecord]

def __init__(self, *records: KernelSessionRecord):
"""Initialize a record list."""
self._records = []
for record in records:
Expand Down Expand Up @@ -252,14 +258,26 @@ async def session_exists(self, path):
exists = True
return exists

def new_session_id(self):
def new_session_id(self) -> str:
"""Create a uuid for a new session"""
return str(uuid.uuid4())

async def create_session(
self, path=None, name=None, type=None, kernel_name=None, kernel_id=None
):
"""Creates a session and returns its model"""
self,
path: Optional[str] = None,
name: Optional[ModelName] = None,
type: Optional[str] = None,
kernel_name: Optional[KernelName] = None,
kernel_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Creates a session and returns its model
name: ModelName(str)
Usually the model name, like the filename associated with current
kernel.
"""
session_id = self.new_session_id()
record = KernelSessionRecord(session_id=session_id)
self._pending_sessions.update(record)
Expand All @@ -277,15 +295,52 @@ async def create_session(
self._pending_sessions.remove(record)
return result

def get_kernel_env(self, path):
"""Return the environment variables that need to be set in the kernel"""
def get_kernel_env(
self, path: Optional[str], name: Optional[ModelName] = None
) -> Dict[str, str]:
"""Return the environment variables that need to be set in the kernel
path : str
the url path for the given session.
name: ModelName(str), optional
Here the name is likely to be the name of the associated file
with the current kernel at startup time.
"""
if name is not None:
cwd = self.kernel_manager.cwd_for_path(path)
path = os.path.join(cwd, name)
assert isinstance(path, str)
return {**os.environ, "JPY_SESSION_NAME": path}

async def start_kernel_for_session(self, session_id, path, name, type, kernel_name):
"""Start a new kernel for a given session."""
async def start_kernel_for_session(
self,
session_id: str,
path: Optional[str],
name: Optional[ModelName],
type: Optional[str],
kernel_name: Optional[KernelName],
) -> str:
"""Start a new kernel for a given session.
session_id : str
uuid for the session; this method must be given a session_id
path : str
the path for the given session - seem to be a session id sometime.
name : str
Usually the model name, like the filename associated with current
kernel.
type : str
the type of the session
kernel_name : str
the name of the kernel specification to use. The default kernel name will be used if not provided.
"""
# allow contents manager to specify kernels cwd
kernel_path = await ensure_async(self.contents_manager.get_kernel_path(path=path))
kernel_env = self.get_kernel_env(path)

kernel_env = self.get_kernel_env(path, name)
kernel_id = await self.kernel_manager.start_kernel(
path=kernel_path,
kernel_name=kernel_name,
Expand All @@ -306,9 +361,9 @@ async def save_session(self, session_id, path=None, name=None, type=None, kernel
uuid for the session; this method must be given a session_id
path : str
the path for the given session
name: str
name : str
the name of the session
type: string
type : str
the type of the session
kernel_id : str
a uuid for the kernel associated with this session
Expand Down Expand Up @@ -405,7 +460,7 @@ async def update_session(self, session_id, **kwargs):
query = "UPDATE session SET %s WHERE session_id=?" % (", ".join(sets))
self.cursor.execute(query, list(kwargs.values()) + [session_id])

def kernel_culled(self, kernel_id):
async def kernel_culled(self, kernel_id: str) -> bool:
"""Checks if the kernel is still considered alive and returns true if its not found."""
return kernel_id not in self.kernel_manager

Expand Down
17 changes: 10 additions & 7 deletions jupyter_server/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import sys
import warnings
from contextlib import contextmanager
from typing import NewType
from urllib.parse import (
SplitResult,
quote,
Expand All @@ -17,14 +18,16 @@
urlsplit,
urlunsplit,
)
from urllib.request import pathname2url # noqa
from urllib.request import pathname2url # noqa: F401

from _frozen_importlib_external import _NamespacePath
from jupyter_core.utils import ensure_async
from packaging.version import Version
from tornado.httpclient import AsyncHTTPClient, HTTPClient, HTTPRequest
from tornado.netutil import Resolver

ApiPath = NewType("ApiPath", str)


def url_path_join(*pieces):
"""Join components of url into a relative url
Expand Down Expand Up @@ -109,19 +112,19 @@ def samefile_simple(path, other_path):
return path.lower() == other_path.lower() and path_stat == other_path_stat


def to_os_path(path, root=""):
def to_os_path(path: ApiPath, root: str = "") -> str:
"""Convert an API path to a filesystem path
If given, root will be prepended to the path.
root must be a filesystem path already.
"""
parts = path.strip("/").split("/")
parts = str(path).strip("/").split("/")
parts = [p for p in parts if p != ""] # remove duplicate splits
path = os.path.join(root, *parts)
return os.path.normpath(path)
path_ = os.path.join(root, *parts)
return os.path.normpath(path_)


def to_api_path(os_path, root=""):
def to_api_path(os_path: str, root: str = "") -> ApiPath:
"""Convert a filesystem path to an API path
If given, root will be removed from the path.
Expand All @@ -132,7 +135,7 @@ def to_api_path(os_path, root=""):
parts = os_path.strip(os.path.sep).split(os.path.sep)
parts = [p for p in parts if p != ""] # remove duplicate splits
path = "/".join(parts)
return path
return ApiPath(path)


def check_version(v, check):
Expand Down
Loading

0 comments on commit a66e8df

Please sign in to comment.