Erste Schritte mit dem ML Diagnostics SDK

Das ML Diagnostics Python SDK kann in ML-Arbeitslasten eingebunden werden, um Modellmesswerte, ‑konfigurationen und ‑profile auf Google Cloudzu erfassen und zu verwalten. Verwenden Sie das SDK für die programmatische Profilerstellung und um Modellmesswerte anzusehen. Modellmesswerte sind nicht verfügbar, wenn Sie die On-Demand-Profilerstellung ohne das SDK verwenden.

In diesem Leitfaden erfahren Sie, wie Sie Machine-Learning-Ausführungen erstellen, Arbeitslastmesswerte und ‑konfigurationen erfassen und verwalten, verwaltete XProf-Ressourcen bereitstellen und die programmatische und On-Demand-Profilerstellung aktivieren.

Weitere Informationen zur Verwendung des ML Diagnostics SDK finden Sie im Repository google-cloud-mldiagnostics.

ML Diagnostics SDK installieren

Installieren Sie die google-cloud-mldiagnostics Bibliothek:

pip install google-cloud-mldiagnostics

Importieren Sie die folgenden Pakete in Ihren ML-Arbeitslastcode:

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

Cloud Logging aktivieren

Das SDK verwendet das Standardmodul logging von Python, um Messwerte und Konfigurationsinformationen auszugeben. Wenn Sie diese Logs an Cloud Logging weiterleiten möchten, installieren und konfigurieren Sie die Bibliothek google-cloud-logging. So können Sie SDK-Logs, protokollierte Messwerte und Ihre eigenen Anwendungsprotokolle in der Google Cloud Console ansehen.

Installieren Sie die Bibliothek google-cloud-logging:

pip install google-cloud-logging

Konfigurieren Sie das Logging in Ihrem Skript, indem Sie den Cloud Logging-Handler an den Python-Root-Logger anhängen. Fügen Sie die folgenden Zeilen am Anfang Ihres Python-Skripts ein:

  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.")

Detailliertes Logging aktivieren

Die Logging-Ebene ist standardmäßig auf INFO festgelegt. Wenn Sie detailliertere Logs vom SDK erhalten möchten, z. B. Details zur Machine-Learning-Ausführung, legen Sie die Logging-Ebene nach dem Aufruf von setup_logging() auf DEBUG fest:

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.")

Wenn DEBUG aktiviert ist, erhalten Sie zusätzliche SDK-Diagnosedaten in Cloud Logging. Beispiel:

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', ...}

Machine-Learning-Ausführung erstellen

Wenn Sie die ML Diagnostics-Plattform verwenden möchten, müssen Sie zuerst eine Machine-Learning-Ausführung erstellen. Dazu müssen Sie Ihre ML-Arbeitslast mit dem SDK instrumentieren, um Logging durchzuführen, Messwerte zu erfassen und die Profilverfolgung zu aktivieren.

Das folgende Beispiel initialisiert Cloud Logging, erstellt eine Machine-Learning-Ausführung (MLRun), erfasst Messwerte und erfasst ein Profil:

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)

Im Codebeispiel werden die folgenden Variablen verwendet:

Variable Anforderung Beschreibung
name Erforderlich Eine Kennung für die jeweilige Ausführung. Das SDK erstellt automatisch eine machine-learning-run-id um sicherzustellen, dass die Ausführungs namen eindeutig sind.
run_group Optional Eine Kennung, mit der mehrere Ausführungen derselben Testreihe gruppiert werden können. Beispielsweise könnten alle Ausführungen, die mit einer TPU-Slice-Größenabstimmung verknüpft sind, derselben Gruppe angehören.
project Optional Wenn nicht angegeben, wird das Projekt aus der Google Cloud CLI extrahiert.
region Erforderlich Alle Cluster Director-Standorte werden unterstützt, außer us-east5. Dieses Flag kann für jeden Befehl mit einem Argument oder mit dem Befehl gcloud config set compute/region festgelegt werden.
configs Optional Schlüssel-Wert-Paare mit Konfigurationsparametern für die Ausführung. Wenn Konfigurationen nicht definiert sind, werden Standardkonfigurationen für Software und System angezeigt, aber die Konfigurationen für die ML-Arbeitslast nicht.
gcs_path Bedingt erforderlich Der Google Cloud Storage-Speicherort, an dem alle Profile gespeichert werden. Beispiel: gs://my-bucket oder gs://my-bucket/folder1. Nur erforderlich, wenn das SDK für die Profilerstellung verwendet wird.
on-demand-xprof Optional Startet xprofz daemon auf Port 9999, um die On-Demand Profilerstellung zu aktivieren. Sie können sowohl die On-Demand-Profilerstellung als auch die programmatische Profilerstellung im selben Code aktivieren, solange sie nicht gleichzeitig ausgeführt werden.

