Integrate kbatch

qhub/initialize.py

@@ -58,6 +58,9 @@
     "monitoring": {
         "enabled": True,
     },
+    "kbatch": {
+        "enabled": True,
+    },
     "cdsdashboards": {
         "enabled": True,
         "cds_hide_user_named_servers": True,

qhub/schema.py

@@ -74,6 +74,13 @@ class HelmExtension(Base):
     overrides: typing.Optional[typing.Dict]
 
 
+# ============== kbatch =============
+
+
+class KBatch(Base):
+    enabled: bool
+
+
 # ============== Monitoring =============
 
 
@@ -452,6 +459,7 @@ class Main(Base):
     theme: Theme
     profiles: Profiles
     environments: typing.Dict[str, CondaEnvironment]
+    kbatch: typing.Optional[KBatch]
     monitoring: typing.Optional[Monitoring]
     clearml: typing.Optional[ClearML]
     tf_extensions: typing.Optional[typing.List[QHubExtension]]

qhub/stages/input_vars.py

@@ -251,6 +251,7 @@ def stage_07_kubernetes_services(stage_outputs, config):
         "dask-gateway-profiles": config["profiles"]["dask_worker"],
         # monitoring
         "monitoring-enabled": config["monitoring"]["enabled"],
+        "kbatch-enabled": config["kbatch"]["enabled"],
         # prefect
         "prefect-enabled": config.get("prefect", {}).get("enabled", False),
         "prefect-token": config.get("prefect", {}).get("token", ""),

qhub/template/stages/07-kubernetes-services/jupyterhub.tf

@@ -112,7 +112,10 @@ module "jupyterhub" {
 
   services = concat([
     "dask-gateway"
-  ], (var.prefect-enabled ? ["prefect"] : []))
+    ],
+    (var.prefect-enabled ? ["prefect"] : []),
+    (var.kbatch-enabled ? ["kbatch"] : [])
+  )
 
   general-node-group = var.node_groups.general
   user-node-group  = var.node_groups.user

qhub/template/stages/07-kubernetes-services/kbatch.tf

@@ -0,0 +1,23 @@
+# ======================= VARIABLES ======================
+variable "kbatch-enabled" {
+  description = "kbatch enabled or disabled"
+  type = bool
+}
+
+
+# ====================== RESOURCES =======================
+module "kbatch" {
+  count = var.kbatch-enabled ? 1 : 0
+
+  source = "./modules/kubernetes/services/kbatch"
+
+  namespace            = var.environment
+  external-url         = var.endpoint
+
+  jupyterhub_api_token = module.jupyterhub.services.kbatch.api_token
+  node-group   = var.node_groups.user
+
+  dask-gateway-address = module.dask-gateway.config.gateway.address
+  dask-gateway-proxy-address = module.dask-gateway.config.gateway.proxy_address
+  dask-worker-image = var.dask-worker-image
+}

qhub/template/stages/07-kubernetes-services/modules/kubernetes/services/jupyterhub/main.tf

@@ -83,6 +83,8 @@ resource "helm_release" "jupyterhub" {
             name = service
             admin = true
             api_token = random_password.service_token[service].result
+            # come back and clean up
+            url = "http://kbatch-kbatch-proxy.${var.external-url}.svc.cluster.local"
           }
         }
 

qhub/template/stages/07-kubernetes-services/modules/kubernetes/services/kbatch/main.tf

@@ -0,0 +1,81 @@
+# Get this name dynamtically
+locals {
+  kbatch_service_account_name = "kbatch-kbatch-proxy"
+}
+
+resource "helm_release" "kbatch" {
+  name       = "kbatch"
+  namespace  = var.namespace
+  repository = "https://kbatch-dev.github.io/helm-chart"
+  chart      = "kbatch-proxy"
+  version    = "0.3.1"
+
+  values = concat([
+    file("${path.module}/values.yaml"),
+    jsonencode({
+      app = {
+        jupyterhub_api_token = var.jupyterhub_api_token
+        jupyterhub_api_url = "https://${var.external-url}/hub/api/"
+        extra_env = {
+          KBATCH_PREFIX = "/services/kbatch"
+          KBATCH_JOB_EXTRA_ENV = jsonencode({
+            DASK_GATEWAY__AUTH__TYPE = "jupyterhub"
+            DASK_GATEWAY__CLUSTER__OPTIONS__IMAGE = "${var.dask-worker-image.name}:${var.dask-worker-image.tag}"
+            DASK_GATEWAY__ADDRESS = "${var.dask-gateway-address}"            
+            DASK_GATEWAY__PROXY_ADDRESS = "${var.dask-gateway-proxy-address}"
+          })
+        }
+      }
+      image = {
+          tag = "0.3.1"
+      }
+    })
+  ])
+
+  set_sensitive {
+    name  = "jupyterHubToken"
+    value = var.jupyterhub_api_token
+  }
+
+  set {
+    name  = "kbatchImage"
+    value = var.image
+  }
+
+  set {
+    name  = "namespace"
+    value = var.namespace
+  }
+
+}
+
+resource "kubernetes_cluster_role" "kbatch" {
+  metadata {
+    name = "${var.name}-kbatch"
+  }
+
+  rule {
+    api_groups = ["", "batch"]
+    resources  = ["*"]
+    verbs      = ["get", "watch", "list", "patch", "create"]
+  }
+}
+
+
+resource "kubernetes_cluster_role_binding" "kbatch" {
+  metadata {
+    name = "${var.name}-kbatch"
+  }
+
+  role_ref {
+    api_group = "rbac.authorization.k8s.io"
+    kind      = "ClusterRole"
+    name      = kubernetes_cluster_role.kbatch.metadata.0.name
+  }
+  subject {
+    kind      = "ServiceAccount"
+    name      = local.kbatch_service_account_name
+    namespace = var.namespace
+    api_group = ""
+  }
+}

qhub/template/stages/07-kubernetes-services/modules/kubernetes/services/kbatch/values.yaml

@@ -0,0 +1 @@
+# https://github.com/kbatch-dev/helm-chart/blob/main/kbatch/values.yaml

qhub/template/stages/07-kubernetes-services/modules/kubernetes/services/kbatch/variables.tf

@@ -0,0 +1,43 @@
+variable "name" {
+  description = "name prefix to assign to kbatch"
+  type        = string
+  default     = "qhub"
+}
+
+variable "jupyterhub_api_token" {
+  type    = string
+  default = ""
+}
+
+variable "namespace" {
+  type    = string
+  default = "dev"
+}
+
+variable "image" {
+  type    = string
+  default = ""
+}
+
+variable "node-group" {
+  description = "Node key value pair for bound resources"
+  type = object({
+    key   = string
+    value = string
+  })
+}
+
+variable "external-url" {
+  description = "External url that jupyterhub cluster is accessible"
+  type        = string
+}
+
+variable "overrides" {
+  description = "kbatch helm chart list of overrides"
+  type        = list(string)
+  default     = []
+}
+
+variable "dask-gateway-address" {}
+variable "dask-gateway-proxy-address" {}
+variable "dask-worker-image" {}

qhub/template/stages/07-kubernetes-services/modules/kubernetes/services/kbatch/versions.tf

@@ -0,0 +1,13 @@
+terraform {
+  required_providers {
+    helm = {
+      source  = "hashicorp/helm"
+      version = "2.1.2"
+    }
+    kubernetes = {
+      source  = "hashicorp/kubernetes"
+      version = "2.7.1"
+    }
+  }
+  required_version = ">= 1.0"
+}