Model training guide

ai-ml/model-train/manifests/01-volumes/bucket.yaml

@@ -0,0 +1,32 @@
+---
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: gcs-fuse-pv
+spec:
+  accessModes:
+  - ReadWriteMany
+  capacity:
+    storage: 16Gi
+  storageClassName: example-storage-class
+  mountOptions:
+    - implicit-dirs
+  csi:
+    driver: gcsfuse.csi.storage.gke.io
+    volumeHandle: <PROJECT_ID>-<CLUSTER_PREFIX>-model-train
+    volumeAttributes:
+      fileCacheCapacity: 5Gi
+      fileCacheForRangeRead: "true"
+---
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: gcs-fuse-claim
+spec:
+  accessModes:
+  - ReadWriteMany
+  resources:
+    requests:
+      storage: 16Gi
+  volumeName: gcs-fuse-pv
+  storageClassName: example-storage-class

ai-ml/model-train/manifests/01-volumes/cloudbuild.yaml

@@ -0,0 +1,15 @@
+steps:
+- name: 'gcr.io/google.com/cloudsdktool/cloud-sdk'
+  entrypoint: /bin/bash
+  args: 
+    - '-c'
+    - |
+      gcloud compute ssh --tunnel-through-iap --quiet cloudbuild@${_INSTANCE_NAME} --zone=${_ZONE} --command="\
+      sudo mkdir -p /mnt/disks/ram-disk && \
+      sudo mount -t tmpfs -o size=16g tmpfs /mnt/disks/ram-disk && \
+      sudo mkfs.ext4 -F /dev/disk/by-id/google-local-ssd-block0 && \
+      sudo mkdir -p /mnt/disks/ssd0 && \
+      sudo mount /dev/disk/by-id/google-local-ssd-block0 /mnt/disks/ssd0"
+substitutions:
+  _ZONE: us-central1-a
+  _INSTANCE_NAME: model-train-vm

ai-ml/model-train/manifests/01-volumes/volumes.yaml

@@ -0,0 +1,69 @@
+---
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: local-ssd-pv
+spec:
+  capacity:
+    storage: 16Gi
+  accessModes: ["ReadWriteOnce"]
+  persistentVolumeReclaimPolicy: Retain
+  storageClassName: local-storage
+  local:
+    path: /mnt/disks/ssd0
+  nodeAffinity:
+    required:
+      nodeSelectorTerms:
+      - matchExpressions:
+        - key: "node_pool"
+          operator: "In"
+          values:
+          - "model-train-pool"
+---
+  kind: PersistentVolumeClaim
+  apiVersion: v1
+  metadata:
+    name: local-ssd-claim
+  spec:
+    accessModes:
+    - ReadWriteOnce
+    storageClassName: local-storage
+    volumeName: local-ssd-pv
+    resources:
+      requests:
+        storage: 16Gi
+---
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: ram-disk-pv
+spec:
+  capacity:
+    storage: 16Gi
+  accessModes: ["ReadWriteOnce"]
+  persistentVolumeReclaimPolicy: Retain
+  storageClassName: local-storage
+  local:
+    path: /mnt/disks/ram-disk
+  nodeAffinity:
+    required:
+      nodeSelectorTerms:
+      - matchExpressions:
+        - key: "node_pool"
+          operator: "In"
+          values:
+          - "model-train-pool"
+---
+  kind: PersistentVolumeClaim
+  apiVersion: v1
+  metadata:
+    name: ram-disk-claim
+  spec:
+    accessModes:
+    - ReadWriteOnce
+    storageClassName: local-storage
+    volumeName: ram-disk-pv
+    resources:
+      requests:
+        storage: 16Gi
+

ai-ml/model-train/manifests/02-dataset/job.yaml

@@ -0,0 +1,93 @@
+---
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: bucket-access
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: download-script
+data:
+  download.sh: |-
+    #!/usr/bin/bash -x
+    apt-get update -y && \
+      apt-get install -y --no-install-recommends \
+      git git-lfs rsync
+    git lfs install 
+    cd /local-ssd
+    echo "Saving dataset into Local SSD..."
+    time git clone --depth=1 "$DATASET_REPO"
+    echo "Saving dataset into Ram disk..."
+    time rsync --info=progress2 -a /local-ssd/dataset/dataset/ /ram-disk/dataset/
+    echo "Saving dataset into Bucket..."
+    time gsutil -q -m cp -r /local-ssd/dataset/dataset/ gs://$BUCKET_NAME/dataset/
+    echo "Dataset was successfully saved in all storages!"
+---
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: dataset-downloader
+  labels:
+    app: dataset-downloader
+spec:
+  ttlSecondsAfterFinished: 120
+  template:
+    metadata:
+      labels:
+        app: dataset-downloader
+    spec:
+      restartPolicy: OnFailure
+      serviceAccountName: bucket-access
+      containers:
+      - name: gcloud
+        image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
+        resources:
+          requests:
+            cpu: "2"
+            memory: "12Gi"
+          limits:
+            cpu: "2"
+            memory: "12Gi"
+        command:
+        - /scripts/download.sh
+        env:
+        - name: BUCKET_NAME
+          value: <PROJECT_ID>-<CLUSTER_PREFIX>-model-train
+        - name: DATASET_REPO
+          value: "https://huggingface.co/datasets/dganochenko/dataset"
+        - name: TIMEFORMAT
+          value: "%0lR"
+        volumeMounts:
+        - name: local-ssd-storage
+          mountPath: /local-ssd
+        - name: ram-disk-storage
+          mountPath: /ram-disk
+        - name: scripts-volume
+          mountPath: "/scripts/"
+          readOnly: true
+      volumes:
+      - name: scripts-volume
+        configMap:
+          defaultMode: 0700
+          name: download-script
+      - name: local-ssd-storage
+        persistentVolumeClaim:
+          claimName: local-ssd-claim
+      - name: ram-disk-storage
+        persistentVolumeClaim:
+          claimName: ram-disk-claim
+      - name: gcs-fuse-storage
+        persistentVolumeClaim:    
+          claimName: gcs-fuse-claim
+      tolerations:
+      - key: "nvidia.com/gpu"
+        operator: "Equal"
+        value: "present"
+        effect: NoSchedule
+      - key: "app.stateful/component"
+        operator: "Equal"
+        value: "model-train"
+        effect: NoSchedule
+      nodeSelector:
+        cloud.google.com/gke-accelerator: nvidia-l4