add custom k3s image for gpu support

.dockerignore

@@ -5,5 +5,6 @@
 !go.mod
 !go.sum
 !sdk
+!docker
 **/__pycache__
 *.pyc

Makefile

@@ -12,6 +12,7 @@ setup-sdk:
 	poetry install -C sdk
 
 k3d-up:
+	docker build . -f ./docker/Dockerfile.k3d -t localhost:5001/rancher/k3s:latest
 	k3d cluster create --config hack/k3d.yaml
 	kubectl config set contexts.k3d-beam.namespace beam
 	okteto context use k3d-beam --namespace beam

docker/Dockerfile.k3d

@@ -0,0 +1,37 @@
+# syntax=docker/dockerfile:1.6
+FROM rancher/k3s:v1.28.5-k3s1 as k3s
+FROM nvidia/cuda:12.3.1-base-ubuntu20.04
+
+ENV CRI_CONFIG_FILE=/var/lib/rancher/k3s/agent/etc/crictl.yaml
+ENV PATH="$PATH:/bin/aux"
+
+RUN <<EOT
+set -eu
+echo 'debconf debconf/frontend select Noninteractive' | debconf-set-selections
+apt-get update
+apt-get -y install gnupg2 curl
+curl -sL https://nvidia.github.io/nvidia-container-runtime/gpgkey | apt-key add -
+curl -sL https://nvidia.github.io/nvidia-container-runtime/ubuntu20.04/nvidia-container-runtime.list | tee /etc/apt/sources.list.d/nvidia-container-runtime.list
+apt-get update
+apt-get -y install nvidia-container-toolkit-base nvidia-container-toolkit nvidia-container-runtime util-linux
+
+mkdir -vp /etc && echo 'hosts: files dns' > /etc/nsswitch.conf
+chmod 1777 /tmp
+mkdir -vp /var/lib/rancher/k3s/agent/etc/containerd/
+
+apt-get clean
+apt-get autoremove -y
+apt-get autopurge -y
+rm -rf /var/lib/apt/lists/* /var/log/*
+EOT
+
+COPY --from=k3s /bin /bin
+COPY docker/config.toml.tmpl /var/lib/rancher/k3s/agent/etc/containerd/config.toml.tmpl
+
+VOLUME /var/lib/kubelet
+VOLUME /var/lib/rancher/k3s
+VOLUME /var/lib/cni
+VOLUME /var/log
+
+ENTRYPOINT ["/bin/k3s"]
+CMD ["agent"]

docker/Dockerfile.worker

@@ -3,14 +3,10 @@ ARG BASE_STAGE=dev
 
 FROM golang:1.21-bullseye AS golang
 
-RUN <<EOT
-set -eux
-apt-get update
-apt-get install -y --no-install-recommends curl git
-EOT
+RUN apt-get update && apt-get install -y curl git
 
 
-# Skopeo
+# skopeo
 # ========================
 FROM golang AS skopeo
 
@@ -30,13 +26,13 @@ make install
 EOT
 
 
-# RUNC
+# runc
 # ========================
 FROM golang AS runc
 
 WORKDIR /workspace
 
-RUN apt-get install -y --no-install-recommends libseccomp-dev
+RUN apt-get install -y libseccomp-dev
 
 RUN <<EOT
 set -eux
@@ -47,7 +43,22 @@ make install
 EOT
 
 
-# Beam Worker
+# nvidia-container-toolkit
+# ========================
+FROM golang AS nvidia-container-toolkit
+
+WORKDIR /workspace
+
+RUN apt update && apt install -y build-essential
+
+RUN <<EOT
+git clone -b np/update --single-branch https://github.com/beam-cloud/nvidia-container-toolkit.git .
+make build
+make binaries
+EOT
+
+
+# beam worker
 # ========================
 FROM golang AS worker
 
@@ -60,64 +71,36 @@ COPY . .
 RUN go build -o /usr/local/bin/worker ./cmd/worker/main.go
 
 
-# NVIDIA CUDA - Final Stage
+# final image
 # ========================
 FROM nvidia/cuda:12.3.1-base-ubuntu20.04 AS release
 FROM release AS dev
 
 FROM ${BASE_STAGE} AS final
-
 ENV DEBIAN_FRONTEND="noninteractive"
 WORKDIR /workspace
 
-RUN rm /etc/apt/sources.list.d/cuda.list && apt-get update && apt-get install -y wget git curl \
-    libseccomp-dev libsndfile1 libsndfile1-dev \
-    libaio-dev asciidoc xmlto libzmq3-dev iptables \
-    build-essential git libprotobuf-dev libprotobuf-c-dev \
-    protobuf-c-compiler protobuf-compiler \
-    pkg-config libbsd-dev iproute2 \
-    libnftnl-dev libcap-dev libnet1-dev libnl-3-dev \
-    python3 python3.8 python3-distutils python3-pip \
-    curl gpg fuse3
+RUN apt-get update && \
+    apt-get install -y curl gpg fuse3 && \
+    curl -fsSL https://download.opensuse.org/repositories/devel:/tools:/criu/xUbuntu_20.04/Release.key | gpg --dearmor -o /usr/share/keyrings/criu.gpg && \
+    echo 'deb [signed-by=/usr/share/keyrings/criu.gpg] https://download.opensuse.org/repositories/devel:/tools:/criu/xUbuntu_20.04 /' > /etc/apt/sources.list.d/criu.list && \
+    curl -fsSL https://nvidia.github.io/nvidia-container-runtime/gpgkey | apt-key add - && \
+    curl -fsSL https://nvidia.github.io/nvidia-container-runtime/ubuntu20.04/nvidia-container-runtime.list | tee /etc/apt/sources.list.d/nvidia-container-runtime.list && \
+    apt-get update
 
-COPY --from=golang /usr/local/go/ /usr/local/go/
-ENV PATH="/usr/local/go/bin:${PATH}"
+RUN curl -sSL https://d.juicefs.com/install | sh -
+RUN apt-get install -y --no-install-recommends criu nvidia-container-toolkit-base nvidia-container-toolkit
 
-RUN <<EOT
-set -eux
-
-# JuiceFS
-curl -sSL https://d.juicefs.com/install | sh -
-
-# nvidia-container-toolkit repo
-# curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
-# echo 'deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://nvidia.github.io/libnvidia-container/stable/deb/$(ARCH) /' > /etc/apt/sources.list.d/nvidia-container-toolkit.list
-curl -s -L https://nvidia.github.io/nvidia-container-runtime/gpgkey | apt-key add -
-curl -s -L https://nvidia.github.io/nvidia-container-runtime/ubuntu20.04/nvidia-container-runtime.list | tee /etc/apt/sources.list.d/nvidia-container-runtime.list
-apt-get update && apt-get install -y nvidia-container-runtime fuse3 libfuse3-dev
-
-# Build and install custom nvidia-container-toolkit
-git clone https://github.com/beam-cloud/nvidia-container-toolkit.git
-cd /workspace/nvidia-container-toolkit && make build && make binaries
-cp /workspace/nvidia-container-toolkit/nvidia-container-runtime* /usr/bin/
-
-# criu repo
-curl -fsSL https://download.opensuse.org/repositories/devel:/tools:/criu/xUbuntu_20.04/Release.key | gpg --dearmor -o /usr/share/keyrings/criu.gpg
-echo 'deb [signed-by=/usr/share/keyrings/criu.gpg] https://download.opensuse.org/repositories/devel:/tools:/criu/xUbuntu_20.04 /' > /etc/apt/sources.list.d/criu.list
-
-apt-get update
-apt-get install -y --no-install-recommends nvidia-container-runtime criu
-apt-get remove -y curl gpg
-apt-get clean
-apt-get autoremove -y
-apt-get autopurge -y
-rm -rf /var/lib/apt/lists/* /var/log/*
-EOT
+RUN apt-get remove -y curl gpg && \
+    apt-get clean && apt-get autoremove -y && apt-get autopurge -y && \
+    rm -rf /var/lib/apt/lists/* /var/log/*
 
 COPY --from=runc /usr/local/sbin/runc /usr/local/sbin/runc
 COPY --from=skopeo /usr/local/bin/skopeo /usr/local/bin/skopeo
 COPY --from=skopeo /workspace/default-policy.json /etc/containers/policy.json
+COPY --from=nvidia-container-toolkit /workspace/nvidia-container-runtime* /usr/bin/
 COPY --from=worker /usr/local/bin/worker /usr/local/bin/worker
 COPY ./sdk/src/beam /workspace/sdk
 
-VOLUME ["/usr/lib/x86_64-linux-gnu", "/usr/lib/aarch64-linux-gnu"]
+VOLUME "/usr/lib/x86_64-linux-gnu"
+VOLUME "/usr/lib/aarch64-linux-gnu"

docker/config.toml.tmpl

@@ -0,0 +1,118 @@
+version = 2
+
+[plugins."io.containerd.internal.v1.opt"]
+  path = "{{ .NodeConfig.Containerd.Opt }}"
+[plugins."io.containerd.grpc.v1.cri"]
+  stream_server_address = "127.0.0.1"
+  stream_server_port = "10010"
+  enable_selinux = {{ .NodeConfig.SELinux }}
+  enable_unprivileged_ports = {{ .EnableUnprivileged }}
+  enable_unprivileged_icmp = {{ .EnableUnprivileged }}
+
+{{- if .DisableCgroup}}
+  disable_cgroup = true
+{{end}}
+{{- if .IsRunningInUserNS }}
+  disable_apparmor = true
+  restrict_oom_score_adj = true
+{{end}}
+
+{{- if .NodeConfig.AgentConfig.PauseImage }}
+  sandbox_image = "{{ .NodeConfig.AgentConfig.PauseImage }}"
+{{end}}
+
+{{- if .NodeConfig.AgentConfig.Snapshotter }}
+[plugins."io.containerd.grpc.v1.cri".containerd]
+  snapshotter = "{{ .NodeConfig.AgentConfig.Snapshotter }}"
+  disable_snapshot_annotations = {{ if eq .NodeConfig.AgentConfig.Snapshotter "stargz" }}false{{else}}true{{end}}
+  default_runtime_name = "nvidia"
+{{ if eq .NodeConfig.AgentConfig.Snapshotter "stargz" }}
+{{ if .NodeConfig.AgentConfig.ImageServiceSocket }}
+[plugins."io.containerd.snapshotter.v1.stargz"]
+cri_keychain_image_service_path = "{{ .NodeConfig.AgentConfig.ImageServiceSocket }}"
+[plugins."io.containerd.snapshotter.v1.stargz".cri_keychain]
+enable_keychain = true
+{{end}}
+{{ if .PrivateRegistryConfig }}
+{{ if .PrivateRegistryConfig.Mirrors }}
+[plugins."io.containerd.snapshotter.v1.stargz".registry.mirrors]{{end}}
+{{range $k, $v := .PrivateRegistryConfig.Mirrors }}
+[plugins."io.containerd.snapshotter.v1.stargz".registry.mirrors."{{$k}}"]
+  endpoint = [{{range $i, $j := $v.Endpoints}}{{if $i}}, {{end}}{{printf "%q" .}}{{end}}]
+{{if $v.Rewrites}}
+  [plugins."io.containerd.snapshotter.v1.stargz".registry.mirrors."{{$k}}".rewrite]
+{{range $pattern, $replace := $v.Rewrites}}
+    "{{$pattern}}" = "{{$replace}}"
+{{end}}
+{{end}}
+{{end}}
+{{range $k, $v := .PrivateRegistryConfig.Configs }}
+{{ if $v.Auth }}
+[plugins."io.containerd.snapshotter.v1.stargz".registry.configs."{{$k}}".auth]
+  {{ if $v.Auth.Username }}username = {{ printf "%q" $v.Auth.Username }}{{end}}
+  {{ if $v.Auth.Password }}password = {{ printf "%q" $v.Auth.Password }}{{end}}
+  {{ if $v.Auth.Auth }}auth = {{ printf "%q" $v.Auth.Auth }}{{end}}
+  {{ if $v.Auth.IdentityToken }}identitytoken = {{ printf "%q" $v.Auth.IdentityToken }}{{end}}
+{{end}}
+{{ if $v.TLS }}
+[plugins."io.containerd.snapshotter.v1.stargz".registry.configs."{{$k}}".tls]
+  {{ if $v.TLS.CAFile }}ca_file = "{{ $v.TLS.CAFile }}"{{end}}
+  {{ if $v.TLS.CertFile }}cert_file = "{{ $v.TLS.CertFile }}"{{end}}
+  {{ if $v.TLS.KeyFile }}key_file = "{{ $v.TLS.KeyFile }}"{{end}}
+  {{ if $v.TLS.InsecureSkipVerify }}insecure_skip_verify = true{{end}}
+{{end}}
+{{end}}
+{{end}}
+{{end}}
+{{end}}
+
+{{- if not .NodeConfig.NoFlannel }}
+[plugins."io.containerd.grpc.v1.cri".cni]
+  bin_dir = "{{ .NodeConfig.AgentConfig.CNIBinDir }}"
+  conf_dir = "{{ .NodeConfig.AgentConfig.CNIConfDir }}"
+{{end}}
+
+[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.runc]
+  runtime_type = "io.containerd.runc.v2"
+
+[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.runc.options]
+  SystemdCgroup = {{ .SystemdCgroup }}
+
+{{ if .PrivateRegistryConfig }}
+{{ if .PrivateRegistryConfig.Mirrors }}
+[plugins."io.containerd.grpc.v1.cri".registry.mirrors]{{end}}
+{{range $k, $v := .PrivateRegistryConfig.Mirrors }}
+[plugins."io.containerd.grpc.v1.cri".registry.mirrors."{{$k}}"]
+  endpoint = [{{range $i, $j := $v.Endpoints}}{{if $i}}, {{end}}{{printf "%q" .}}{{end}}]
+{{if $v.Rewrites}}
+  [plugins."io.containerd.grpc.v1.cri".registry.mirrors."{{$k}}".rewrite]
+{{range $pattern, $replace := $v.Rewrites}}
+    "{{$pattern}}" = "{{$replace}}"
+{{end}}
+{{end}}
+{{end}}
+
+{{range $k, $v := .PrivateRegistryConfig.Configs }}
+{{ if $v.Auth }}
+[plugins."io.containerd.grpc.v1.cri".registry.configs."{{$k}}".auth]
+  {{ if $v.Auth.Username }}username = {{ printf "%q" $v.Auth.Username }}{{end}}
+  {{ if $v.Auth.Password }}password = {{ printf "%q" $v.Auth.Password }}{{end}}
+  {{ if $v.Auth.Auth }}auth = {{ printf "%q" $v.Auth.Auth }}{{end}}
+  {{ if $v.Auth.IdentityToken }}identitytoken = {{ printf "%q" $v.Auth.IdentityToken }}{{end}}
+{{end}}
+{{ if $v.TLS }}
+[plugins."io.containerd.grpc.v1.cri".registry.configs."{{$k}}".tls]
+  {{ if $v.TLS.CAFile }}ca_file = "{{ $v.TLS.CAFile }}"{{end}}
+  {{ if $v.TLS.CertFile }}cert_file = "{{ $v.TLS.CertFile }}"{{end}}
+  {{ if $v.TLS.KeyFile }}key_file = "{{ $v.TLS.KeyFile }}"{{end}}
+  {{ if $v.TLS.InsecureSkipVerify }}insecure_skip_verify = true{{end}}
+{{end}}
+{{end}}
+{{end}}
+
+{{range $k, $v := .ExtraRuntimes}}
+[plugins."io.containerd.grpc.v1.cri".containerd.runtimes."{{$k}}"]
+  runtime_type = "{{$v.RuntimeType}}"
+[plugins."io.containerd.grpc.v1.cri".containerd.runtimes."{{$k}}".options]
+  BinaryName = "{{$v.BinaryName}}"
+{{end}}

hack/k3d.yaml

@@ -3,7 +3,7 @@ apiVersion: k3d.io/v1alpha5
 kind: Simple
 metadata:
   name: beam
-image: docker.io/rancher/k3s:v1.28.5-k3s1
+image: localhost:5001/rancher/k3s:latest
 servers: 1
 
 ports:

internal/worker/metrics_linux.go

@@ -9,7 +9,9 @@ import (
 )
 
 func (wm *WorkerMetrics) InitNvml() {
-	wm.nvmlActive = nvml.Init() == nvml.SUCCESS
+	// TODO: investigate segmentation violation
+	// wm.nvmlActive = nvml.Init() == nvml.SUCCESS
+	wm.nvmlActive = false
 }
 
 func (wm *WorkerMetrics) Shutdown() {

manifests/k3d/nvidia-device-plugin.yaml

@@ -21,6 +21,7 @@ spec:
         operator: Exists
         effect: NoSchedule
       priorityClassName: system-node-critical
+      runtimeClassName: nvidia
       containers:
       - image: nvcr.io/nvidia/k8s-device-plugin:v0.14.3
         name: nvidia-device-plugin-ctr