google-cloud-mldiagnostics 0.1.0__py3-none-any.whl

google_cloud_mldiagnostics/__init__.py

@@ -0,0 +1,13 @@
+"""Init module for ml diagnostics."""
+
+# Import the upfront API for users
+from . import api
+from . import core
+
+
+__path__ = __import__('pkgutil').extend_path(__path__, __name__)
+
+
+machinelearning_run = api.machinelearning_run
+metrics = api.metrics
+xprof = core.xprof.Xprof

google_cloud_mldiagnostics/api/__init__.py

@@ -0,0 +1,7 @@
+"""API for interacting with the diagnostics library."""
+
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.api import metrics
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.api import mlrun
+
+
+machinelearning_run = mlrun.machinelearning_run

google_cloud_mldiagnostics/api/metrics.py

@@ -0,0 +1,38 @@
+"""Module for recording metrics."""
+
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.core import metrics
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.custom_types import metric_types
+
+
+_metrics_recorder = metrics.metrics_recorder
+
+
+def record(
+    metric_name: metric_types.MetricType or str,
+    value: int | float,
+    step: int | None = None,
+    labels: dict[str, str] | None = None,
+    record_on_all_hosts: bool = False,
+) -> None:
+  """Record a single metric value using the active run.
+
+  Args:
+      metric_name: Name of metric to record.
+      value: Metric value.
+      step: Optional step number (auto-incremented if not provided).
+      labels: Optional additional labels.
+      record_on_all_hosts: Whether to record metrics on all hosts.
+
+  Raises:
+      RecordingError: If no active run or recording fails.
+
+  Example:
+      metrics.record(MetricType.TF_FLOPS, per_device_tf_flops)
+      metrics.record(MetricType.LEARNING_RATE, learning_rate)
+      metrics.record(MetricType.LEARNING_RATE, learning_rate, step=1)
+  """
+  is_enum = isinstance(metric_name, metric_types.MetricType)
+  metric_name = metric_name.value if is_enum else metric_name
+  _metrics_recorder.record(
+      metric_name, value, step, labels, record_on_all_hosts
+  )

google_cloud_mldiagnostics/api/mlrun.py

@@ -0,0 +1,71 @@
+"""Module for creating and managing ML runs."""
+from typing import Any
+
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.core import create_mlrun
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.custom_types import exceptions
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.custom_types import mlrun_types
+
+
+# Main SDK function - this is the primary interface users will import
+def machinelearning_run(
+    run_group: str,
+    name: str | None = None,
+    configs: dict[str, Any] | None = None,
+    gcs_path: str | None = None,
+    gcp_project: str | None = None,
+    gcp_region: str | None = None,
+    metrics_record_interval_sec: float = 10.0,
+) -> mlrun_types.MLRun:
+  """Create a new machine learning run.
+
+  This is the main entry point for the SDK that users will call to create ML
+  runs.
+
+  Args:
+      run_group: The run set this run belongs to
+      name: The name of the run (optional), system will generate a unique ID if
+        not provided by users.
+      configs: dict of configuration parameters
+      gcs_path: GCS path for storing run artifacts
+      gcp_project: The GCP project ID
+      gcp_region: The GCP region
+      metrics_record_interval_sec: The interval in seconds
+        for recording system metrics backend (tpu duty cycle,
+        tpu tensorcore utilization, hbm utilization,
+        host cpu utilization, host memory utilization).
+
+  Returns:
+      MLRun: A new ML run instance
+
+  Example:
+      from google_cloud_mldiagnostics import machinelearning_run
+
+      my_run = machinelearning_run(
+            run_group="training_set_v1",
+            name="experiment_1",
+            configs={"epochs": 100, "batch_size": 32},
+            gcs_path="gs://my-bucket/experiments"
+        )
+
+      # Update configs using dict methods
+      my_run.configs.update({"epochs": 300, "optimizer": "adam"})
+
+      # Update configs using attribute notation
+      my_run.configs.batch_size = 64
+  """
+  if not run_group:
+    raise exceptions.MLRunConfigurationError(
+        "run_group is required and must be provided. The run_group parameter"
+        " helps identify and organize your ML workloads. Please provide a"
+        " meaningful run_group name (e.g., 'training_v1',"
+        " 'experiment_batch_1')."
+    )
+  return create_mlrun.initialize_mlrun(
+      run_group=run_group,
+      name=name,
+      configs=configs,
+      gcs_path=gcs_path,
+      gcp_project=gcp_project,
+      gcp_region=gcp_region,
+      metrics_record_interval_sec=metrics_record_interval_sec,
+  )

