איך מתחילים לעבוד עם ML Diagnostics SDK

אפשר לשלב את ML Diagnostics Python SDK עם עומסי עבודה של ML כדי לאסוף ולנהל מדדי מודלים, הגדרות ופרופילים ב- Google Cloud. אפשר להשתמש ב-SDK כדי ליצור פרופילים באופן פרוגרמטי ולראות את מדדי המודל. מדדי המודל לא זמינים אם משתמשים בפרופיל לפי דרישה בלי ה-SDK.

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

מידע נוסף על שימוש ב-ML Diagnostics SDK זמין במאגר google-cloud-mldiagnostics.

התקנה של ML Diagnostics SDK

מתקינים את google-cloud-mldiagnostics הספרייה:

pip install google-cloud-mldiagnostics

מייבאים את החבילות הבאות לקוד של עומס העבודה של ה-ML:

from google_cloud_mldiagnostics import machinelearning_run
from google_cloud_mldiagnostics import metrics
from google_cloud_mldiagnostics import xprof

הפעלת Cloud Logging

ה-SDK משתמש במודול logging הסטנדרטי של Python כדי להפיק מדדים ומידע על ההגדרות. כדי לנתב את היומנים האלה אל Cloud Logging, צריך להתקין ולהגדיר את ספריית google-cloud-logging. כך תוכלו לראות את יומני ה-SDK, את המדדים שנרשמו ואת יומני האפליקציה שלכם במסוף Google Cloud .

מתקינים את הספרייה google-cloud-logging:

pip install google-cloud-logging

מגדירים את הרישום ביומן בסקריפט על ידי צירוף של handler של Cloud Logging ל-logger הבסיסי של Python. מוסיפים את השורות הבאות לתחילת הסקריפט של Python:

  import logging
  import google.cloud.logging

  # Instantiate a Cloud Logging client
  logging_client = google.cloud.logging.Client()

  # Attach the Cloud Logging handler to the Python root logger
  logging_client.setup_logging()

  # Standard logging calls will go to Cloud Logging
  logging.info("SDK logs and application logs will appear in Cloud Logging.")

הפעלת רישום מפורט ביומן

כברירת מחדל, רמת הרישום ביומן מוגדרת כ-INFO. כדי לקבל יומנים מפורטים יותר מ-SDK, כמו פרטים על הרצת למידת מכונה, צריך להגדיר את רמת הרישום ל-DEBUG אחרי הקריאה ל-setup_logging():

import logging
import google.cloud.logging

logging_client = google.cloud.logging.Client()
logging_client.setup_logging()
logging.getLogger().setLevel(logging.DEBUG) # Enable DEBUG level logs

logging.debug("This is a debug message.")
logging.info("This is an info message.")

אם האפשרות DEBUG מופעלת, תקבלו אבחון SDK נוסף ב-Cloud Logging. לדוגמה:

DEBUG:google_cloud_mldiagnostics.core.global_manager:current run details:
{'name': 'projects/my-gcp-project/locations/us-central1/mlRuns/my-run-12345',
'gcs_path': 'gs://my-bucket/profiles', ...}

יצירת הפעלה של למידת מכונה

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

הדוגמה הבסיסית הבאה מאתחלת את Cloud Logging, יוצרת הפעלה של למידת מכונה (MLRun), מתעדת מדדים ויוצרת פרופיל:

import logging
import os
import google.cloud.logging
from google_cloud_mldiagnostics import machinelearning_run, metrics, xprof, metric_types

# 1. Set up Cloud Logging
# Make sure to pip install google-cloud-logging
logging_client = google.cloud.logging.Client()
logging_client.setup_logging()
# Optional: Set logging level to DEBUG for more detailed SDK logs
logging.getLogger().setLevel(logging.DEBUG)

# 2. Define and start machinelearning run
try:
    run = machinelearning_run(
          name="<run_name>",
          run_group="<run_group>",
          configs={ "epochs": 100, "batch_size": 32 },
          project="<some_project>",
          region="<some_zone>",
          gcs_path="gs://<some_bucket>",
          on_demand_xprof=True,
        )
    logging.info(f"MLRun created: {run.name}")

    # 3. Collect metrics during your run
    metrics.record(metric_types.MetricType.LOSS, 0.123, step=1)
    logging.info("Loss metric recorded.")

    # 4. Capture profiles programmatically
    with xprof():
        # ... your code to profile here ...
        pass
    logging.info("Profile captured.")

except Exception as e:
    logging.error(f"Error during MLRun: {e}", exc_info=True)

