Skip to content

Commit 0f3a9ba

Browse files
author
Zachary Seguin
committed
feat(notebooks): Add volume and GPU support on Notebook creation
1 parent d72939d commit 0f3a9ba

File tree

1 file changed

+74
-37
lines changed

1 file changed

+74
-37
lines changed

‎notebooks.go

Lines changed: 74 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package main
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"io/ioutil"
@@ -30,28 +31,35 @@ const (
3031
)
3132

3233
type volumerequest struct {
33-
Type volumetype `json:"type"`
34-
Name string `json:"name"`
35-
TemplatedName string `json:"templatedName"`
36-
Class string `json:"class"`
37-
ExtraFields map[string]interface{} `json:"extraFields"`
38-
Path string `json:"path"`
34+
Type volumetype `json:"type"`
35+
Name string `json:"name"`
36+
TemplatedName string `json:"templatedName"`
37+
Class string `json:"class"`
38+
ExtraFields map[string]interface{} `json:"extraFields"`
39+
Path string `json:"path"`
40+
Size resource.Quantity `json:"size"`
41+
Mode corev1.PersistentVolumeAccessMode `json:"mode"`
42+
}
43+
44+
type gpurequest struct {
45+
Quantity resource.Quantity `json:"num"`
46+
Vendor string `json:"vendor"`
3947
}
4048

4149
type newnotebookrequest struct {
42-
Name string `json:"name"`
43-
Namespace string `json:"namespace"`
44-
Image string `json:"image"`
45-
CustomImage string `json:"customImage"`
46-
CustomImageCheck bool `json:"customImageCheck"`
47-
CPU resource.Quantity `json:"cpu"`
48-
Memory resource.Quantity `json:"memory"`
49-
// TODO: GPU
50-
NoWorkspace bool `json:"noWorkspace"`
51-
Workspace volumerequest `json:"workspace"`
52-
DataVolumes []volumerequest `json:"datavols"`
53-
EnableSharedMemory bool `json:"shm"`
54-
Configurations []string `json:"configurations"`
50+
Name string `json:"name"`
51+
Namespace string `json:"namespace"`
52+
Image string `json:"image"`
53+
CustomImage string `json:"customImage"`
54+
CustomImageCheck bool `json:"customImageCheck"`
55+
CPU resource.Quantity `json:"cpu"`
56+
Memory resource.Quantity `json:"memory"`
57+
GPUs gpurequest `json:"gpus"`
58+
NoWorkspace bool `json:"noWorkspace"`
59+
Workspace volumerequest `json:"workspace"`
60+
DataVolumes []volumerequest `json:"datavols"`
61+
EnableSharedMemory bool `json:"shm"`
62+
Configurations []string `json:"configurations"`
5563
}
5664

