פריסת אפליקציית Ray Serve עם מודל דיפוזיה יציבה ב-Google Kubernetes Engine‏ (GKE) באמצעות TPU

Standard

במדריך הזה נדגים איך פורסים מודל Stable Diffusion ב-Google Kubernetes Engine ‏ (GKE) ומציגים אותו באמצעות TPU,‏ Ray Serve ותוסף Ray Operator.

המדריך הזה מיועד ללקוחות של AI גנרטיבי, למשתמשים חדשים או קיימים ב-GKE, למהנדסי למידת מכונה, למהנדסי MLOps‏ (DevOps) או למנהלי פלטפורמות שרוצים להשתמש ביכולות של Kubernetes לתזמור קונטיינרים כדי להציג מודלים באמצעות Ray.

מידע על Ray ו-Ray Serve

‫Ray היא מסגרת מחשוב מבוססת קוד פתוח שניתנת להרחבה, ומיועדת לאפליקציות של AI/ML. ‫Ray Serve היא ספרייה לפרסום מודלים ב-Ray, שמשמשת להרחבה ולפרסום של מודלים בסביבה מבוזרת. מידע נוסף זמין במאמר בנושא Ray Serve במסמכי התיעוד של Ray.

מידע על מעבדי TPU

יחידות לעיבוד טנסורים (TPU) הן מאיצי חומרה ייעודיים שנועדו להאיץ באופן משמעותי את האימון וההסקה של מודלים של למידת מכונה בקנה מידה גדול. השימוש ב-Ray עם יחידות TPU מאפשר לכם להרחיב בצורה חלקה אפליקציות ML עם ביצועים גבוהים. למידע נוסף על יחידות TPU, קראו את המאמר מבוא ל-Cloud TPU במסמכי התיעוד של Cloud TPU.

מידע על תגובה לפעולה מאתר אחר (webhook) של אתחול TPU ב-KubeRay

כחלק מתוסף Ray Operator,‏ GKE מספק ווּבּהוּקים לאימות ולשינוי, שמטפלים בתזמון של TPU Pod ובמשתני סביבה מסוימים של TPU שנדרשים על ידי מסגרות כמו JAX לאתחול של קונטיינרים. ה-webhook של KubeRay TPU משנה את ה-Pods עם התווית app.kubernetes.io/name: kuberay שמבקשת מעבדי TPU עם המאפיינים הבאים:

• ‫TPU_WORKER_ID: מספר שלם ייחודי לכל Pod של Worker בפרוסת ה-TPU.
• ‫TPU_WORKER_HOSTNAMES: רשימה של שמות מארחים ב-DNS של כל העובדים ב-TPU שצריכים לתקשר אחד עם השני בתוך הפרוסה. המשתנה הזה מוזרק רק ל-TPU Pods בקבוצה עם כמה מארחים.
• ‫replicaIndex: תווית Pod שמכילה מזהה ייחודי של הרפליקה של קבוצת ה-workers שאליה שייך ה-Pod. האפשרות הזו שימושית לקבוצות של workers בכמה מארחים, שבהן כמה Pods של worker עשויים להשתייך לאותה רפליקה. היא משמשת את Ray כדי להפעיל התאמה אוטומטית לעומס בכמה מארחים.
• ‫TPU_NAME: מחרוזת שמייצגת את ה-PodSlice של GKE TPU שאליו שייך ה-Pod הזה. הערך שלה זהה לערך של התווית replicaIndex.
• podAffinity: מוודא שמערכת GKE מתזמנת את ה-TPU Pods עם תוויות תואמות של replicaIndex באותו מאגר צמתים. כך GKE יכול לבצע שינוי גודל של מערכי TPU מרובי-מארחים באופן אטומי לפי מאגרי צמתים, ולא לפי צמתים בודדים.

מטרות

• יוצרים אשכול GKE עם מאגר צמתים של TPU.
• פריסת אשכול Ray עם מעבדי TPU.
• פורסים משאב בהתאמה אישית של RayService.
• יצירת אינטראקציה עם שרת המודל Stable Diffusion.

עלויות

במסמך הזה משתמשים ברכיבים הבאים של Google Cloud, והשימוש בהם כרוך בתשלום:

• GKE

בעזרת מחשבון העלויות אתם יכולים להעריך את העלויות בהתאם לשימוש הצפוי.

משתמשים חדשים של Google Cloud ? יכול להיות שאתם זכאים לתקופת ניסיון בחינם.

כשמסיימים את המשימות שמתוארות במסמך הזה אפשר למחוק את המשאבים שיצרתם כדי להימנע מחיובים נוספים. מידע נוסף זמין בקטע הסרת המשאבים.

לפני שמתחילים