Die folgenden Konfigurationen werden automatisch vom SDK erfasst und müssen nicht in machinelearning_run angegeben werden:

  • Softwarekonfigurationen: Framework, Framework-Version, XLA-Flags.
  • Systemkonfigurationen: Gerätetyp, Anzahl der Slices, Slice-Größe, Anzahl der Hosts.

Projekt- und Regionsinformationen werden als Metadaten für die Machine-Learning-Ausführung gespeichert. Die für die Machine-Learning-Ausführung verwendete Region muss nicht mit der für die Arbeitslastausführung verwendeten Region übereinstimmen.

Konfigurationen schreiben

Viele Arbeitslasten enthalten zu viele Konfigurationen, um sie direkt in der Definition machinelearning_run zu definieren. In diesen Fällen können Sie Konfigurationen für Ihre Ausführung mit JSON oder YAML schreiben.

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",
)

Messwerte erfassen

Sie können mit dem SDK Modellmesswerte, Messwerte zur Modellleistung und Systemmesswerte erfassen. Sie können Visualisierungen dieser Messwerte als Durchschnittswerte und mit Zeitreihendiagrammen erstellen.

Das SDK bietet zwei Funktionen zum Erfassen von Messwerten: metrics.record() zum Erfassen einzelner Datenpunkte und metrics.record_metrics() zum Erfassen mehrerer Messwerte in einem Batch. Beide Funktionen schreiben Messwerte in Cloud Logging, sodass sie visualisiert und analysiert werden können.

So erfassen Sie einen einzelnen Messwert:

# 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)

So erfassen Sie mehrere Messwerte:

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)

Die folgenden Systemmesswerte werden automatisch vom SDK aus den Bibliotheken libTPU, psutil und JAX erfasst:

  • TPU-TensorCore-Auslastung
  • TPU-Arbeitszyklus
  • HBM-Auslastung
  • Host-CPU-Auslastung
  • Hostarbeitsspeicherauslastung

Sie müssen diese Messwerte nicht manuell angeben. Für diese Systemmesswerte ist time die Standard-X-Achse.

Die folgenden vordefinierten Messwertschlüssel werden automatisch in der Google Cloud Console angezeigt, wenn sie zugewiesen sind. Diese Messwerte werden nicht automatisch berechnet. Es handelt sich um vordefinierte Schlüssel, denen Sie Werte zuweisen können.

  • Messwertschlüssel für die Modellqualität: LEARNING_RATE, LOSS, GRADIENT_NORM, TOTAL_WEIGHTS.
  • Leistungsmesswertschlüssel für das Modell: STEP_TIME, THROUGHPUT, LATENCY, MFU, TFLOPS.

Die vordefinierten Messwerte sowie andere benutzerdefinierte Messwerte können mit time oder sowohl time als auch step als X-Achse erfasst werden. Sie können jeden benutzerdefinierten Messwert in der Arbeitslast erfassen.

Im folgenden Beispiel wird ein einzelner Messwert für die Arbeitslast erfasst, den Sie auf dem Tab Model Metrics (Modellmesswerte) für die jeweilige Machine-Learning-Ausführung ansehen können:

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

Wenn Sie mehrere Messwerte in einem Aufruf erfassen möchten, verwenden Sie die Methode record_metrics. Beispiel:

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)

Profile erfassen

Sie können XProf-Profile Ihrer ML-Arbeitslast mit der programmatischen oder der On-Demand-Profilerstellung (manuelle Profilerstellung) erfassen. Bei der programmatischen Profilerstellung werden Profilerstellungsbefehle direkt in Ihren Machine-Learning-Code eingebettet und Sie geben explizit an, wann die Aufzeichnung von Daten gestartet und beendet werden soll. Die On-Demand-Profilerstellung erfolgt in Echtzeit. Sie lösen den Profiler aus, während die Arbeitslast bereits aktiv ausgeführt wird.

Die SDK-Befehle zum Erfassen von Profilen sind Framework-unabhängig, da alle Profilerstellungsbefehle auf Framework-Ebene automatisch in die Profilerstellungsbefehle von ML Diagnostics integriert werden. Das bedeutet, dass Ihr Profilerstellungscode nicht vom verwendeten Framework abhängig ist.

Programmatische Profilerstellung

Bei der programmatischen Profilerstellung müssen Sie Ihren Modellcode mit Anmerkungen versehen und angeben, wo Profile erfasst werden sollen. In der Regel erfassen Sie ein Profil für einige Trainingsschritte oder für einen bestimmten Codeblock in Ihrem Modell.

