Add more context with structured logging for metric uploads to GCM Monitoring.

PiperOrigin-RevId: 805110455

Author: dipannita08

ml_goodput_measurement/src/gcp_metrics.py

@@ -61,7 +61,9 @@ def create_time_series(
 
     return series
 
-  def send_metrics(self, metrics: list[Dict[str, Any]]):
+  def send_metrics(
+      self, metrics: list[Dict[str, Any]], context: Dict[str, Any] | None = None
+  ):
     """Sends multiple metrics to GCP Monitoring in a batch with dynamic resources.
 
     Args:
@@ -75,12 +77,12 @@ def send_metrics(self, metrics: list[Dict[str, Any]]):
                  - 'resource_type': str
                  - 'resource_labels': dict
     """
+    time_series_list = []
     try:
       now = time.time()
       seconds = int(now)
       nanos = int((now - seconds) * 10**9)
 
-      time_series_list = []
       for metric in metrics:
         try:
           metric_labels = metric.get("metric_labels", {})
@@ -96,11 +98,29 @@ def send_metrics(self, metrics: list[Dict[str, Any]]):
           )
           time_series_list.append(series)
         except Exception as e:  # pylint: disable=broad-exception-caught
-          logger.error("Failed to create time series: %s", e)
+          logger.error(
+              "Failed to create time series for metric '%s': %s",
+              metric.get("metric_type", "UNKNOWN"),
+              e,
+          )
+
+      if not time_series_list:
+        logger.warning(
+            "No valid time series were created, skipping GCM upload."
+        )
+        return
+
       self.client.create_time_series(
           name=self.project_name, time_series=time_series_list
       )
-      logger.info("Sent %d Goodput metrics to GCM Monitoring.", len(metrics))
+      log_data = {
+          "message": "Sent Goodput metrics to GCM Monitoring.",
+          "metrics_count": len(time_series_list),
+      }
+      if context:
+        log_data.update(context)
+
+      logger.info(log_data)
 
     except GoogleAPIError as e:
       logger.error("Failed to send Goodput metrics to GCM Monitoring: %s", e)

ml_goodput_measurement/src/monitoring.py

@@ -451,7 +451,13 @@ def _upload_goodput_metrics_to_gcm(
 
     # Send metrics to Google Cloud Monitoring.
     if metrics_sender and gcm_metrics:
-      metrics_sender.send_metrics(gcm_metrics)
+      log_context = {
+          'job_name': config.get('job_name', 'unknown-job'),
+          'pid': os.getpid(),
+          'worker': multiprocessing.current_process().name,
+          'metrics_type': 'cumulative',
+      }
+      metrics_sender.send_metrics(gcm_metrics, context=log_context)
 
   except Exception as e:  # pylint: disable=broad-exception-caught
     logger.error(
@@ -502,7 +508,13 @@ def _send_step_deviation_metric_to_gcp(
         },
     }]
     if metrics_sender:
-      metrics_sender.send_metrics(perf_metric)
+      log_context = {
+          'job_name': config.get('job_name', 'unknown-job'),
+          'pid': os.getpid(),
+          'worker': multiprocessing.current_process().name,
+          'metrics_type': 'step-time-deviation',
+      }
+      metrics_sender.send_metrics(perf_metric, context=log_context)
   except Exception as e:  # pylint: disable=broad-exception-caught
     logger.error('Error sending step deviation to GCM: %s', e)
 
@@ -585,7 +597,14 @@ def _upload_interval_goodput_metrics_to_gcm(
       })
 
     if metrics_sender and gcm_metrics:
-      metrics_sender.send_metrics(gcm_metrics)
+      log_context = {
+          'job_name': config.get('job_name', 'unknown-job'),
+          'pid': os.getpid(),
+          'worker': multiprocessing.current_process().name,
+          'metrics_type': 'rolling-window',
+          'window_size': str(window_size),
+      }
+      metrics_sender.send_metrics(gcm_metrics, context=log_context)
 
   except Exception as e:  # pylint: disable=broad-exception-caught
     logger.error(

ml_goodput_measurement/tests/gcp_metrics_test.py

@@ -72,8 +72,9 @@ def test_create_time_series(self):
           time_series.points[0].value.distribution_value, value
       )
 
+  @patch(f"{gcp_metrics.__name__}.logger")
   @patch("time.time")
-  def test_send_metrics(self, mock_time):
+  def test_send_metrics(self, mock_time, mock_logger):
     # Set a fixed return value for the mocked time.time()
     mock_time.return_value = 1677347200.5
 
@@ -94,8 +95,8 @@ def test_send_metrics(self, mock_time):
             "resource_labels": {"loc": "eu"},
         },
     ]
-
-    self.metrics_sender.send_metrics(metrics_to_send)
+    log_context = {"job_name": "test-job", "pid": 12345}
+    self.metrics_sender.send_metrics(metrics_to_send, context=log_context)
 
     # Verify that create_time_series was called with the correct arguments
     expected_name = f"projects/{self.project_id}"
@@ -126,6 +127,16 @@ def test_send_metrics(self, mock_time):
       self.assertEqual(actual.resource.labels, expected.resource.labels)
       self.assertEqual(actual.metric.labels, expected.metric.labels)
 
+    mock_logger.info.assert_called_once()
+    expected_log_payload = {
+        "message": "Sent Goodput metrics to GCM Monitoring.",
+        "metrics_count": 2,
+        "job_name": "test-job",
+        "pid": 12345,
+    }
+    actual_log_payload = mock_logger.info.call_args.args[0]
+    self.assertEqual(actual_log_payload, expected_log_payload)
+
   @patch("cloud_goodput.ml_goodput_measurement.src.gcp_metrics.logger.error")
   def test_send_metrics_failure(self, mock_logging_error):