בדוגמת הקוד נעשה שימוש במשתנים הבאים:

משתנה	דרישה	תיאור
name	חובה	מזהה של הריצה הספציפית. ה-SDK יוצר באופן אוטומטי machine-learning-run-id כדי לוודא שהשמות של ההרצות יהיו ייחודיים.
run_group	אופציונלי	מזהה שיכול לעזור לקבץ כמה הרצות ששייכות לאותו ניסוי. לדוגמה, כל ההרצות שמשויכות לסריקה של גודל פרוסת TPU יכולות להיות שייכות לאותה קבוצה.
project	אופציונלי	אם לא מציינים את הפרויקט, הוא מחולץ מ-Google Cloud CLI.
region	חובה	כל המיקומים של Cluster Director נתמכים, מלבד us-east5. אפשר להגדיר את הדגל הזה באמצעות ארגומנט לכל פקודה, או באמצעות הפקודה: gcloud config set compute/region.
configs	אופציונלי	צמדי מפתח/ערך שמכילים פרמטרים של הגדרות להרצה. אם לא מוגדרות הגדרות, יופיעו הגדרות ברירת המחדל של התוכנה והמערכת, אבל לא הגדרות של עומסי עבודה של ML.
gcs_path	חובה מותנית	Google Cloud מיקום האחסון שבו נשמרים כל הפרופילים. לדוגמה: gs://my-bucket או gs://my-bucket/folder1. חובה רק אם משתמשים ב-SDK לצורך איסוף נתונים של פרופילים.
on-demand-xprof	אופציונלי	מתחיל xprofz daemon ביציאה 9999 כדי להפעיל פרופילים לפי דרישה. אפשר להפעיל גם פרופיל לפי דרישה וגם פרופיל תוכניתי באותו קוד, כל עוד הם לא מתרחשים בו-זמנית.

ההגדרות הבאות נאספות באופן אוטומטי על ידי ה-SDK ואין צורך לציין אותן ב-machinelearning_run:

• הגדרות תוכנה: מסגרת, גרסת מסגרת, דגלי XLA.
• הגדרות מערכת: סוג המכשיר, מספר הפרוסות, גודל הפרוסה, מספר המארחים.

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

כתיבת הגדרות

עומסי עבודה רבים מכילים יותר מדי הגדרות כדי להגדיר אותן ישירות בהגדרה של machinelearning_run. במקרים כאלה, אפשר לכתוב הגדרות לריצה באמצעות JSON או YAML.

import yaml
import json

# Read the YAML file
with open('config.yaml', 'r') as yaml_file:
  # Parse YAML into a Python dictionary
  yaml_data = yaml.safe_load(yaml_file)

# Define machinelearning run
machinelearning_run(
  name="RUN_NAME",
  run_group="GROUP_NAME",
  configs=yaml_data,
  project="PROJECT_NAME",
  region="ZONE",
  gcs_path="gs://BUCKET_NAME",
)

איסוף מדדים

אתם יכולים לאסוף מדדי מודל, מדדי ביצועים של המודל ומדדי מערכת באמצעות ה-SDK. אתם יכולים ליצור המחשות של המדדים האלה כערכים ממוצעים וכתרשימי ציר זמן.

ה-SDK מספק שתי פונקציות לתיעוד מדדים: metrics.record() לתיעוד נקודות נתונים נפרדות ו-metrics.record_metrics() לתיעוד כמה מדדים באצווה אחת. שתי הפונקציות כותבות מדדים ל-Cloud Logging, וכך מאפשרות הצגה חזותית וניתוח.

כדי להקליט מדד יחיד:

# Record a metric only with time as the x-axis
metrics.record(metric_types.MetricType.LOSS, 0.123)

# Record a metric with time and step as the x-axis
metrics.record(metric_types.MetricType.LOSS, 0.123, step=1)

כדי להקליט כמה מדדים:

from google_cloud_mldiagnostics import metric_types
# User codes
# machinelearning_run should be called
# ......

for step in range(num_steps):
  if (step + 1) % 10 == 0:
    metrics.record_metrics([
        # Model quality metrics
        {"metric_name": metric_types.MetricType.LEARNING_RATE, "value": step_size},
        {"metric_name": metric_types.MetricType.LOSS, "value": loss},
        {"metric_name": metric_types.MetricType.GRADIENT_NORM, "value": gradient},
        {"metric_name": metric_types.MetricType.TOTAL_WEIGHTS, "value": total_weights},
        # Model performance metrics
        {"metric_name": metric_types.MetricType.STEP_TIME, "value": step_time},
        {"metric_name": metric_types.MetricType.THROUGHPUT, "value": throughput},
        {"metric_name": metric_types.MetricType.LATENCY, "value": latency},
        {"metric_name": metric_types.MetricType.TFLOPS, "value": tflops},
        {"metric_name": metric_types.MetricType.MFU, "value": mfu},
    ], step=step+1)