Sie können die programmatische Profilerstellung mit dem ML Diagnostics SDK auf folgende Arten durchführen:

  • API-basierte Erfassung: Steuern Sie die Profilerstellung mit start() und stop() Methoden.
  • Decorator-basierte Erfassung: Fügen Sie Funktionen mit @xprof(run) hinzu, um die automatische Profilerstellung zu aktivieren.
  • Kontextmanager: Verwenden Sie mit xprof() für die bereichsbasierte Profilerstellung, die die Vorgänge start() und stop() automatisch verarbeitet.

Sie können denselben Code für die Profilerstellung in allen Frameworks verwenden. Alle Profilerstellungssitzungen werden in dem Cloud Storage-Bucket erfasst, das in der Machine-Learning-Ausführung definiert ist.

# 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()

Profilerstellung für mehrere Hosts (Prozesse)

Bei der programmatischen Profilerstellung startet das SDK die Profilerstellung auf jedem Host (Prozess), auf dem ML-Arbeitslastcode ausgeführt wird. Wenn die Liste der Knoten nicht angegeben wird, werden alle Hosts berücksichtigt.

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

Standardmäßig führt der Aufruf der Methode prof.start() ohne das Argument session_id auf mehreren Hosts zu separaten Trace-Sitzungen, eine für jeden Host. Wenn Sie Traces von verschiedenen Hosts in einer einzelnen, einheitlichen Sitzung für mehrere Hosts in XProf gruppieren möchten, muss die Methode prof.start() auf allen beteiligten Hosts mit demselben Argument session_id aufgerufen werden. Beispiel:

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

So aktivieren Sie die Profilerstellung für bestimmte Hosts:

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

On-Demand-Profilerstellung

Verwenden Sie die On-Demand-Profilerstellung, wenn Sie Profile ad hoc erfassen möchten oder wenn die programmatische Profilerstellung noch nicht aktiviert ist. Die On-Demand-Profilerstellung ist hilfreich, wenn während der Ausführung Probleme mit der Modellleistung auftreten und Sie in diesen Momenten Profile erfassen möchten, um die Probleme zu diagnostizieren.

Modellmesswerte (die vom SDK ausgegeben werden) sind bei der On-Demand-Profilerstellung nicht verfügbar. TPU-Systemmesswerte sind weiterhin in GKE-TPU-Dashboards verfügbar.

Wenn Sie die On-Demand-Profilerstellung aktivieren möchten, konfigurieren Sie die Ausführung mit On-Demand-Unterstützung:

# 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
)

Sie können denselben Code für die Profilerstellung in allen Frameworks verwenden. Alle Profilerstellungssitzungen werden in dem Cloud Storage-Bucket erfasst, das in der Machine-Learning-Ausführung definiert ist.

Für die On-Demand-Profilerstellung in GKE stellen Sie connection-operator und injection-webhook von GKE im GKE-Cluster bereit. So kann Ihre Machine-Learning-Ausführung die GKE-Knoten finden, auf denen sie ausgeführt wird, und die Drop-down-Liste für die On-Demand-Profilerstellung kann diese Knoten automatisch ausfüllen. Weitere Informationen finden Sie unter GKE-Cluster konfigurieren.

Arbeitslast für GKE verpacken

Sie können ein Dockerfile verwenden, um eine Anwendung zu verpacken, die das ML Diagnostics SDK verwendet. Installieren Sie das Paket google-cloud-logging für die Cloud Logging-Integration. Beispiel:

# 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"]

Arbeitslast bereitstellen

Nachdem Sie das SDK in Ihre Arbeitslast eingebunden haben, verpacken Sie die Arbeitslast in einem Image und erstellen Sie Ihre YAML-Datei mit dem angegebenen Image. Für die programmatische Profilerstellung kennzeichnen Sie die Arbeitslast in der YAML-Datei mit managed-mldiagnostics-gke=true. Für die On-Demand-Profilerstellung ist dieses Label nicht erforderlich.

Für GKE:

kubectl apply -f YAML_FILE_NAME

Für Compute Engine stellen Sie eine SSH-Verbindung zur VM her und führen Sie den Python-Code für Ihre Arbeitslast aus:

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

Nachdem Sie die Arbeitslast bereitgestellt haben, suchen Sie nach dem Namen Ihres Jobs, indem Sie nach dem Namespace Ihrer Arbeitslast suchen:

kubectl get job -n YOUR_NAMESPACE

Sie finden den Namen und den Link der Ausführung in Ihren kubectl-Logs, indem Sie den Jobnamen und den Namespace übergeben. Sie müssen den Arbeitslastcontainer angeben (z. B. -c workload), da der ML Diagnostics-Sidecar sein eigenes Logging verwaltet.

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