# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: stable-diffusion-deployment
  labels:
    app: stable-diffusion
spec:
  replicas: 1
  selector:
    matchLabels:
      app: stable-diffusion
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/cpu-limit: 500m
        gke-gcsfuse/memory-limit: 100Mi
        gke-gcsfuse/ephemeral-storage-limit: 50Gi
      labels:
        app: stable-diffusion
    spec:
      # nodeSelector:
      #   cloud.google.com/gke-gpu-sharing-strategy: time-sharing
      #   cloud.google.com/gke-max-shared-clients-per-gpu: "2"
      affinity:
        nodeAffinity:
          preferredDuringSchedulingIgnoredDuringExecution:
          - weight: 1
            preference:
              matchExpressions:
              - key: cloud.google.com/gke-spot
                operator: In
                values:
                - "true"
      volumes:
        - name: stable-diffusion-storage
          persistentVolumeClaim:
            claimName: gcs-fuse-csi-static-pvc
      containers:
      - name: stable-diffusion-webui
        image: $image_url # e.g. us-central1-docker.pkg.dev/<PROJECT_ID>/<BUILD_REGIST>/sd-webui:0.1
        imagePullPolicy: Always
        env:
        - name: SAFETENSORS_FAST_GPU
          value: "1"
        - name: LD_PRELOAD
          value: "/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4"
        resources:
          limits:
            cpu: "1000m"
            memory: "15Gi"
            nvidia.com/gpu: "1"
        ports:
          - containerPort: 7860
        # Make sure subPath folders are already in the root dir of the bucket
        # If you want to mount specific path of the bucket, use --only-dir
        # https://cloud.google.com/kubernetes-engine/docs/how-to/persistent-volumes/cloud-storage-fuse-csi-driver#mounting-flags
        volumeMounts:
          - mountPath: "/stable-diffusion-webui/models/"
            name: stable-diffusion-storage
            subPath: models/
          - mountPath: "/stable-diffusion-webui/embeddings/"
            name: stable-diffusion-storage
            subPath: embeddings/
      # Follow this for service account setup,
      # https://github.com/GoogleCloudPlatform/gcs-fuse-csi-driver/blob/main/docs/authentication.md
      serviceAccountName: $k8s_serviceaccount