feat: Use Amazon S3 as Stable Diffusion Model Storage

.gitignore

@@ -60,5 +60,6 @@ site
 # node modules
 node_modules
 gen-ai/inference/stable-diffusion-rayserve-gpu/locust/__pycache__/*
+gen-ai/inference/stable-diffusion-rayserve-gpu/stable-diffusion-2/*
 website/package-lock.json
 website/package.json

ai-ml/jark-stack/terraform/addons.tf

@@ -82,8 +82,18 @@ module "eks_blueprints_addons" {
     vpc-cni = {
       preserve = true
     }
+
+    aws-mountpoint-s3-csi-driver = {
+      service_account_role_arn = module.s3_csi_driver_irsa.iam_role_arn
+      configuration_values     = <<-EOF
+      node:
+        tolerateAllTaints: true
+      EOF
+    }
   }
 
+
+
   #---------------------------------------
   # AWS Load Balancer Controller Add-on
   #---------------------------------------
@@ -354,6 +364,28 @@ module "data_addons" {
   ]
 }
 
+#---------------------------------------------------------------
+# IRSA for Mountpoint for Amazon S3 CSI Driver
+#---------------------------------------------------------------
+module "s3_csi_driver_irsa" {
+  source           = "terraform-aws-modules/iam/aws//modules/iam-role-for-service-accounts-eks"
+  version          = "~> 5.34"
+  role_name_prefix = format("%s-%s-", local.name, "s3-csi-driver")
+  role_policy_arns = {
+    # WARNING: Demo purpose only. Bring your own IAM policy with least privileges
+    s3_csi_driver = "arn:aws:iam::aws:policy/AmazonS3FullAccess"
+  }
+  oidc_providers = {
+    main = {
+      provider_arn               = module.eks.oidc_provider_arn
+      namespace_service_accounts = ["kube-system:s3-csi-driver-sa"]
+    }
+  }
+  tags = local.tags
+}
+
+
+
 
 #---------------------------------------------------------------
 # Additional Resources
@@ -399,3 +431,8 @@ data "aws_iam_policy_document" "karpenter_controller_policy" {
     sid       = "KarpenterControllerAdditionalPolicy"
   }
 }
+
+resource "aws_s3_bucket" "model_storage" {
+  count         = var.create_s3_bucket ? 1 : 0
+  bucket_prefix = "model-storage-"
+}

ai-ml/jark-stack/terraform/outputs.tf

@@ -2,3 +2,11 @@ output "configure_kubectl" {
   description = "Configure kubectl: make sure you're logged in with the correct AWS profile and run the following command to update your kubeconfig"
   value       = "aws eks --region ${var.region} update-kubeconfig --name ${var.name}"
 }
+output "grafana_secret_name" {
+  description = "The name of the secret containing the Grafana admin password."
+  value       = aws_secretsmanager_secret.grafana.name
+}
+output "model_s3_bucket" {
+  description = "The S3 bucket name for storing ML models"
+  value       = length(aws_s3_bucket.model_storage) > 0 ? aws_s3_bucket.model_storage[0].bucket : ""
+}

ai-ml/jark-stack/terraform/variables.tf

@@ -56,5 +56,10 @@ variable "bottlerocket_data_disk_snpashot_id" {
   description = "Bottlerocket Data Disk Snapshot ID"
   type        = string
   default     = ""
+}
 
+variable "create_s3_bucket" {
+  description = "Create S3 Bucket for Model Storage"
+  default     = false
+  type        = bool
 }

ai-ml/jark-stack/terraform/vpc.tf

@@ -51,3 +51,30 @@ module "vpc" {
 
   tags = local.tags
 }
+
+
+module "vpc_endpoints" {
+  source  = "terraform-aws-modules/vpc/aws//modules/vpc-endpoints"
+  version = "~> 5.0"
+
+  # create = var.enable_vpc_endpoints
+
+  create = true
+
+  vpc_id             = module.vpc.vpc_id
+  security_group_ids = [module.vpc_endpoints_sg.security_group_id]
+
+  endpoints = merge({
+    s3 = {
+      service         = "s3"
+      service_type    = "Gateway"
+      route_table_ids = module.vpc.private_route_table_ids
+      tags = {
+        Name = "${local.name}-s3"
+      }
+    }
+    }
+  )
+
+  tags = local.tags
+}

gen-ai/inference/stable-diffusion-rayserve-gpu/Dockerfile

@@ -1,8 +1,8 @@
 # https://hub.docker.com/layers/rayproject/ray-ml/2.10.0-py310-gpu/images/sha256-4181ed53b0b25a758b155312ca6ab29a65cb78cd57296d42cfbe4806a2b77df4?context=explore
-# docker buildx build --platform=linux/amd64 -t ray2.10.0-py310-gpu-stablediffusion:v1.0 -f Dockerfile .
+# docker buildx build --platform=linux/amd64 -t ray2.24.0-py310-gpu-stablediffusion:v1.0 -f Dockerfile .
 
 # Use Ray base image
-FROM rayproject/ray-ml:2.10.0-py310-gpu
+FROM docker.io/rayproject/ray:2.33.0-py311-gpu
 
 # Maintainer label
 LABEL maintainer="DoEKS"
@@ -14,13 +14,14 @@ ENV DEBIAN_FRONTEND=non-interactive
 USER $USER
 
 # Install Ray Serve and other Python packages with specific versions
-RUN pip install --no-cache-dir requests torch "diffusers==0.12.1" "transformers==4.25.1"
+RUN pip install --no-cache-dir requests torch "diffusers==0.29.2" "transformers==4.42.4" "accelerate==0.30.1"
 
 # Set a working directory
 WORKDIR /serve_app
 
 # Copy your Ray Serve script into the container
-COPY ray_serve_stablediffusion.py /serve_app/ray_serve_stablediffusion.py
+# The serving script can moved to the ConfigMap, or copied in the container image
+# COPY ray_serve_sd.py /serve_app/ray_serve_sd.py
 
 # Set the PYTHONPATH environment variable
 ENV PYTHONPATH=/serve_app:$PYTHONPATH

gen-ai/inference/stable-diffusion-rayserve-gpu/download-models-to-s3.yaml

@@ -0,0 +1,119 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: stablediffusion
+---
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: s3-pv-model-storage
+spec:
+  capacity:
+    storage: 1200Gi # ignored, required
+  accessModes:
+    - ReadWriteMany # supported options: ReadWriteMany / ReadOnlyMany
+  mountOptions:
+    - allow-overwrite
+    - allow-delete
+    - region us-west-2
+    - prefix stable-diffusion-2/
+  csi:
+    driver: s3.csi.aws.com # required
+    volumeHandle: s3-csi-driver-volume
+    volumeAttributes:
+      bucketName: <YOUR_BUCKET_NAME> # Replace bucketName with S3 bucket name created from `terraform output`
+---
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: s3-model-storage-claim
+  namespace: stablediffusion
+spec:
+  accessModes:
+    - ReadWriteMany # supported options: ReadWriteMany / ReadOnlyMany
+  storageClassName: "" # required for static provisioning
+  resources:
+    requests:
+      storage: 1200Gi # ignored, required
+  volumeName: s3-pv-model-storage
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: shell-script-configmap
+data:
+  download_models.py: |
+    from diffusers import StableDiffusionPipeline
+    import torch
+
+    # Set the model you want to download
+    model_name = "stabilityai/stable-diffusion-2"
+    model_directory = "/serve_app/temp-stable-diffusion-2"
+
+    # Load the model
+    pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16, cache_dir=model_directory)
+
+    # Save the model to the local directory
+    pipe.save_pretrained(model_directory)
+
+    print(f"Model saved to {model_directory}")
+
+    import shutil, logging
+
+    # Source directory (the one you want to copy)
+    src_dir = model_directory
+
+    # Destination directory (where you want to copy to)
+    dst_dir = "/serve_app/stable-diffusion-2/"
+
+    shutil.rmtree(f"{model_directory}/.locks")
+    shutil.rmtree(f"{model_directory}/models--stabilityai--stable-diffusion-2")
+
+    # Copy the directory recursively
+    try:
+        shutil.copytree(src_dir, dst_dir, dirs_exist_ok=True)
+        logging.info("Directory copied successfully")
+    except shutil.Error as e:
+        logging.error(f"Error copying directory: {e}")
+    except OSError as e:
+        logging.error(f"OS error occurred: {e}")
+    except Exception as e:
+        logging.error(f"Unexpected error occurred: {e}")
+
+    import logging
+
+    notice_message = """
+    NOTICE: shutil.copytree() may generate errors when used with Amazon S3.
+    This is because it calls copystat(src, dst) at:
+        /home/ray/anaconda3/lib/python3.11/shutil.py, line 527, in _copytree
+    Amazon S3 does not support modifying metadata, which causes these errors.
+    """
+
+    logging.warning(notice_message.strip())
+---
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: shell-script-job
+spec:
+  template:
+    spec:
+      containers:
+      - name: shell-script-container
+        image: public.ecr.aws/data-on-eks/ray-serve-gpu-stablediffusion:2.33.0-py311-gpu
+        command: ["python"]
+        args: ["/scripts/download_models.py"]
+        volumeMounts:
+        - name: script-volume
+          mountPath: /scripts
+        - mountPath: /serve_app/stable-diffusion-2
+          name: cache-dir
+      volumes:
+      - name: cache-dir
+        persistentVolumeClaim:
+          claimName: s3-model-storage-claim
+      - name: script-volume
+        configMap:
+          name: shell-script-configmap
+      restartPolicy: Never
+  backoffLimit: 4