Heavy notebook

.github/workflows/merge.yml

@@ -16,12 +16,11 @@ jobs:
 
     services:
       mongodb:
-        image: bitnami/mongodb:latest
+        image: mongodb/mongodb-community-server
         env:
-          MONGODB_ROOT_PASSWORD: "admin123456789"
-          MONGODB_USERNAME: "bridge-user"
-          MONGODB_PASSWORD: "admin123456789"
-          MONGODB_DATABASE: "bridge"
+          MONGO_INITDB_ROOT_PASSWORD: "admin123456789"
+          MONGO_INITDB_ROOT_USERNAME: "bridge-user"
+          MONGO_INITDB_DATABASE: "bridge"
         ports:
           - 27017:27017
       keydb:

config/notebook.toml

@@ -17,6 +17,11 @@ notebook_env = [
 secret = "ibmdpdev-openad-pull-secret"
 start_up_url = "lab/tree/start_menu.ipynb"
 max_idle_time = 86400
+[open_ad_workbench.scheduling]
+cpu_heavy = "4"
+mem_heavy = "32Gi"
+toleration_key = "worker/cpu"
+toleration_value = "large"
 [open_ad_workbench.env]
 PROXY_URL = "https://open.accelerate.science/proxy"
 
@@ -36,5 +41,10 @@ notebook_env = [
 	# base-url is taken care by bridge
 ]
 max_idle_time = 86400
+[datascience_notebook.scheduling]
+cpu_heavy = "4"
+mem_heavy = "8Gi"
+toleration_key = "worker/cpu"
+toleration_value = ""
 [datascience_notebook.env]
 PROXY_URL = ""

config/services_sample.toml

@@ -19,6 +19,13 @@
 	check = false
 	mcp = false
 
+	# This is a supplementary entry to notebook for heavier resource
+	[services.notebook_heavy]
+	url = ""
+	readiness = ""
+	check = false
+	mcp = false
+
 	[services.owui]
 	url = ""
 	readiness = ""

justfile

@@ -82,12 +82,6 @@ build-front:
 
 # --- Local Development Services ---
 local-mongo:
-	podman run -d --rm --name mongodb \
-	-e MONGODB_ROOT_PASSWORD="admin123456789" \
-	-e MONGODB_USERNAME="bridge-user" -e MONGODB_PASSWORD="admin123456789" -e MONGODB_DATABASE="bridge" \
-	-p 27017:27017 bitnami/mongodb:latest
-
-local-mongo-arm:
 	podman run -d --rm --name mongodb \
 	-e MONGO_INITDB_ROOT_USERNAME="guardian-user" \
 	-e MONGO_INITDB_ROOT_PASSWORD="admin123456789" \

src/config/mod.rs

@@ -71,6 +71,15 @@ pub struct Notebook {
     pub args: Option<Vec<String>>,
     pub start_up_url: Option<String>,
     pub max_idle_time: Option<u64>,
+    pub scheduling: Option<Scheduling>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct Scheduling {
+    pub cpu_heavy: String,
+    pub mem_heavy: String,
+    pub toleration_key: String,
+    pub toleration_value: String,
 }
 
 const OIDC_PROVIDER: [OpenIDProvider; 2] = [OpenIDProvider::W3, OpenIDProvider::IbmId];

src/kube/mod.rs

@@ -5,7 +5,7 @@ pub use models::*;
 mod notebook;
 #[cfg(feature = "notebook")]
 pub use notebook::{
-    LifecycleStream, Medium, NOTEBOOK_NAMESPACE, Notebook, NotebookSpec, PVCSpec,
+    LifecycleStream, Medium, NOTEBOOK_NAMESPACE, Notebook, NotebookSpec, PVCSpec, Toleration,
     notebook_lifecycle,
 };
 

src/kube/notebook/mod.rs

@@ -22,18 +22,42 @@ pub struct NotebookSpec {
     template: NotebookTemplateSpec,
 }
 
+static CPU_HEAVY_DEFAULT: &str = "4";
+static MEM_HEAVY_DEFAULT: &str = "8Gi";
+
 impl NotebookSpec {
     pub fn new(
         name: String,
         notebook_image_name: &str,
         volume_name: String,
+        tolerations: Option<Vec<Toleration>>,
         notebook_start_url: &mut Option<String>,
         max_idle_time: &mut Option<u64>,
         env_to_add: Vec<(String, String)>,
     ) -> Self {
         let notebook_image = CONFIG.notebooks.get(notebook_image_name).unwrap();
         let mut notebook_env = notebook_image.notebook_env.clone().unwrap_or_default();
 
+        // get resource limit from notebook config
+        let (cpu, mem) = {
+            if tolerations.is_some() {
+                (
+                    notebook_image
+                        .scheduling
+                        .as_ref()
+                        .map(|v| v.cpu_heavy.as_ref())
+                        .unwrap_or(CPU_HEAVY_DEFAULT),
+                    notebook_image
+                        .scheduling
+                        .as_ref()
+                        .map(|v| v.mem_heavy.as_ref())
+                        .unwrap_or(MEM_HEAVY_DEFAULT),
+                )
+            } else {
+                ("2", "4Gi")
+            }
+        };
+
         notebook_env.push(format!(
             "--ServerApp.base_url='notebook/{}/{}'",
             *NOTEBOOK_NAMESPACE, name
@@ -68,12 +92,12 @@ impl NotebookSpec {
                         image: notebook_image.url.clone(),
                         resources: Some(ResourceRequirements {
                             requests: BTreeMap::from([
-                                ("cpu".to_string(), "2".to_string()),
-                                ("memory".to_string(), "4Gi".to_string()),
+                                ("cpu".to_string(), cpu.to_string()),
+                                ("memory".to_string(), mem.to_string()),
                             ]),
                             limits: BTreeMap::from([
-                                ("cpu".to_string(), "2".to_string()),
-                                ("memory".to_string(), "4Gi".to_string()),
+                                ("cpu".to_string(), cpu.to_string()),
+                                ("memory".to_string(), mem.to_string()),
                             ]),
                         }),
                         image_pull_policy: notebook_image.pull_policy.clone(),
@@ -86,6 +110,7 @@ impl NotebookSpec {
                         workingdir: notebook_image.working_dir.clone(),
                         env: Some(env),
                     }],
+                    tolerations,
                     image_pull_secrets: notebook_image
                         .secret
                         .clone()
@@ -114,9 +139,31 @@ pub struct PodSpec {
     containers: Vec<ContainerSpec>,
     #[serde(rename = "imagePullSecrets")]
     image_pull_secrets: Option<Vec<ImagePullSecret>>,
+    tolerations: Option<Vec<Toleration>>,
     volumes: Option<Vec<VolumeSpec>>,
 }
 
+#[derive(Deserialize, Serialize, Clone, Debug, JsonSchema)]
+pub struct Toleration {
+    pub effect: Option<std::string::String>,
+    pub key: Option<std::string::String>,
+    pub operator: Option<std::string::String>,
+    pub toleration_seconds: Option<i64>,
+    pub value: Option<std::string::String>,
+}
+
+impl Toleration {
+    pub fn new(key: String, value: String) -> Self {
+        Self {
+            effect: Some("NoSchedule".to_string()),
+            key: Some(key),
+            operator: Some("Equal".to_string()),
+            toleration_seconds: None,
+            value: Some(value),
+        }
+    }
+}
+
 #[derive(Deserialize, Serialize, Clone, Debug, JsonSchema)]
 pub struct ImagePullSecret {
     name: String,
@@ -228,6 +275,7 @@ mod test {
             name,
             "open_ad_workbench",
             volume_name,
+            None,
             &mut start_url,
             &mut max_idle_time,
             vec![],

src/web/route/notebook/mod.rs

@@ -34,14 +34,16 @@ use crate::{
         mongo::{DB, ObjectID},
     },
     errors::{BridgeError, Result},
-    kube::{KubeAPI, NOTEBOOK_NAMESPACE, Notebook, NotebookSpec, PVCSpec},
+    kube::{KubeAPI, NOTEBOOK_NAMESPACE, Notebook, NotebookSpec, PVCSpec, Toleration},
     web::{
         bridge_middleware::{CookieCheck, Htmx, NotebookCookieCheck},
         helper::{self, bson, forwarding},
     },
 };
 
 pub const NOTEBOOK_SUB_NAME: &str = "notebook";
+pub const NOTEBOOK_SUB_HEAVY_NAME: &str = "notebook_heavy";
+const NOTEBOOK_CFG_NAME: &str = "open_ad_workbench";
 const NOTEBOOK_PORT: &str = "8888";
 const NOTEBOOK_TOKEN_LIFETIME: usize = const { 60 * 60 * 24 * 30 };
 const PVC_DELETE_ATTEMPT: u8 = 9;
@@ -172,6 +174,27 @@ async fn notebook_create(
             ));
         }
 
+        // check for heavy notebook add-on
+        let tolerations = if group
+            .subscriptions
+            .contains(&NOTEBOOK_SUB_HEAVY_NAME.to_string())
+        {
+            let notebook_image = CONFIG
+                .notebooks
+                .get(NOTEBOOK_CFG_NAME)
+                .and_then(|v| v.scheduling.as_ref())
+                .map(|v| (&v.toleration_key, &v.toleration_value));
+
+            if let Some(kv) = notebook_image {
+                let (key, value) = (kv.0.to_string(), kv.1.to_string());
+                Some(vec![Toleration::new(key, value)])
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+
         let scp = if user.groups.is_empty() {
             vec!["".to_string()]
         } else {
@@ -260,8 +283,9 @@ async fn notebook_create(
             &name,
             NotebookSpec::new(
                 name.clone(),
-                "open_ad_workbench",
+                NOTEBOOK_CFG_NAME,
                 pvc_name,
+                tolerations,
                 &mut start_up_url,
                 &mut max_idle_time,
                 vec![("PROXY_KEY".to_string(), notebook_token)],