google_cloud_mldiagnostics/api/test_metrics.py

@@ -0,0 +1,32 @@
+"""Unit tests for the API metrics module."""
+
+import unittest
+from unittest import mock
+
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.api import metrics
+
+
+class TestMetricsAPI(unittest.TestCase):
+  """Test cases for the metrics API module."""
+
+  @mock.patch(
+      "google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.api."
+      "metrics._metrics_recorder"
+  )
+  def test_record_function_minimal_parameters(self, mock_recorder):
+    """Test record function with minimal parameters."""
+    # Setup
+    mock_metric_type = mock.Mock()
+    mock_metric_type.value = "test_metric"
+
+    # Call the API function
+    metrics.record(metric_name=mock_metric_type, value=42.5)
+
+    # Verify delegation to the underlying recorder
+    mock_recorder.record.assert_called_once_with(
+        mock_metric_type, 42.5, None, None, False
+    )
+
+
+if __name__ == "__main__":
+  unittest.main()

google_cloud_mldiagnostics/api/test_mlrun.py

@@ -0,0 +1,97 @@
+"""Unit tests for the machinelearning_run module."""
+
+import json
+import unittest
+from unittest import mock
+
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.api import mlrun as api_mlrun
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.custom_types import exceptions
+
+
+class TestMachineLearningRun(unittest.TestCase):
+  """Test cases for the machinelearning_run function."""
+
+  def setUp(self):
+    """Set up test fixtures before each test method."""
+    super().setUp()
+    self.valid_run_group = "test_run_group"
+    self.valid_name = "test_run_name"
+    self.valid_configs = {"epochs": 100, "batch_size": 32}
+    self.valid_gcs_path = "gs://test-bucket/path"
+    self.mock_create_mlrun = mock.patch(
+        "google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.core.create_mlrun.initialize_mlrun"
+    ).start()
+
+  def tearDown(self):
+    """Clean up after each test method."""
+    mock.patch.stopall()
+    super().tearDown()
+
+  def test_machinelearning_run_with_all_parameters(self):
+    """Test creating ML run with all parameters provided."""
+    api_mlrun.machinelearning_run(
+        run_group=self.valid_run_group,
+        name=self.valid_name,
+        configs=self.valid_configs,
+        gcs_path=self.valid_gcs_path
+    )
+
+    self.mock_create_mlrun.assert_called_once_with(
+        run_group=self.valid_run_group,
+        name=self.valid_name,
+        configs=self.valid_configs,
+        gcs_path=self.valid_gcs_path,
+        gcp_project=None,
+        gcp_region=None,
+        metrics_record_interval_sec=10.0,
+    )
+
+  def test_machinelearning_run_none_run_group_raises_error(self):
+    """Test that None run_group raises MLRunConfigurationError."""
+    with self.assertRaises(exceptions.MLRunConfigurationError) as context:
+      api_mlrun.machinelearning_run(run_group=None)
+
+    self.assertIn(
+        "run_group is required and must be provided. The run_group parameter"
+        " helps identify and organize your ML workloads. Please provide a"
+        " meaningful run_group name (e.g., 'training_v1',"
+        " 'experiment_batch_1').",
+        str(context.exception),
+    )
+    self.assertIn("identify and organize", str(context.exception))
+
+  def test_yaml_json_workflow_configs(self):
+    """Test ML run creation with configs loaded from a JSON representation of YAML."""
+    self.valid_run_group = "test_run_group"
+    self.valid_name = "test_run_name"
+    self.valid_configs = {"epochs": 100, "batch_size": 32}
+    self.valid_gcs_path = "gs://test-bucket/path"
+
+    # Simulate YAML data that was loaded and converted to JSON
+    yaml_data = {
+        "epochs": 200,
+        "batch_size": 64,
+        "learning_rate": 0.001,
+        "model_type": "transformer"
+    }
+    json_data = json.dumps(yaml_data)
+
+    api_mlrun.machinelearning_run(
+        run_group=self.valid_run_group,
+        name=self.valid_name,
+        configs=json.loads(json_data),  # Parse JSON back to dict for this test
+        gcs_path=self.valid_gcs_path
+    )
+
+    self.mock_create_mlrun.assert_called_once_with(
+        run_group=self.valid_run_group,
+        name=self.valid_name,
+        configs=yaml_data,  # Verify that the YAML data was correctly converted
+        gcs_path=self.valid_gcs_path,
+        gcp_project=None,
+        gcp_region=None,
+        metrics_record_interval_sec=10.0,
+    )
+
+if __name__ == "__main__":
+  unittest.main()

