Erste Schritte mit dem ML Diagnostics SDK

Das ML Diagnostics Python SDK kann in ML-Arbeitslasten eingebunden werden, um Arbeitslastmesswerte, ‑konfigurationen und ‑profile auf Google Cloudzu erfassen und zu verwalten. In diesem Leitfaden erfahren Sie, wie Sie Machine-Learning-Läufe 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 den Code Ihrer ML-Arbeitslast:

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 Standard-Python-Modul logging, um Messwerte und Konfigurationsinformationen auszugeben. Wenn Sie diese Logs an Cloud Logging weiterleiten möchten, installieren und konfigurieren Sie die google-cloud-logging-Bibliothek. So können Sie SDK-Logs, protokollierte Messwerte und Ihre eigenen Anwendungslogs in der Google Cloud Console ansehen.

Installieren Sie die google-cloud-logging-Bibliothek:

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 am Anfang Ihres Python-Skripts die folgenden Zeilen 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 zu Machine-Learning-Ausführungen, legen Sie die Logging-Ebene nach dem Aufrufen 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 einen Machine-Learning-Lauf erstellen. Dazu müssen Sie Ihren ML-Arbeitslast mit dem SDK instrumentieren, um Logging durchzuführen, Messwerte zu erfassen und Profil-Tracing zu aktivieren.

Das folgende Beispiel initialisiert Cloud Logging, erstellt einen Machine-Learning-Lauf (MLRun), erfasst Messwerte und erstellt 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 den jeweiligen Lauf. Das SDK erstellt automatisch eine machine-learning-run-id, um sicherzustellen, dass die Namen der Läufe eindeutig sind.
run_group	Optional	Eine Kennung, mit der mehrere Läufe derselben Testgruppe zusammengefasst werden können. Beispielsweise können alle Läufe, die mit einem Sweep der TPU-Slice-Größe verknüpft sind, zur selben Gruppe gehören.
project	Optional	Wenn nichts angegeben ist, wird das Projekt aus der Google Cloud CLI extrahiert.
region	Erforderlich	Alle Cluster Director-Standorte werden unterstützt, mit Ausnahme von 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 den Lauf. Wenn keine Konfigurationen definiert sind, werden Standardkonfigurationen für Software und System angezeigt, aber nicht die Konfigurationen für ML-Arbeitslasten.
gcs_path	Bedingt erforderlich	Der Google Cloud Speicherort, an dem alle Profile gespeichert werden. Beispiel: gs://my-bucket oder gs://my-bucket/folder1 Nur erforderlich, wenn das SDK zum Erfassen von Profilen verwendet wird.
on-demand-xprof	Optional	Startet xprofz daemon auf Port 9999, um die On-Demand-Profilerstellung zu aktivieren. Sie können sowohl On-Demand- als auch programmatische Profilerstellung im selben Code aktivieren, sofern sie nicht gleichzeitig erfolgen.

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 Machine-Learning-Ausführungen gespeichert. Die Region, die für den Lauf für maschinelles Lernen verwendet wird, muss nicht mit der Region übereinstimmen, die für den Arbeitslastlauf verwendet wird.

Konfigurationen schreiben

Viele Arbeitslasten enthalten zu viele Konfigurationen, um sie direkt in der machinelearning_run-Definition zu definieren. In diesen Fällen können Sie Konfigurationen für Ihren Lauf in 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

Mit dem SDK können Sie Modellmesswerte, Modellleistungsmesswerte und Systemmesswerte erfassen. Sie können diese Messwerte als Durchschnittswerte und mit Zeitreihendiagrammen visualisieren.

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 einzelnen Batch. Beide Funktionen schreiben Messwerte in Cloud Logging, was die Visualisierung und Analyse ermöglicht.

So zeichnen Sie einen einzelnen Messwert auf:

# 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 zeichnen Sie mehrere Messwerte auf:

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 Zeit die Standard-X-Achse.

