feat: configurable template

Author: mxab

example/jupyterhub/job.hcl.j2

@@ -0,0 +1,93 @@
+job "{{ job_name }}" {
+
+    type = "service"
+    datacenters = {{ datacenters | tojson}}
+
+    meta {
+        jupyterhub_user = "{{ username }}"
+        {% if notebook_name %}
+        jupyterhub_notebook = "{{ notebook_name }}"
+        {% endif %}
+    }
+    group "nb" {
+
+        {% if policies %}
+        vault {
+            policies = {{ policies  | tojson }}
+        }
+        {% endif %}
+        {% if volume_data %}
+        volume "{{ volume_data.volume_name }}" {
+            type      = "{{ volume_data.type }}"
+            read_only = false
+            source    = "{{ volume_data.source }}"
+
+            {% if volume_data.type == "csi" %}
+            attachment_mode = "file-system"
+            access_mode     = "single-node-writer"
+            {% endif %}
+        }
+        {% endif %}
+
+        network {
+            port "notebook" {
+                to = 8888
+            }
+        }
+
+        task "nb" {
+            driver = "docker"
+
+            config {
+                image = "{{ image }}"
+                ports = [ "notebook" ]
+
+                args = {{ args | tojson }}
+            }
+            env {
+                {% for key, value in env|dictsort %}
+                {{ key }} = {{ value | tojson}}
+                {% endfor %}
+                JUPYTER_ENABLE_LAB="yes"
+                # GRANT_SUDO="yes"
+            }
+
+            resources {
+                cpu    = {{ cpu }}
+                memory = {{ memory }}
+            }
+
+            {% if volume_data %}
+            volume_mount {
+                volume      = "{{ volume_data.volume_name }}"
+                destination = "{{ volume_data.destination }}"
+                read_only   = false
+            }
+            {% endif %}
+
+            template {
+                data = <<-EOH
+                [[ with nomadVar "openai" ]]
+                OPENAI_AI_KEY=[[ .openai_ai_key ]]
+                [[ end ]]
+                EOH
+                destination = "secrets/openai.env"
+                env         = true
+                left_delimiter = "[["
+                right_delimiter = "]]"
+            }
+        }
+
+        service {
+            name = "{{ service_name }}"
+            provider = "{{ service_provider }}"
+            port = "notebook"
+             check {
+                name     = "alive"
+                type     = "tcp"
+                interval = "10s"
+                timeout  = "2s"
+            }
+        }
+    }
+}

example/jupyterhub/jupyterhub.nomad

