Added GPU enabled sandbox image. (v2?)

docker/sandbox-bundled/Dockerfile.gpu

@@ -0,0 +1,72 @@
+# syntax=docker/dockerfile:1.4-labs
+
+###### BUILD FLYTE
+FROM --platform=${BUILDPLATFORM} mgoltzsche/podman:minimal AS builder
+
+ARG TARGETARCH
+ENV TARGETARCH "${TARGETARCH}"
+
+WORKDIR /build
+
+COPY images/manifest.txt images/preload ./
+RUN --security=insecure ./preload manifest.txt
+
+FROM --platform=${BUILDPLATFORM} golang:1.19-bullseye AS bootstrap
+
+ARG TARGETARCH
+ENV CGO_ENABLED 0
+ENV GOARCH "${TARGETARCH}"
+ENV GOOS linux
+
+WORKDIR /flyteorg/build
+COPY bootstrap/go.mod bootstrap/go.sum ./
+RUN go mod download
+COPY bootstrap/ ./
+RUN --mount=type=cache,target=/root/.cache/go-build --mount=type=cache,target=/root/go/pkg/mod \
+    go build -o dist/flyte-sandbox-bootstrap cmd/bootstrap/main.go
+
+###### GET K3S
+# ARG K3S_TAG=v1.26.4-k3s1
+FROM rancher/k3s:v1.26.4-k3s1 as k3s
+
+FROM nvidia/cuda:11.8.0-base-ubuntu22.04
+
+ENV CRICTL_VERSION="v1.26.0" 
+ENV FLYTE_GPU "ENABLED"
+ARG TARGETARCH
+
+RUN apt-get update \
+    && apt-get -y install gnupg2 curl nvidia-container-toolkit \
+    && chmod 1777 /tmp \
+    && mkdir -p /var/lib/rancher/k3s/agent/etc/containerd \
+    && mkdir -p /var/lib/rancher/k3s/server/manifests \
+    && curl -L https://github.com/kubernetes-sigs/cri-tools/releases/download/$CRICTL_VERSION/crictl-${CRICTL_VERSION}-linux-amd64.tar.gz --output crictl-${CRICTL_VERSION}-linux-amd64.tar.gz \
+    && tar zxvf crictl-$CRICTL_VERSION-linux-amd64.tar.gz -C /usr/local/bin \
+    && rm -f crictl-$CRICTL_VERSION-linux-amd64.tar.gz \
+    && echo "alias kubectl='k3s kubectl'" >> /root/.bashrc
+
+COPY --from=k3s /bin /bin
+COPY --from=k3s /etc /etc
+
+# Provide custom containerd configuration to configure the nvidia-container-runtime
+COPY config.toml.tmpl /var/lib/rancher/k3s/agent/etc/containerd/config.toml.tmpl
+
+# Deploy the nvidia driver plugin on startup
+COPY device-plugin-daemonset.yaml /var/lib/rancher/k3s/server/manifests/nvidia-device-plugin-daemonset.yaml
+
+COPY --from=builder /build/images/ /var/lib/rancher/k3s/agent/images/
+COPY --from=bootstrap /flyteorg/build/dist/flyte-sandbox-bootstrap /bin/
+COPY images/tar/${TARGETARCH}/ /var/lib/rancher/k3s/agent/images/
+COPY manifests/ /var/lib/rancher/k3s/server/manifests-staging/
+COPY bin/ /bin/
+
+VOLUME /var/lib/kubelet
+VOLUME /var/lib/rancher/k3s
+VOLUME /var/lib/cni
+VOLUME /var/log
+
+ENV PATH="$PATH:/bin/aux"
+ENV CRI_CONFIG_FILE=/var/lib/rancher/k3s/agent/etc/crictl.yaml
+
+ENTRYPOINT [ "/bin/k3d-entrypoint.sh" ]
+CMD [ "server", "--disable=traefik", "--disable=servicelb" ]

docker/sandbox-bundled/Makefile

@@ -33,6 +33,11 @@ manifests:
 		--load-restrictor=LoadRestrictionsNone \
 		kustomize/complete-agent > manifests/complete-agent.yaml
 
+.PHONY: manifests-gpu
+manifests-gpu: manifests
+	cat kustomize/gpu-operator.yaml >> manifests/complete.yaml
+
+
 .PHONY: build
 build: flyte manifests
 	[ -n "$(shell docker buildx ls | awk '/^flyte-sandbox / {print $$1}')" ] || \