google_cloud_mldiagnostics/clients/control_plane_client.py

@@ -0,0 +1,127 @@
+"""Client for sending requests to Diagon Control Plane."""
+
+import re
+from typing import Any, Dict, List, Optional
+
+import google.auth
+from google.auth.transport import requests as google_auth_requests
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.utils import host_utils
+import requests
+
+
+class ControlPlaneClient:
+  """Client for communicating with Google Cloud Hypercompute Cluster ML Run service."""
+
+  def __init__(
+      self,
+      project_id: str = "supercomputer-testing",
+      location: str = "us-central1",
+      base_url: str = "https://autopush-hypercomputecluster.sandbox.googleapis.com/v1alpha",
+  ):
+    """Initializes a new ControlPlaneClient.
+
+    Args:
+        project_id: Google Cloud project ID
+        location: Google Cloud location/region
+        base_url: Base URL for the API endpoint
+    """
+    self.project_id = project_id
+    self.location = location
+    self.base_url = base_url
+    self.ml_runs_url = f"{base_url}/projects/{project_id}/locations/{location}/machineLearningRuns"
+
+    # Initialize Google Cloud credentials
+    self.credentials, _ = google.auth.default()
+
+  def _get_access_token(self) -> str:
+    """Get Google Cloud access token for authentication."""
+    if not self.credentials.valid:
+      self.credentials.refresh(google_auth_requests.Request())
+
+    return self.credentials.token
+
+  def _get_headers(self) -> Dict[str, str]:
+    """Get HTTP headers with authentication."""
+    return {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {self._get_access_token()}",
+    }
+
+  def create_ml_run(
+      self,
+      name: str,
+      display_name: str,
+      run_phase: str,
+      configs: Optional[Dict[str, Any]] = None,
+      tools: Optional[List[Dict[str, Any]]] = None,
+      metrics: Optional[Dict[str, str]] = None,
+      artifacts: Optional[Dict[str, str]] = None,
+      run_group: Optional[str] = None,
+      labels: Optional[Dict[str, str]] = None,
+  ) -> Dict[str, Any]:
+    """Create a new ML run using the Google Cloud API.
+
+    Args:
+        name: Name of the run
+        display_name: Display name for the run
+        run_phase: Phase of the run (ACTIVE, COMPLETE, FAILED)
+        configs: Configuration settings (userConfigs, softwareConfigs,
+          hardwareConfigs)
+        tools: List of tools to enable (e.g., XProf, NSys)
+        metrics: Metrics for the run (e.g., avgStep, avgLatency)
+        artifacts: Artifacts configuration (e.g., gcsPath)
+        run_group: Run group grouping identifier
+        labels: Custom labels for the run
+
+    Returns:
+        Response from the API as a dictionary
+
+    Raises:
+        requests.exceptions.RequestException: If the HTTP request fails
+    """
+    payload = {"displayName": display_name, "name": name}
+
+    if configs:
+      payload["configs"] = configs
+
+    if metrics:
+      payload["metrics"] = metrics
+
+    if artifacts:
+      payload["artifacts"] = artifacts
+
+    if run_group:
+      payload["runSet"] = run_group
+
+    if labels:
+      payload["labels"] = labels
+
+    if run_phase:
+      payload["runPhase"] = run_phase
+
+    if tools:
+      payload["tools"] = tools
+
+    # Sanitize the name for machineLearningRunId
+    sanitized_name = host_utils.sanitize_identifier(name)
+
+    params = {"machine_learning_run_id": sanitized_name}
+
+    response = requests.post(
+        self.ml_runs_url,
+        headers=self._get_headers(),
+        params=params,
+        json=payload,
+    )
+
+    # Raise an exception for HTTP error status codes
+    response.raise_for_status()
+
+    return response.json()
+
+  def update_ml_run(
+      self
+  ) -> None:
+    """Update an existing ML run using the Google Cloud API."""
+    del self
+    pass

