Skip to content

Commit 8956795

Browse files
pbaillypbailly
authored
[9.0.1] Improve Jupyter Notebook REST API [ch61633] (#124)
* Use the new REST API for Jupyter Notebook State. [ch61633] * Apply review. [ch61633] * Prevent empty save. * Rework API [ch61633] * Add missing file. [ch61633] * Separe legacy API from the new API. * Improve code consistency. * Improve session handling. * Tweak after code review. * Remove warnings.
1 parent 9de20c6 commit 8956795

File tree

4 files changed

+211
-79
lines changed

4 files changed

+211
-79
lines changed

‎dataikuapi/dss/jupyternootebook.py‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
from .discussion import DSSObjectDiscussions
2+
3+
class DSSJupyterNotebook(object):
4+
def __init__(self, client, project_key, notebook_name):
5+
self.client = client
6+
self.project_key = project_key
7+
self.notebook_name = notebook_name
8+
9+
def unload(self, session_id=None):
10+
"""
11+
Stop this Jupyter notebook and release its resources
12+
"""
13+
sessions = self.get_sessions()
14+
if sessions is None:
15+
raise Exception("Notebook isn't running")
16+
if len(sessions) == 0:
17+
raise Exception("Notebook isn't running")
18+
if session_id is None:
19+
if len(sessions) > 1:
20+
raise Exception("Several sessions of the notebook are running, choose one")
21+
else:
22+
session_id = sessions[0].get('sessionId', None)
23+
return self.client._perform_json("DELETE",
24+
"/projects/%s/jupyter-notebooks/%s/sessions/%s" % (self.project_key, self.notebook_name, session_id))
25+
26+
def get_sessions(self, as_objects=False):
27+
"""
28+
Get the list of running sessions of this Jupyter notebook
29+
30+
:param boolean as_objects: if True, each returned item will be a :class:`dataikuapi.dss.notebook.DSSNotebookSession`
31+
:rtype: list of :class:`dataikuapi.dss.notebook.DSSNotebookSession` or list of dict
32+
"""
33+
sessions = self.client._perform_json("GET",
34+
"/projects/%s/jupyter-notebooks/%s/sessions" % (self.project_key, self.notebook_name))
35+
36+
if as_objects:
37+
return [DSSNotebookSession(self.client, session) for session in sessions]
38+
else:
39+
return sessions
40+
41+
def get_content(self):
42+
"""
43+
Get the content of this Jupyter notebook (metadata, cells, nbformat)
44+
"""
45+
raw_content = self.client._perform_json("GET", "/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name))
46+
return DSSNotebookContent(self.client, self.project_key, self.notebook_name, raw_content)
47+
48+
def delete(self):
49+
"""
50+
Delete this Jupyter notebook and stop all of its active sessions.
51+
"""
52+
return self.client._perform_json("DELETE",
53+
"/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name))
54+
55+
########################################################
56+
# Discussions
57+
########################################################
58+
def get_object_discussions(self):
59+
"""
60+
Get a handle to manage discussions on the notebook
61+
62+
:returns: the handle to manage discussions
63+
:rtype: :class:`dataikuapi.discussion.DSSObjectDiscussions`
64+
"""
65+
return DSSObjectDiscussions(self.client, self.project_key, "JUPYTER_NOTEBOOK", self.notebook_name)
66+
67+
class DSSNotebookContent(object):
68+
"""
69+
Content of a Jupyter Notebook. Do not create this directly, use :meth:`DSSJupyterNotebook.get_content`
70+
"""
71+
72+
"""
73+
A Python/R/Scala notebook on the DSS instance
74+
"""
75+
def __init__(self, client, project_key, notebook_name, content):
76+
self.client = client
77+
self.project_key = project_key
78+
self.notebook_name = notebook_name
79+
self.content = content
80+
81+
def get_raw(self):
82+
"""
83+
Get the content of this Jupyter notebook (metadata, cells, nbformat)
84+
:rtype: a dict containing the full content of a notebook
85+
"""
86+
return self.content
87+
88+
def get_metadata(self):
89+
"""
90+
Get the metadata associated to this Jupyter notebook
91+
:rtype: dict with metadata
92+
"""
93+
return self.content["metadata"]
94+
95+
def get_cells(self):
96+
"""
97+
Get the cells associated to this Jupyter notebook
98+
:rtype: list of cells
99+
"""
100+
return self.content["cells"]
101+
102+
def save(self):
103+
"""
104+
Save the content of this Jupyter notebook
105+
"""
106+
return self.client._perform_json("PUT",
107+
"/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name),
108+
body=self.content)
109+
110+
class DSSNotebookSession(object):
111+
"""
112+
Metadata associated to the session of a Jupyter Notebook. Do not create this directly, use :meth:`DSSJupyterNotebook.get_sessions()`
113+
"""
114+
115+
def __init__(self, client, session):
116+
self.client = client
117+
self.project_key = session.get("projectKey")
118+
self.notebook_name = session.get("notebookName")
119+
self.session_creator = session.get("sessionCreator")
120+
self.session_creator_display_name = session.get("sessionCreatorDisplayName")
121+
self.session_unix_owner = session.get("sessionUnixOwner")
122+
self.session_id = session.get("sessionId")
123+
self.kernel_id = session.get("kernelId")
124+
self.kernel_pid = session.get("kernelPid")
125+
self.kernel_connections = session.get("kernelConnections")
126+
self.kernel_last_activity_time = session.get("kernelLastActivityTime")
127+
self.kernel_execution_state = session.get("kernelExecutionState")
128+
self.session_start_time = session.get("sessionStartTime")
129+
130+
131+
def unload(self):
132+
"""
133+
Stop this Jupyter notebook and release its resources
134+
"""
135+
return self.client._perform_json("DELETE",
136+
"/projects/%s/jupyter-notebooks/%s/sessions/%s" % (self.project_key, self.notebook_name, self.session_id))

‎dataikuapi/dss/notebook.py‎

Lines changed: 30 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -2,90 +2,63 @@
22

33
class DSSNotebook(object):
44
"""
5+
Deprecated. Use DSSJupyterNotebook
56
A Python/R/Scala notebook on the DSS instance
67
"""
7-
def __init__(self, client, project_key, notebook_name, state=None, content=None):
8-
self.client = client
9-
self.project_key = project_key
10-
self.notebook_name = notebook_name
11-
self.state = state
12-
self.content = content
13-
self.state_is_peek = True
8+
def __init__(self, client, project_key, notebook_name, state=None):
9+
self.client = client
10+
self.project_key = project_key
11+
self.notebook_name = notebook_name
12+
self.state = state
13+
self.state_is_peek = True
1414

1515
def unload(self, session_id=None):
1616
"""
17+
Deprecated. Use DSSJupyterNotebook
1718
Stop this Jupyter notebook and release its resources
1819
"""
19-
sessions = self.get_sessions()
20-
if sessions is None:
20+
state = self.get_state()
21+
if state is None:
2122
raise Exception("Notebook isn't running")
22-
if len(sessions) == 0:
23+
if state.get('activeSessions', None) is None:
24+
raise Exception("Notebook isn't running")
25+
if len(state['activeSessions']) == 0:
2326
raise Exception("Notebook isn't running")
2427
if session_id is None:
25-
if len(sessions) > 1:
28+
if len(state['activeSessions']) > 1:
2629
raise Exception("Several sessions of the notebook are running, choose one")
2730
else:
28-
session_id = sessions[0].get('sessionId', None)
29-
return self.client._perform_json("DELETE",
30-
"/projects/%s/jupyter-notebooks/%s/sessions/%s" % (self.project_key, self.notebook_name, session_id))
31+
session_id = state['activeSessions'][0].get('sessionId', None)
32+
return self.client._perform_json("DELETE", "/projects/%s/notebooks/" % self.project_key, params={'notebookName' : self.notebook_name, 'sessionId' : session_id})
3133

32-
def get_state(self, refresh=False):
33-
"""
34-
Get the status of this Jupyter notebook
3534

36-
:param bool refresh: if True, get the status of the notebook from the backend
35+
def get_state(self):
36+
"""
37+
Deprecated. Use DSSJupyterNotebook
38+
Get the metadata associated to this Jupyter notebook
3739
"""
38-
notebook_states = self.client._perform_json("GET",
39-
"/projects/%s/jupyter-notebooks/" % self.project_key,
40-
params={"active": False})
41-
if self.state is None or refresh:
42-
for notebook_state in notebook_states:
43-
if notebook_state.get("name") == self.notebook_name:
44-
self.state = notebook_state
45-
break
40+
if self.state is None:
41+
self.state = self.client._perform_json("GET", "/projects/%s/notebooks/" % self.project_key, params={'notebookName' : self.notebook_name})
4642
return self.state
4743

4844
def get_sessions(self):
4945
"""
46+
Deprecated. Use DSSJupyterNotebook
5047
Get the list of running sessions of this Jupyter notebook
5148
"""
52-
53-
if self.state is None:
54-
self.state = {}
55-
sessions = self.client._perform_json("GET",
56-
"/projects/%s/jupyter-notebooks/%s/sessions" % (self.project_key, self.notebook_name))
57-
self.state["activeSessions"] = sessions
58-
return sessions
59-
60-
def get_content(self):
61-
"""
62-
Get the content of this Jupyter notebook (metadata, cells, nbformat)
63-
"""
64-
if self.content is None:
65-
self.content = self.client._perform_json("GET",
66-
"/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name))
67-
return self.content
68-
69-
def save(self):
70-
"""
71-
Save the content of this Jupyter notebook
72-
"""
73-
return self.client._perform_json("PUT",
74-
"/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name),
75-
body=self.content)
76-
77-
def delete(self):
78-
"""
79-
Delete this Jupyter notebook and stop all of its active sessions.
80-
"""
81-
return self.client._perform_json("DELETE",
82-
"/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name))
49+
state = self.get_state()
50+
if state is None:
51+
raise Exception("Notebook isn't running")
52+
if state.get('activeSessions', None) is None:
53+
raise Exception("Notebook isn't running")
54+
return state['activeSessions']
8355

8456
########################################################
8557
# Discussions
8658
########################################################
8759
def get_object_discussions(self):
8860
"""
61+
Deprecated. Use DSSJupyterNotebook
8962
Get a handle to manage discussions on the notebook
9063
9164
:returns: the handle to manage discussions

‎dataikuapi/dss/project.py‎

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import time, warnings, sys, os.path as osp
22
from .dataset import DSSDataset, DSSDatasetListItem, DSSManagedDatasetCreationHelper
3+
from .jupyternootebook import DSSJupyterNotebook, DSSNotebookContent
4+
from .notebook import DSSNotebook
35
from .streaming_endpoint import DSSStreamingEndpoint, DSSStreamingEndpointListItem, DSSManagedStreamingEndpointCreationHelper
46
from .recipe import DSSRecipeListItem, DSSRecipe
57
from . import recipe
@@ -11,15 +13,13 @@
1113
from .continuousactivity import DSSContinuousActivity
1214
from .apiservice import DSSAPIService
1315
from .future import DSSFuture
14-
from .notebook import DSSNotebook
1516
from .macro import DSSMacro
1617
from .wiki import DSSWiki
1718
from .discussion import DSSObjectDiscussions
1819
from .ml import DSSMLTask
1920
from .analysis import DSSAnalysis
2021
from .flow import DSSProjectFlow
2122
from .app import DSSAppManifest
22-
from ..utils import DataikuException
2323

2424

2525
class DSSProject(object):
@@ -843,12 +843,11 @@ def list_jupyter_notebooks(self, as_objects=True, active=False):
843843
:returns: The list of the notebooks - see as_objects for more information
844844
:rtype: list
845845
"""
846-
notebooks = self.client._perform_json("GET", "/projects/%s/jupyter-notebooks/" % self.project_key,
847-
params={"active": active})
846+
notebook_names = self.client._perform_json("GET", "/projects/%s/jupyter-notebooks/" % self.project_key, params={"active": active})
848847
if as_objects:
849-
return [DSSNotebook(self.client, notebook_state['projectKey'], notebook_state['name'], state=notebook_state) for notebook_state in notebooks]
848+
return [DSSJupyterNotebook(self.client, self.project_key, notebook_name) for notebook_name in notebook_names]
850849
else:
851-
return notebooks
850+
return notebook_names
852851

853852
def get_jupyter_notebook(self, notebook_name):
854853
"""
@@ -860,9 +859,7 @@ def get_jupyter_notebook(self, notebook_name):
860859
:returns: A handle to interact with this jupyter notebook
861860
:rtype: :class:`~dataikuapi.dss.notebook.DSSNotebook` jupyter notebook handle
862861
"""
863-
notebook_content = self.client._perform_json("GET",
864-
"/projects/%s/jupyter-notebooks/%s" % (self.project_key, notebook_name))
865-
return DSSNotebook(self.client, self.project_key, notebook_name, content=notebook_content)
862+
return DSSJupyterNotebook(self.client, self.project_key, notebook_name)
866863

867864
def create_jupyter_notebook(self, notebook_name, notebook_content):
868865
"""
@@ -877,10 +874,8 @@ def create_jupyter_notebook(self, notebook_name, notebook_content):
877874
:returns: A handle to interact with the newly created jupyter notebook
878875
:rtype: :class:`~dataikuapi.dss.notebook.DSSNotebook` jupyter notebook handle
879876
"""
880-
created_notebook_content = self.client._perform_json("POST",
881-
"/projects/%s/jupyter-notebooks/%s" % (self.project_key, notebook_name),
882-
body=notebook_content)
883-
return DSSNotebook(self.client, self.project_key, notebook_name, content=created_notebook_content)
877+
self.client._perform_json("POST", "/projects/%s/jupyter-notebooks/%s" % (self.project_key, notebook_name), body=notebook_content)
878+
return self.get_jupyter_notebook(notebook_name)
884879

885880
########################################################
886881
# Continuous activities
@@ -1279,19 +1274,20 @@ def sync_datasets_acls(self):
12791274
########################################################
12801275
# Notebooks
12811276
########################################################
1282-
1277+
12831278
def list_running_notebooks(self, as_objects=True):
12841279
"""
1280+
Deprecated. Use :meth:`DSSProject.list_jupyter_notebooks`
12851281
List the currently-running notebooks
12861282
12871283
Returns:
12881284
list of notebooks. Each object contains at least a 'name' field
12891285
"""
1290-
list = self.client._perform_json("GET", "/projects/%s/notebooks/active" % self.project_key)
1286+
notebook_list = self.client._perform_json("GET", "/projects/%s/notebooks/active" % self.project_key)
12911287
if as_objects:
1292-
return [DSSNotebook(self.client, notebook['projectKey'], notebook['name'], notebook) for notebook in list]
1288+
return [DSSNotebook(self.client, notebook['projectKey'], notebook['name'], notebook) for notebook in notebook_list]
12931289
else:
1294-
return list
1290+
return notebook_list
12951291

12961292
########################################################
12971293
# Tags

0 commit comments

Comments
 (0)