5765
type notebookresponse struct {
@@ -226,27 +234,51 @@ func (s *server) GetNotebooks(w http.ResponseWriter, r *http.Request) {
226234
s.respond(w, r, resp)
227235
}
228236

229-
func (s *server) handleVolume(req volumerequest, notebook *notebooksv1.Notebook) error {
230-
if req.Type == VolumeTypeExisting {
231-
notebook.Spec.Template.Spec.Volumes = append(notebook.Spec.Template.Spec.Volumes, corev1.Volume{
232-
Name: req.Name,
233-
VolumeSource: corev1.VolumeSource{
234-
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
235-
ClaimName: req.Name,
237+
func (s *server) handleVolume(ctx context.Context, req volumerequest, notebook *notebooksv1.Notebook) error {
238+
if req.Type == VolumeTypeNew {
239+
// Create the PVC
240+
pvc := corev1.PersistentVolumeClaim{
241+
ObjectMeta: v1.ObjectMeta{
242+
Name: req.Name,
243+
Namespace: notebook.Namespace,
244+
},
245+
Spec: corev1.PersistentVolumeClaimSpec{
246+
AccessModes: []corev1.PersistentVolumeAccessMode{req.Mode},
247+
Resources: corev1.ResourceRequirements{
248+
Requests: corev1.ResourceList{
249+
corev1.ResourceStorage: req.Size,
250+
},
236251
},
237252
},
238-
})
253+
}
239254

240-
notebook.Spec.Template.Spec.Containers[0].VolumeMounts = append(notebook.Spec.Template.Spec.Containers[0].VolumeMounts, corev1.VolumeMount{
241-
Name: req.Name,
242-
MountPath: req.Path,
243-
})
244-
} else if req.Type == VolumeTypeNew {
245-
return fmt.Errorf("unsupported volume type %q", req.Type)
246-
} else {
255+
// Add the storage class, if set and not set to an "empty" value
256+
if req.Class != "" && req.Class != "{none}" && req.Class != "{empty}" {
257+
pvc.Spec.StorageClassName = &req.Class
258+
}
259+
260+
if _, err := s.clientsets.kubernetes.CoreV1().PersistentVolumeClaims(notebook.Namespace).Create(ctx, &pvc, v1.CreateOptions{}); err != nil {
261+
return err
262+
}
263+
} else if req.Type != VolumeTypeExisting {
247264
return fmt.Errorf("unknown volume type %q", req.Type)
248265
}
249266

267+
// Add the volume and volume mount ot the notebook spec
268+
notebook.Spec.Template.Spec.Volumes = append(notebook.Spec.Template.Spec.Volumes, corev1.Volume{
269+
Name: req.Name,
270+
VolumeSource: corev1.VolumeSource{
271+
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
272+
ClaimName: req.Name,
273+
},
274+
},
275+
})
276+
277+
notebook.Spec.Template.Spec.Containers[0].VolumeMounts = append(notebook.Spec.Template.Spec.Containers[0].VolumeMounts, corev1.VolumeMount{
278+
Name: req.Name,
279+
MountPath: req.Path,
280+
})
281+
250282
return nil
251283
}
252284

@@ -276,7 +308,6 @@ func (s *server) NewNotebook(w http.ResponseWriter, r *http.Request) {
276308

277309
// Setup the notebook
278310
// TODO: Work with default CPU/memory limits from config
279-
// TODO: Add GPU support
280311
notebook := notebooksv1.Notebook{
281312
ObjectMeta: v1.ObjectMeta{
282313
Name: req.Name,
@@ -310,15 +341,15 @@ func (s *server) NewNotebook(w http.ResponseWriter, r *http.Request) {
310341
// Add workspace volume
311342
if !req.NoWorkspace {
312343
req.Workspace.Path = WorkspacePath
313-
err = s.handleVolume(req.Workspace, &notebook)
344+
err = s.handleVolume(r.Context(), req.Workspace, &notebook)
314345
if err != nil {
315346
s.error(w, r, err)
316347
return
317348
}
318349
}
319350

320351
for _, volreq := range req.DataVolumes {
321-
err = s.handleVolume(volreq, &notebook)
352+
err = s.handleVolume(r.Context(), volreq, &notebook)
322353
if err != nil {
323354
s.error(w, r, err)
324355
return
@@ -342,6 +373,12 @@ func (s *server) NewNotebook(w http.ResponseWriter, r *http.Request) {
342373
})
343374
}
344375

376+
// Add GPU
377+
if !req.GPUs.Quantity.IsZero() {
378+
notebook.Spec.Template.Spec.Containers[0].Resources.Requests[corev1.ResourceName(req.GPUs.Vendor)] = req.GPUs.Quantity
379+
notebook.Spec.Template.Spec.Containers[0].Resources.Limits[corev1.ResourceName(req.GPUs.Vendor)] = req.GPUs.Quantity
380+
}
381+
345382
log.Printf("creating notebook %q for %q", notebook.ObjectMeta.Name, namespace)
346383

347384
// Submit the notebook to the API server

0 commit comments

Comments
 (0)