Die folgenden vordefinierten Messwertschlüssel werden automatisch in derGoogle Cloud Console angezeigt, sofern sie zugewiesen sind. Diese Messwerte werden nicht automatisch berechnet, sondern sind vordefinierte Schlüssel, denen Sie Werte zuweisen können.

• Schlüssel für Messwerte zur Modellqualität: LEARNING_RATE, LOSS, GRADIENT_NORM, TOTAL_WEIGHTS.
• Schlüssel für Messwerte zur Modellleistung: STEP_TIME, THROUGHPUT, LATENCY, MFU, TFLOPS.

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

Im folgenden Beispiel wird ein einzelner Messwert für die Arbeitslast erfasst, den Sie auf dem Tab Modellmesswerte für den jeweiligen Machine-Learning-Lauf 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)

Aufnahmeprofile

Sie können XProf-Profile Ihrer ML-Arbeitslast mit programmatischer oder On-Demand-Erfassung (manuelle Erfassung) erfassen. Bei der programmatischen Erfassung werden Profiling-Befehle 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-Erfassung 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 frameworkunabhängig, da alle Profiling-Befehle auf Framework-Ebene automatisch in die Profiling-Befehle von ML Diagnostics integriert werden. Das bedeutet, dass Ihr Profiling-Code nicht vom verwendeten Framework abhängt.

Programmatische Profilerfassung

Bei der programmatischen Erfassung müssen Sie Ihren Modellcode mit Anmerkungen versehen und angeben, an welcher Stelle 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 das Profil programmatisch mit dem ML Diagnostics SDK auf folgende Arten erfassen:

• API-basierte Erhebung: Profilerstellung mit den Methoden start() und stop() steuern.
• Dekoratorbasierte Erfassung: Funktionen mit @xprof(run) für die automatische Profilerstellung annotieren.
• Context Manager: Wird mit xprof() für die bereichsbasierte Profilerstellung verwendet, bei der start()- und stop()-Vorgänge automatisch verarbeitet werden.

Sie können denselben Profilerfassungscode für alle Frameworks verwenden. Alle Profilsitzungen werden im Cloud Storage-Bucket erfasst, der im Lauf des maschinellen Lernens 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()

Profiling von mehreren Hosts (Prozessen)

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

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

Wenn Sie die Methode prof.start() standardmäßig ohne das Argument session_id auf mehreren Hosts aufrufen, werden separate Tracing-Sitzungen erstellt, eine für jeden Host. Wenn Sie Traces von verschiedenen Hosts in einer einzigen, einheitlichen Multi-Host-Sitzung in XProf gruppieren möchten, muss die Methode prof.start() mit dem gleichen session_id-Argument auf allen beteiligten Hosts 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-Profilerfassung

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

Wenn Sie die On-Demand-Profilerstellung aktivieren möchten, konfigurieren Sie den Lauf 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 Profilerfassungscode für alle Frameworks verwenden. Alle Profilsitzungen werden im Cloud Storage-Bucket erfasst, der im Machine-Learning-Lauf definiert ist.

Für das On-Demand-Profiling in GKE müssen Sie connection-operator und injection-webhook im GKE-Cluster bereitstellen. So kann Ihr Machine-Learning-Lauf die GKE-Knoten finden, auf denen er ausgeführt wird, und das Drop-down-Menü für die On-Demand-Erfassung 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 integriert haben, verpacken Sie die Arbeitslast in einem Image und erstellen Sie Ihre YAML-Datei mit dem angegebenen Image. Kennzeichnen Sie die Arbeitslast in der YAML-Datei mit managed-mldiagnostics-gke=true.

Für GKE:

kubectl apply -f YAML_FILE_NAME

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

source venv/bin/activate
python3.11 WORKLOAD_FILE_NAME

Nachdem Sie die Arbeitslast bereitgestellt haben, suchen Sie nach dem Namespace der Arbeitslast, um den Jobnamen zu finden:

kubectl get job -n YOUR_NAMESPACE

Sie finden den Namen und den Link des Laufs 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 übernimmt.

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