ב-Cloud Shell מותקנת מראש התוכנה שדרושה למדריך הזה, כולל kubectl ו-ה-CLI של gcloud. אם אתם לא משתמשים ב-Cloud Shell, התקינו את ה-CLI של gcloud.

נכנסים לחשבון Google Cloud . אנחנו ממליצים למשתמשים חדשים ב- Google Cloud ליצור חשבון כדי שיוכלו להעריך את הביצועים של המוצרים שלנו בתרחישים מהעולם האמיתי. לקוחות חדשים מקבלים בחינם גם קרדיט בשווי 300$ להרצה, לבדיקה ולפריסה של עומסי העבודה.

התקינו את ה-CLI של Google Cloud.

הערה: אם התקנתם את ה-CLI של gcloud, השתמשו בפקודה gcloud components update כדי לבדוק אם הגרסה העדכנית מותקנת.

אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

gcloud init

יוצרים או בוחרים פרויקט ב- Google Cloud .

התפקידים שנדרשים כדי לבחור או ליצור פרויקט

• Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
• יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים?

הערה: אם אתם לא מתכננים לשמור את המשאבים שתיצרו בתהליך הזה, תוכלו ליצור פרויקט חדש במקום לבחור באחד מהפרויקטים הקיימים. בסיום התהליך תוכלו למחוק את הפרויקט ולהסיר את כל המשאבים שמשויכים אליו.

• יוצרים פרויקט ב- Google Cloud :

  gcloud projects create PROJECT_ID

  מחליפים את PROJECT_ID בשם של הפרויקט ב- Google Cloud שאתם יוצרים.

• בוחרים את הפרויקט שיצרתם: Google Cloud

  gcloud config set project PROJECT_ID

  מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

מוודאים שהחיוב מופעל בפרויקט ב- Google Cloud .

מפעילים את GKE API, אם הוא עדיין לא מופעל:

התפקידים שנדרשים כדי להפעיל ממשקי API

כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק Service Usage' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים

gcloud services enable container.googleapis.com

התקינו את ה-CLI של Google Cloud.

הערה: אם התקנתם את ה-CLI של gcloud, השתמשו בפקודה gcloud components update כדי לבדוק אם הגרסה העדכנית מותקנת.

אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

gcloud init

יוצרים או בוחרים פרויקט ב- Google Cloud .

התפקידים שנדרשים כדי לבחור או ליצור פרויקט

• Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
• יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים?

הערה: אם אתם לא מתכננים לשמור את המשאבים שתיצרו בתהליך הזה, תוכלו ליצור פרויקט חדש במקום לבחור באחד מהפרויקטים הקיימים. בסיום התהליך תוכלו למחוק את הפרויקט ולהסיר את כל המשאבים שמשויכים אליו.

• יוצרים פרויקט ב- Google Cloud :

  gcloud projects create PROJECT_ID

  מחליפים את PROJECT_ID בשם של הפרויקט ב- Google Cloud שאתם יוצרים.

• בוחרים את הפרויקט שיצרתם: Google Cloud

  gcloud config set project PROJECT_ID

  מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

מוודאים שהחיוב מופעל בפרויקט ב- Google Cloud .

מפעילים את GKE API, אם הוא עדיין לא מופעל:

התפקידים שנדרשים כדי להפעיל ממשקי API

כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק Service Usage' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים

gcloud services enable container.googleapis.com

1. מעניקים תפקידים לחשבון המשתמש. מריצים את הפקודה הבאה לכל אחד מהתפקידים הבאים ב-IAM: roles/container.clusterAdmin, roles/container.admin

   gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

   מחליפים את מה שכתוב בשדות הבאים:

   • ‫PROJECT_ID: מזהה הפרויקט.
   • ‫USER_IDENTIFIER: המזהה של חשבון המשתמש . לדוגמה, [e-mail].
   • ‫ROLE: תפקיד ה-IAM שאתם מקצים לחשבון המשתמש.

מוודאים שיש מספיק מכסה

מוודאים שיש לפרויקט מכסת TPU מספיקה באזור או בתחום (zone) של Compute Engine. Google Cloud מידע נוסף זמין במאמר הבטחת מכסות מספיקות של TPU ו-GKE במאמרי העזרה של Cloud TPU. יכול להיות שתצטרכו גם להגדיל את המכסות שלכם עבור:

• Persistent Disk SSD‏ (GB)
• כתובות IP בשימוש

הכנת הסביבה

כדי להכין את הסביבה:

1. מפעילים סשן של Cloud Shell מGoogle Cloud המסוף על ידי לחיצה על Activate Cloud Shell . Google Cloud ייפתח סשן בחלונית התחתונה של מסוף Google Cloud .