@@ -43,6 +48,16 @@ build: flyte manifests
 	docker buildx build --builder flyte-sandbox --allow security.insecure --load \
 		--tag flyte-sandbox:latest .
 
+.PHONY: build-gpu
+build-gpu: flyte manifests-gpu
+	[ -n "$(shell docker buildx ls | awk '/^flyte-sandbox / {print $$1}')" ] || \
+		docker buildx create --name flyte-sandbox \
+		--driver docker-container --driver-opt image=moby/buildkit:master \
+		--buildkitd-flags '--allow-insecure-entitlement security.insecure' \
+		--platform linux/arm64,linux/amd64
+	docker buildx build --builder flyte-sandbox --allow security.insecure --load \
+                --tag flyte-sandbox-gpu:latest -f Dockerfile.gpu .
+
 # Port map
 # 6443 - k8s API server
 # 30000 - Docker Registry

docker/sandbox-bundled/bin/k3d-entrypoint-cgroupv2.sh

@@ -14,8 +14,12 @@ if [ -f /sys/fs/cgroup/cgroup.controllers ]; then
     # move the processes from the root group to the /init group,
   # otherwise writing subtree_control fails with EBUSY.
   mkdir -p /sys/fs/cgroup/init
-  busybox xargs -rn1 < /sys/fs/cgroup/cgroup.procs > /sys/fs/cgroup/init/cgroup.procs || :
+  if command -v busybox >/dev/null 2>&1; then
+    busybox xargs -rn1 < /sys/fs/cgroup/cgroup.procs > /sys/fs/cgroup/init/cgroup.procs || :
+  else
+    xargs -rn1 < /sys/fs/cgroup/cgroup.procs > /sys/fs/cgroup/init/cgroup.procs || :
+  fi
   # enable controllers
-  sed -e 's/ / +/g' -e 's/^/+/' <"/sys/fs/cgroup/cgroup.controllers" >"/sys/fs/cgroup/cgroup.subtree_control"
+  sed -e 's/ / +/g' -e 's/^/+/' < /sys/fs/cgroup/cgroup.controllers > /sys/fs/cgroup/cgroup.subtree_control
   echo "[$(date -Iseconds)] [CgroupV2 Fix] Done"
 fi

docker/sandbox-bundled/bin/k3d-entrypoint-gpu-check.sh

@@ -0,0 +1,15 @@
+#!/bin/sh
+
+if [ -n "${FLYTE_GPU}"  ]; then
+  echo "GPU Enabled - checking if it's available"
+  nvidia-smi
+  if [ $? -eq 0 ]; then
+    echo "nvidia-smi working"
+  else
+    >&2 echo "NVIDIA not available, enable it in docker like so: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/user-guide.html"
+    exit 255 
+  fi
+
+else
+  echo "GPU not enabled"
+fi

docker/sandbox-bundled/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]
+  default_runtime_name = "nvidia"
+  snapshotter = "{{ .NodeConfig.AgentConfig.Snapshotter }}"
+  disable_snapshot_annotations = {{ if eq .NodeConfig.AgentConfig.Snapshotter "stargz" }}false{{else}}true{{end}}
+{{ 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}}

docker/sandbox-bundled/device-plugin-daemonset.yaml

@@ -0,0 +1,41 @@
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: nvidia-device-plugin-daemonset
+  namespace: kube-system
+spec:
+  selector:
+    matchLabels:
+      name: nvidia-device-plugin-ds
+  template:
+    metadata:
+      # Mark this pod as a critical add-on; when enabled, the critical add-on scheduler
+      # reserves resources for critical add-on pods so that they can be rescheduled after
+      # a failure.  This annotation works in tandem with the toleration below.
+      annotations:
+        scheduler.alpha.kubernetes.io/critical-pod: ""
+      labels:
+        name: nvidia-device-plugin-ds
+    spec:
+      tolerations:
+      # Allow this pod to be rescheduled while the node is in "critical add-ons only" mode.
+      # This, along with the annotation above marks this pod as a critical add-on.
+      - key: CriticalAddonsOnly
+        operator: Exists
+      containers:
+      - env:
+        - name: DP_DISABLE_HEALTHCHECKS
+          value: xids
+        image: nvidia/k8s-device-plugin:1.11
+        name: nvidia-device-plugin-ctr
+        securityContext:
+          allowPrivilegeEscalation: true
+          capabilities:
+            drop: ["ALL"]
+        volumeMounts:
+          - name: device-plugin
+            mountPath: /var/lib/kubelet/device-plugins
+      volumes:
+        - name: device-plugin
+          hostPath:
+            path: /var/lib/kubelet/device-plugins