מדדי המערכת הבאים נאספים באופן אוטומטי על ידי ה-SDK מהספריות libTPU, psutil ו-JAX:

• ניצול ה-TensorCore של TPU
• דיוטי סייקל (Duty cycle) ב-TPU
• ניצול HBM
• ניצול המעבד של המארח
• ניצול הזיכרון של המארח

לא צריך לציין את המדדים האלה באופן ידני. במדדי המערכת האלה, ציר ה-X מוגדר כברירת מחדל לזמן.

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

• מפתחות של מדדי איכות המודל: LEARNING_RATE, LOSS, GRADIENT_NORM, TOTAL_WEIGHTS.
• מפתחות של מדדי ביצועים של המודל: STEP_TIME, THROUGHPUT, LATENCY, MFU, TFLOPS.

אפשר לתעד את המדדים המוגדרים מראש, וגם מדדים אחרים שהוגדרו על ידי המשתמש, עם ציר X בתור time, או גם time וגם step. אפשר להקליט כל מדד מותאם אישית בעומס העבודה.

בדוגמה הבאה מוצג מדד יחיד של עומס העבודה, שאפשר לראות בכרטיסייה Model Metrics של ההרצה הספציפית של למידת מכונה:

metrics.record("custom_metrics_1", step_size, step=step + 1)

כדי להקליט כמה מדדים בשיחה אחת, משתמשים בשיטה record_metrics. לדוגמה:

metrics.record_metrics([
        # Model quality metrics
        {"metric_name": metric_types.MetricType.LEARNING_RATE, "value": step_size},
        {"metric_name": metric_types.MetricType.LOSS, "value": loss},
        {"metric_name": metric_types.MetricType.GRADIENT_NORM, "value": gradient},
        {"metric_name": metric_types.MetricType.TOTAL_WEIGHTS, "value": total_weights},
        # Model performance metrics
        {"metric_name": metric_types.MetricType.STEP_TIME, "value": step_time},
        {"metric_name": metric_types.MetricType.THROUGHPUT, "value": throughput},
        {"metric_name": metric_types.MetricType.LATENCY, "value": latency},
        {"metric_name": metric_types.MetricType.TFLOPS, "value": tflops},
        {"metric_name": metric_types.MetricType.MFU, "value": mfu},
     # Custom metrics
     {"custom_metrics_1", "value":<value>},
     {"custom_metrics_2", "value":<value>},
     {"avg_mtp_acceptance_rate_percent", "value":<value>},
     {"dpo_reward_accuracy", "value":<value>},
    ], step=step+1)

לכידת פרופילים

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

פקודות ה-SDK ללכידת פרופילים לא תלויות במסגרת, כי כל פקודות הפרופילים ברמת המסגרת משולבות אוטומטית בפקודות הפרופילים של ML Diagnostics. המשמעות היא שקוד הפרופיל לא תלוי במסגרת שבה אתם משתמשים.

לכידת פרופילים פרוגרמטית

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

אפשר לבצע לכידה פרוגרמטית של פרופילים באמצעות ML Diagnostics SDK בדרכים הבאות:

• איסוף מבוסס API: שליטה בפרופילים באמצעות השיטות start() ו-stop().
• איסוף מבוסס-Decorator: הוספת הערות לפונקציות באמצעות @xprof(run) ליצירת פרופיל אוטומטית.
• מנהל הקשר: לשימוש עם xprof() ליצירת פרופילים מבוססי-היקף שמטפלים אוטומטית בפעולות start() ו-stop().

אפשר להשתמש באותו קוד לתיעוד פרופיל בכל המסגרות. כל הסשנים של הפרופיל נשמרים בקטגוריה של Cloud Storage שמוגדרת בהרצת למידת המכונה.

# Support collection via APIs
prof = xprof()  # Updates metadata and starts xprofz collector
prof.start()  # Collects traces to bucket
# ..... Your code execution here
# ....
prof.stop()

# Also supports collection via decorators
@xprof()
def abc(self):
    # does something
    pass

# Use xprof as a context manager to automatically start and stop collection
with xprof() as prof:
    # Your training or execution code here
    train_model()
    evaluate_model()

יצירת פרופילים (תהליך) של אתרים עם כמה מארחים

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