2. הגדרת משתני סביבה:

   export PROJECT_ID=PROJECT_ID
   export CLUSTER_NAME=ray-cluster
   export COMPUTE_REGION=us-central2-b
   export CLUSTER_VERSION=CLUSTER_VERSION
   

   מחליפים את מה שכתוב בשדות הבאים:

   • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud.
   • ‫CLUSTER_VERSION: גרסת GKE לשימוש. התאריך צריך להיות 1.30.1 או מאוחר יותר.

3. משכפלים את המאגר ב-GitHub:

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
   

4. עוברים לספריית העבודה:

   cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/stable-diffusion
   

יצירת אשכול עם מאגר צומתי TPU

יוצרים אשכול GKE רגיל עם מאגר צמתים של TPU:

1. יוצרים אשכול במצב רגיל עם Ray Operator מופעל:

   gcloud container clusters create ${CLUSTER_NAME} \
       --addons=RayOperator \
       --machine-type=n1-standard-8 \
       --cluster-version=${CLUSTER_VERSION} \
       --location=${COMPUTE_REGION}
   

2. יוצרים מאגר צמתים של TPU במארח יחיד:

   gcloud container node-pools create tpu-pool \
       --location=${COMPUTE_REGION} \
       --cluster=${CLUSTER_NAME} \
       --machine-type=ct4p-hightpu-4t \
       --num-nodes=1
   

כדי להשתמש ביחידות TPU במצב רגיל, צריך לבחור:

• מיקום ב-Compute Engine עם קיבולת למאיצי TPU
• סוג מכונה שתואם ל-TPU ול-
• הטופולוגיה הפיזית של TPU PodSlice

הגדרת משאב RayCluster עם מעבדי TPU

מגדירים את קובץ המניפסט של RayCluster כדי להכין את עומס העבודה של TPU:

הגדרת TPU nodeSelector

‫GKE משתמש ב-nodeSelectors של Kubernetes כדי לוודא שעומסי עבודה של TPU מתוזמנים בטופולוגיה ובמאיץ המתאימים של TPU. מידע נוסף על בחירת TPU nodeSelectors זמין במאמר פריסת עומסי עבודה של TPU ב-GKE Standard.

מעדכנים את המניפסט ray-cluster.yaml כדי לתזמן את ה-Pod ב-podslice של TPU v4 עם טופולוגיה של 2x2x1:

nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
  cloud.google.com/gke-tpu-topology: 2x2x1

הגדרת משאב קונטיינר של TPU

כדי להשתמש במאיץ TPU, צריך לציין את מספר שבבי ה-TPU ש-GKE צריך להקצות לכל Pod. לשם כך, מגדירים את המשאב google.com/tpu ואת המכסה limits requests בשדה של מאגר ה-TPU במניפסט של RayCluster workerGroupSpecs.

מעדכנים את המניפסט ray-cluster.yaml עם בקשות ומגבלות על מקורות מידע:

resources:
  limits:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"
   requests:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"

הגדרת קבוצת עובדים numOfHosts

בגרסה 1.1.0 של KubeRay נוסף השדה numOfHosts למשאב המותאם אישית של RayCluster, שמציין את מספר מארחי ה-TPU שצריך ליצור לכל רפליקה של קבוצת Worker. בקבוצות של כמה מארחים, הרפליקות נחשבות ל-PodSlices ולא לצומתי עובדים נפרדים, ונוצרים numOfHosts צומתי עובדים לכל רפליקה.

מעדכנים את המניפסט ray-cluster.yaml עם הערכים הבאים:

workerGroupSpecs:
  # Several lines omitted
  numOfHosts: 1 # the number of "hosts" or workers per replica

יצירת משאב מותאם אישית מסוג RayService

יוצרים משאב בהתאמה אישית של RayService:

1. בודקים את המניפסט הבא:

   apiVersion: ray.io/v1
   kind: RayService
   metadata:
     name: stable-diffusion-tpu
   spec:
     serveConfigV2: |
       applications:
         - name: stable_diffusion
           import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion_tpu:deployment
           runtime_env:
             working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
             pip:
               - diffusers==0.7.2
               - flax
               - jax[tpu]==0.4.11
               - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
               - fastapi
     rayClusterConfig:
       rayVersion: '2.9.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             - name: ray-head
               image: rayproject/ray-ml:2.9.0-py310
               ports:
               - containerPort: 6379
                 name: gcs
               - containerPort: 8265
                 name: dashboard
               - containerPort: 10001
                 name: client
               - containerPort: 8000
                 name: serve
               resources:
                 limits:
                   cpu: "2"
                   memory: "8G"
                 requests:
                   cpu: "2"
                   memory: "8G"
       workerGroupSpecs:
       - replicas: 1
         minReplicas: 1
         maxReplicas: 10
         numOfHosts: 1
         groupName: tpu-group
         rayStartParams: {}
         template:
           spec:
             containers:
             - name: ray-worker
               image: rayproject/ray-ml:2.9.0-py310
               resources:
                 limits:
                   cpu: "100"
                   ephemeral-storage: 20Gi
                   google.com/tpu: "4"
                   memory: 200G
                 requests:
                   cpu: "100"
                   ephemeral-storage: 20Gi
                   google.com/tpu: "4"
                   memory: 200G
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
               cloud.google.com/gke-tpu-topology: 2x2x1

   קובץ המניפסט הזה מתאר משאב מותאם אישית של RayService שיוצר משאב RayCluster עם צומת ראשי אחד וקבוצת עובדים של TPU עם טופולוגיה של 2x2x1, כלומר לכל צומת עובד יהיו 4 שבבי TPU מדגם v4.

   צומת ה-TPU שייך ל-Podslice יחיד של TPU v4 עם טופולוגיה של 2x2x1. כדי ליצור קבוצת עובדים עם כמה מארחים, מחליפים את הערכים של gke-tpu nodeSelector, את הערכים של google.com/tpu (מגבלות ובקשות של קונטיינר) ואת הערכים של numOfHosts בהגדרות שלכם לקבוצת עובדים עם כמה מארחים. למידע נוסף על טופולוגיות של TPU עם כמה מארחים, קראו את המאמר ארכיטקטורת המערכת במסמכי התיעוד של Cloud TPU.

2. מחילים את המניפסט על האשכול:

   kubectl apply -f ray-service-tpu.yaml
   

3. מוודאים שהמשאב RayService פועל:

   kubectl get rayservices
   

   הפלט אמור להיראות כך:

   NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
   stable-diffusion-tpu   Running          2
   

   בפלט הזה, Running בעמודה SERVICE STATUS מציין שהמשאב RayService מוכן.

(אופציונלי) הצגת מרכז הבקרה של Ray

אתם יכולים לראות את פריסת Ray Serve ואת היומנים הרלוונטיים ממרכז הבקרה של Ray.

1. יוצרים סשן להעברת פורטים ללוח הבקרה של Ray משירות ה-head של Ray:

   kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265
   

2. בדפדפן אינטרנט, עוברים אל http://localhost:8265/.

3. לוחצים על הכרטיסייה Serve (הצגה).

שליחת הנחיות לשרת המודל

1. מקימים סשן להעברת נתונים דרך יציאה לנקודת הקצה של Serve משירות ה-Ray head:

   kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000
   

2. פותחים סשן חדש ב-Cloud Shell.

3. שולחים הנחיה ליצירת תמונה לפי טקסט לשרת המודל של Stable Diffusion:

   python stable_diffusion_tpu_req.py  --save_pictures
   

   התוצאות של הסקת הדיפוזיה היציבה נשמרות בקובץ בשם diffusion_results.png.

   

מעקב אחרי עומסי העבודה ב-Ray

כדי לראות את הפרטים של משימות Ray, נכנסים לקטע Kubernetes Engine > AI/ML > Jobs במסוף Google Cloud .

הצגת RayJobs ב Google Cloud מסוף

הסרת המשאבים

מחיקת הפרויקט

זהירות: כשמוחקים פרויקט:

• התוכן של הפרויקט נמחק כולו. אם השתמשתם בפרויקט קיים כדי לתרגל את המשימות במדריך הזה, כשתמחקו אותו תימחק גם כל העבודה שביצעתם באותו פרויקט.
• מאבדים את מזהה הפרויקט בהתאמה אישית. כשיצרתם את הפרויקט, יכולתם לבחור לו מזהה פרויקט בהתאמה אישית כדי להשתמש בו בעתיד. כדי לשמור על כתובות ה-URL שמשתמשות במזהה הפרויקט, כמו appspot.com, צריך למחוק את המשאבים שנבחרו בפרויקט אבל לא את הפרויקט כולו.

אם אתם מתכננים להיעזר בכמה ארכיטקטורות, מדריכים או מדריכים למתחילים, מומלץ להשתמש שוב באותו פרויקט כדי לא לחרוג ממכסות הפרויקטים.

כדי למחוק Google Cloud פרויקט:

gcloud projects delete PROJECT_ID

מחיקת משאבים בודדים

כדי למחוק את האשכול, מקלידים:

gcloud container clusters delete ${CLUSTER_NAME}

המאמרים הבאים

• מידע על Ray ב-Kubernetes
• למסמכי העזרה של KubeRay
• כדאי להעמיק את הקריאה ולהכיר דוגמאות לארכיטקטורות, תרשימים ושיטות מומלצות בנושאי Google Cloud. כל אלה זמינים במרכז הארכיטקטורה של Cloud.