docker/sandbox-bundled/kustomize/gpu-operator.yaml

@@ -0,0 +1,9 @@
+---
+apiVersion: helm.cattle.io/v1
+kind: HelmChart
+metadata:
+  name: nvidia-device-plugin
+  namespace: kube-system
+spec:
+  chart: nvidia-device-plugin
+  repo: https://nvidia.github.io/k8s-device-plugin

docker/sandbox-bundled/manifests/complete-agent.yaml

@@ -816,7 +816,7 @@ type: Opaque
 ---
 apiVersion: v1
 data:
-  haSharedSecret: SWRqSVZmendLdVJvamVGRQ==
+  haSharedSecret: UlZsSHR0c2RFMDdkSFJ0Sw==
   proxyPassword: ""
   proxyUsername: ""
 kind: Secret
@@ -1409,7 +1409,7 @@ spec:
     metadata:
       annotations:
         checksum/config: 8f50e768255a87f078ba8b9879a0c174c3e045ffb46ac8723d2eedbe293c8d81
-        checksum/secret: 4e6248034c039b0bc65cf110d26fccc9d88b8a2d49855a0ecba0443f972ca3c3
+        checksum/secret: ff9d2702885cbb03ce89a5659e44f49b781a55ea62d7b5099ecea38d1781bf53
       labels:
         app: docker-registry
         release: flyte-sandbox
@@ -1744,7 +1744,7 @@ spec:
           value: minio
         - name: FLYTE_AWS_SECRET_ACCESS_KEY
           value: miniostorage
-        image: ghcr.io/flyteorg/flyteagent:1.9.1
+        image: ghcr.io/flyteorg/flyteagent:1.10.0
         imagePullPolicy: IfNotPresent
         name: flyteagent
         ports:

docker/sandbox-bundled/manifests/complete.yaml

@@ -805,7 +805,7 @@ type: Opaque
 ---
 apiVersion: v1
 data:
-  haSharedSecret: UUREcXo3a1VGNnlyc1RCWg==
+  haSharedSecret: OVRlcXBXS2NkUTc0czEwQg==
   proxyPassword: ""
   proxyUsername: ""
 kind: Secret
@@ -1366,7 +1366,7 @@ spec:
     metadata:
       annotations:
         checksum/config: 8f50e768255a87f078ba8b9879a0c174c3e045ffb46ac8723d2eedbe293c8d81
-        checksum/secret: cfe2089fc583f69d068c3b3d56e875082a5d926c70b00b32f094d587df7396a5
+        checksum/secret: 17ec194cf72de1676eef76a26280b2056a6b549bb06b25cf333e4b5f62562ab3
       labels:
         app: docker-registry
         release: flyte-sandbox
@@ -1814,3 +1814,12 @@ spec:
   updateStrategy:
     rollingUpdate: {}
     type: RollingUpdate
+---
+apiVersion: helm.cattle.io/v1
+kind: HelmChart
+metadata:
+  name: nvidia-device-plugin
+  namespace: kube-system
+spec:
+  chart: nvidia-device-plugin
+  repo: https://nvidia.github.io/k8s-device-plugin

docker/sandbox-bundled/manifests/dev.yaml

@@ -499,7 +499,7 @@ metadata:
 ---
 apiVersion: v1
 data:
-  haSharedSecret: VjRtTTRpSW94OU9rUVFZTQ==
+  haSharedSecret: Ulh6cjlwamRWMDNOeG9ycw==
   proxyPassword: ""
   proxyUsername: ""
 kind: Secret
@@ -933,7 +933,7 @@ spec:
     metadata:
       annotations:
         checksum/config: 8f50e768255a87f078ba8b9879a0c174c3e045ffb46ac8723d2eedbe293c8d81
-        checksum/secret: 8b25114d47a69e875e3841dace0b87abfda7a16f6833a5cf76b4570092f3a1cd
+        checksum/secret: 1cea0c45cc972dd53d56c1dc3279a3a24106ad8488a3f8559a8237b96ee040d9
       labels:
         app: docker-registry
         release: flyte-sandbox