# starts profiling on all nodes
prof = xprof()
prof.start()
# ...
prof.stop()

כברירת מחדל, קריאה לשיטה prof.start() ללא הארגומנט session_id במספר מארחים מובילה לסשנים נפרדים של מעקב – אחד לכל מארח. כדי לקבץ עקבות ממארחים שונים לסשן מאוחד יחיד של כמה מארחים ב-XProf, צריך לוודא שהמתודה prof.start() נקראת עם אותו ארגומנט session_id בכל המארחים המשתתפים. לדוגמה:

# Use the same session_id on all hosts to group traces
prof = xprof()
prof.start(session_id="profiling_session")
# ...
prof.stop()

כדי להפעיל את יצירת הפרופיל למארחים ספציפיים:

# starts profiling on node with index 0 and 2
prof = xprof(process_index_list=[0,2])
prof.start()
# ...
prof.stop()

לכידת פרופיל לפי דרישה

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

מדדי המודל (שמופקים מ-SDK) לא זמינים כשמשתמשים בפרופיל לפי דרישה. מדדי מערכת TPU עדיין זמינים בלוחות הבקרה של GKE TPU.

כדי להפעיל את האפשרות ליצירת פרופיל לפי דרישה, צריך להגדיר את ההרצה עם תמיכה לפי דרישה:

# Define machinelearning run
machinelearning_run(
    name="<run_name>",
    # specify where profiling data is stored
    gcs_path="gs://<bucket>",
    ...
    # enable on demand profiling, starts xprofz daemon on port 9999
    on_demand_xprof=True
)

אפשר להשתמש באותו קוד לתיעוד פרופיל בכל המסגרות. כל הסשנים של הפרופיל נשמרים בקטגוריה של Cloud Storage שמוגדרת בהרצת למידת המכונה.

כדי ליצור פרופיל לפי דרישה ב-GKE, פורסים את GKE connection-operator ואת injection-webhook באשכול GKE. כך אפשר לוודא שהרצת הלמידה החישובית תוכל לאתר את צמתי GKE שהיא פועלת בהם, ושהתפריט הנפתח של הלכידה לפי דרישה יוכל לאכלס אוטומטית את הצמתים האלה. מידע נוסף זמין במאמר בנושא הגדרת אשכול GKE.

אריזת עומס עבודה ל-GKE

אפשר להשתמש בקובץ Docker כדי לארוז אפליקציה שמשתמשת ב-SDK של ML Diagnostics. מתקינים את חבילת google-cloud-logging לשילוב עם Cloud Logging. לדוגמה:

# Base image (user's choice, e.g., python:3.10-slim, or a base with ML frameworks)
FROM python:3.11-slim

# Install base utilities
RUN pip install --no-cache-dir --upgrade pip

# Install SDK and Logging client
# psutil is installed as a dependency of google-cloud-mldiagnostics
RUN pip install --no-cache-dir \
    google-cloud-mldiagnostics \
    google-cloud-logging

# Optional: For JAX/TPU workloads
# RUN pip install --no-cache-dir "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html &&
#     pip install --no-cache-dir libtpu xprof

# Add your application code
COPY ./app /app
WORKDIR /app

# Run your script
CMD ["python", "your_train_script.py"]

פריסת עומס עבודה

אחרי שמשלבים את ה-SDK עם עומס העבודה, אורזים את עומס העבודה בקובץ אימג' ויוצרים את קובץ ה-YAML עם קובץ האימג' שצוין. לצורך יצירת פרופיל באופן פרוגרמטי, צריך להוסיף את התווית managed-mldiagnostics-gke=true לעומס העבודה בקובץ ה-YAML. לא צריך להוסיף את התווית הזו כדי ליצור פרופיל לפי דרישה.

ב-GKE:

kubectl apply -f YAML_FILE_NAME

ב-Compute Engine, מתחברים למכונה הווירטואלית באמצעות SSH ומריצים את קוד Python של עומס העבודה:

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

אחרי פריסת עומס העבודה, מחפשים את שם העבודה במרחב השמות של עומס העבודה:

kubectl get job -n YOUR_NAMESPACE

אפשר למצוא את שם ההרצה והקישור ביומני kubectl על ידי העברת שם המשימה ומרחב השמות. צריך לציין את מאגר עומסי העבודה (לדוגמה: -c workload) כי קובץ ה-sidecar של ML Diagnostics מטפל ברישום ביומן משלו.

kubectl logs jobs/s5-tpu-slice-0 -n YOUR_NAMESPACE -c workload