@@ -40,46 +40,13 @@ NOMAD_ADDR=http://host.docker.internal:4646
             }
             template {
                 destination = "/local/jupyterhub_config.py"
-
-                data = <<EOF
-import json
-import os
-import socket
-
-from jupyterhub.auth import DummyAuthenticator
-import tarfile
-c.JupyterHub.spawner_class = "nomad-spawner"
-c.JupyterHub.bind_url = "http://0.0.0.0:8000"
-c.JupyterHub.hub_bind_url = "http://0.0.0.0:8081"
-
-c.JupyterHub.hub_connect_url = f"http://{os.environ.get('NOMAD_IP_api')}:{os.environ.get('NOMAD_HOST_PORT_api')}"
-c.JupyterHub.log_level = "DEBUG"
-c.ConfigurableHTTPProxy.debug = True
-
-
-c.JupyterHub.allow_named_servers = True
-c.JupyterHub.named_server_limit_per_user = 5
-
-c.JupyterHub.authenticator_class = DummyAuthenticator
-
-c.NomadSpawner.datacenters = ["dc1", "dc2", "dc3"]
-c.NomadSpawner.csi_plugin_ids = ["nfs", "hostpath-plugin0"]
-c.NomadSpawner.mem_limit = "2G"
-c.NomadSpawner.common_images = ["jupyter/minimal-notebook:2023-06-26"]
-
-def csi_volume_parameters(spawner):
-    if spawner.user_options["volume_csi_plugin_id"] == "nfs":
-        return {
-            "gid" : "1000",
-            "uid" : "1000"
-        }
-    else:
-        return None
-c.NomadSpawner.csi_volume_parameters = csi_volume_parameters
-
-                EOF
-
-
+                data = file("jupyterhub_config.py")
+            }
+            template {
+                destination = "/local/job.hcl.j2"
+                data = file("job.hcl.j2")
+                left_delimiter = "<[{"
+                right_delimiter = "}]>"
             }
 
             resources {

example/jupyterhub/jupyterhub_config.py

@@ -0,0 +1,37 @@
+# flake8: noqa
+import os
+
+from jupyterhub.auth import DummyAuthenticator
+
+c.JupyterHub.spawner_class = "nomad-spawner"
+c.JupyterHub.bind_url = "http://0.0.0.0:8000"
+c.JupyterHub.hub_bind_url = "http://0.0.0.0:8081"
+
+c.JupyterHub.hub_connect_url = (
+    f"http://{os.environ.get('NOMAD_IP_api')}:{os.environ.get('NOMAD_HOST_PORT_api')}"
+)
+c.JupyterHub.log_level = "DEBUG"
+c.ConfigurableHTTPProxy.debug = True
+
+
+c.JupyterHub.allow_named_servers = True
+c.JupyterHub.named_server_limit_per_user = 5
+
+c.JupyterHub.authenticator_class = DummyAuthenticator
+
+c.NomadSpawner.datacenters = ["dc1", "dc2", "dc3"]
+c.NomadSpawner.csi_plugin_ids = ["nfs", "hostpath-plugin0"]
+c.NomadSpawner.mem_limit = "2G"
+c.NomadSpawner.common_images = ["jupyter/minimal-notebook:2023-06-26"]
+
+
+def csi_volume_parameters(spawner):
+    if spawner.user_options["volume_csi_plugin_id"] == "nfs":
+        return {"gid": "1000", "uid": "1000"}
+    else:
+        return None
+
+
+c.NomadSpawner.csi_volume_parameters = csi_volume_parameters
+
+c.NomadSpawner.job_template_path = "/local/job.hcl.j2"

example/jupyterhub/main.tf

@@ -11,9 +11,22 @@ terraform {
 provider "nomad" {
   address = "http://localhost:4646"
 }
-
+# https://github.com/hashicorp/terraform-provider-nomad/issues/288
+# resource "nomad_variable" "dummy_openai_key" {
+#   path = "openai"
+#   items = {
+#     open_ai_key : "sk_dummy-key"
+#   }
+# }
 resource "nomad_job" "jupyterhub" {
   jobspec = templatefile("jupyterhub.nomad", {
 
   })
+
+  hcl2 {
+    enabled = true
+    allow_fs = true
+
+  }
+
 }

jupyterhub_nomad_spawner/job_factory.py

@@ -1,7 +1,14 @@
+import os
 from enum import Enum
 from typing import Dict, List, Optional
 
-from jinja2 import Environment, PackageLoader, select_autoescape
+from jinja2 import (
+    BaseLoader,
+    Environment,
+    FileSystemLoader,
+    PackageLoader,
+    select_autoescape,
+)
 from pydantic import BaseModel
 
 
@@ -49,12 +56,21 @@ class Config:
         use_enum_values = True
 
 
-def create_job(job_data: JobData) -> str:
-    env = Environment(
-        loader=PackageLoader("jupyterhub_nomad_spawner"), autoescape=select_autoescape()
-    )
+def create_job(job_data: JobData, job_template_path: Optional[str] = None) -> str:
+    # if present, split up and use head as filesystem loader and tail as template name
+    loader: BaseLoader
 
-    template = env.get_template("job.hcl.j2")
+    if job_template_path:
+        head_tail = os.path.split(job_template_path)
+        loader = FileSystemLoader(head_tail[0])
+        template_name = head_tail[1]
+
+    else:
+        loader = PackageLoader("jupyterhub_nomad_spawner")
+        template_name = "job.hcl.j2"
+    env = Environment(loader=loader, autoescape=select_autoescape())
+
+    template = env.get_template(template_name)
     job_hcl = template.render(**job_data.dict())
 
     return job_hcl

jupyterhub_nomad_spawner/spawner.py

@@ -12,9 +12,10 @@
 from traitlets import Bool
 from traitlets import Callable as TCallable
 from traitlets import Dict as TDict
+from traitlets import Enum as TEnum
 from traitlets import Int as TInt
 from traitlets import List as TList
-from traitlets import Unicode, Enum as TEnum
+from traitlets import Unicode
 from traitlets import Union as TUnion
 from traitlets import default
 
@@ -334,6 +335,13 @@ def service_name(self) -> str:
             return f"{self.job_name}"
         raise ValueError("notebook_id is not set")
 
+    job_template_path = Unicode(
+        help="""
+        The path to the job template file with jinja2 placeholders
+        see templates/job.hcl.j2
+    """
+    ).tag(config=True)
+
     async def start(self):
         nomad_service_config = build_nomad_config_from_options(self)
         nomad_httpx_client = build_nomad_httpx_client(nomad_service_config)
@@ -363,8 +371,9 @@ async def start(self):
             elif self.vault_policies:
                 policies = self.vault_policies
 
+            self.log.info("scheduling job %s", self.job_name)
             job_hcl = create_job(
-                JobData(
+                job_data=JobData(
                     job_name=self.job_name,
                     username=self.user.name,
                     notebook_name=self.name,
@@ -377,7 +386,8 @@ async def start(self):
                     memory=self.user_options["memory"],
                     volume_data=volume_data,
                     policies=policies,
-                )
+                ),
+                job_template_path=self.job_template_path,
             )
 
             await nomad_service.schedule_job(job_hcl)
@@ -470,9 +480,6 @@ async def address_and_port_from_nomad(
         self.log.info("Getting service %s from nomad", self.service_name)
         return await nomad_service.get_service_address(self.job_name)
 
-        
-        
-
     async def poll(self):
         nomad_httpx_client = build_nomad_httpx_client(
             build_nomad_config_from_options(self)