google_cloud_mldiagnostics/clients/logging_client.py

@@ -0,0 +1,142 @@
+"""Client for writing metrics directly to Google Cloud Logging."""
+
+import datetime
+import logging
+import os
+from typing import List, Optional
+
+from google.auth import credentials
+from google.cloud import logging as cloud_logging
+from google.cloud.logging_v2 import resource
+from google_cloud_mldiagnostics.src.google_cloud_mldiagnostics.custom_types import exceptions
+
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingClient:
+  """Client for writing metrics directly to Google Cloud Logging."""
+
+  def __init__(
+      self,
+      project_id: str,
+      log_name: str = "ml_diagnostics_metric",
+      credentials_path: Optional[str] = None,
+      user_credentials: Optional[credentials.Credentials] = None,
+  ):
+    """Initialize the logging client.
+
+    Args:
+        project_id: GCP project ID
+        log_name: Name of the log to write to
+        credentials_path: Optional path to service account credentials
+        user_credentials: Optional explicit credentials object.
+    """
+    self.project_id = project_id
+    self.log_name = log_name
+
+    try:
+      if user_credentials:
+        # Use explicit credentials if provided
+        self.client = cloud_logging.Client(
+            project=project_id, credentials=user_credentials
+        )
+      elif credentials_path:
+        # Fall back to credentials file
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
+        self.client = cloud_logging.Client(project=project_id)
+      else:
+        # Use default credentials
+        self.client = cloud_logging.Client(project=project_id)
+
+      # Get the logger for this log name
+      self.logger = self.client.logger(log_name)
+
+    except Exception as e:
+      raise exceptions.MLDiagnosticError(
+          f"Failed to initialize logging client: {e}"
+      ) from e
+
+  def write_metric(
+      self,
+      metric_name: str,
+      value: float | int | List[float | int],
+      run_id: str,
+      location: str,
+      step: Optional[int] = None,
+      labels: Optional[dict[str, str]] = None,
+  ):
+    """Write a single metric point to Cloud Logging.
+
+    Args:
+        metric_name: Name of the metric
+        value: Metric value. Can be a single int or float, or a list of ints
+            and floats.
+        run_id: ML run identifier
+        location: ML run region
+        step: Optional step number
+        labels: Optional additional labels
+
+    Returns:
+        True if the metric was written successfully, False otherwise.
+    """
+    try:
+      # Create timestamp
+      current_time = datetime.datetime.now(datetime.timezone.utc)
+
+      # Define the resource for the log entry
+      metric_resource = resource.Resource(
+          type="generic_node",
+          labels={
+              "project_id": self.project_id,
+              "location": location,
+              "namespace": metric_name,
+              "node_id": run_id,
+          },
+      )
+
+      # Build the log payload structure
+      if isinstance(value, (int, float)):
+        value = [value]
+      payload = {"values": value}
+
+      # Add optional fields
+      if step is not None:
+        payload["step_index"] = step
+      if labels:
+        payload.update(labels)
+
+      # Write to Cloud Logging
+      self.logger.log_struct(
+          payload,
+          severity="INFO",
+          timestamp=current_time,
+          resource=metric_resource,
+      )
+
+      logger.info("Successfully written metric to log: %s", metric_name)
+
+    except Exception as e:
+      raise exceptions.MLDiagnosticError(
+          f"Failed to write to Cloud Logging: {e}"
+      ) from e
+
+
+class NoOpLoggingClient(LoggingClient):
+  """A logging client that performs no operations."""
+
+  def __init__(self):  # pylint: disable=super-init-not-called
+    """Initializes the NoOp client. Does nothing."""
+    pass
+
+  def write_metric(
+      self,
+      metric_name: str,
+      value: float | int | List[float | int],
+      run_id: str,
+      location: str,
+      step: Optional[int] = None,
+      labels: Optional[dict[str, str]] = None,
+  ):
+    """This is a no-op and does not write any metrics."""
+    pass