Added Cloud Monitoring metrics to the EE ingestion pipeline.

ci3.8.yml

@@ -37,4 +37,5 @@ dependencies:
  - cython==0.29.34
  - earthengine-api==0.1.329
  - git+https:/dabhicusp/cdsapi-beta-google-weather-tools.git@master#egg=cdsapi # TODO([#474](https:/google/weather-tools/issues/474)): Compatible cdsapi with weather-dl.
+ - google-cloud-monitoring
  - .[test]

ci3.9.yml

@@ -37,4 +37,5 @@ dependencies:
  - cython==0.29.34
  - earthengine-api==0.1.329
  - git+https:/dabhicusp/cdsapi-beta-google-weather-tools.git@master#egg=cdsapi # TODO([#474](https:/google/weather-tools/issues/474)): Compatible cdsapi with weather-dl.
+ - google-cloud-monitoring
  - .[test]

environment.yml

@@ -31,6 +31,7 @@ dependencies:
  - firebase-admin==6.0.1
  - setuptools==70.3.0
  - git+https:/dabhicusp/cdsapi-beta-google-weather-tools.git@master#egg=cdsapi # TODO([#474](https:/google/weather-tools/issues/474)): Compatible cdsapi with weather-dl.
+ - google-cloud-monitoring
  - .
  - ./weather_dl
  - ./weather_mv

weather_mv/loader_pipeline/ee.py

@@ -250,7 +250,7 @@ class ToEarthEngine(ToDataSink):
  initialization_time_regex: str
  forecast_time_regex: str
  ingest_as_virtual_asset: bool
- use_deflate:bool
+ use_deflate: bool
  use_metrics: bool
  topic: str
 
@@ -376,8 +376,8 @@ def expand(self, paths):
  | 'IngestIntoEE' >> IngestIntoEETransform.from_kwargs(**vars(self))
  )
 
- if self.use_metrics:
- output | 'AddMetrics' >> beam.ParDo(AddMetrics())
+ if self.use_metrics and not self.skip_region_validation:
+ output | 'AddMetrics' >> AddMetrics.from_kwargs(**vars(self))
  else:
  (
  paths

weather_mv/loader_pipeline/metrics.py

@@ -22,15 +22,23 @@
 import time
 import typing as t
 from collections import OrderedDict
-from functools import wraps
 
 import apache_beam as beam
 from apache_beam.metrics import metric
+from apache_beam.transforms import window, trigger
+from functools import wraps
+from google.cloud import monitoring_v3
 
 from .sinks import get_file_time, KwargsFactoryMixin
 
 logger = logging.getLogger(__name__)
 
+# For Metrics API retry logic.
+INITIAL_DELAY = 1.0 # Initial delay in seconds.
+MAX_DELAY = 600 # Maximum delay before giving up in seconds.
+NUM_RETRIES = 10 # Number of tries with exponential backoff.
+TASK_QUEUE_WAIT_TIME = 120 # Task queue wait time in seconds.
+
 
 def timeit(func_name: str, keyed_fn: bool = False):
  """Decorator to add time it takes for an element to be processed by a stage.
@@ -115,7 +123,7 @@ def process(self, element) -> t.Iterator[t.Any]:
  yield element, time_dict
 
 
-class AddMetrics(beam.DoFn):
+class AddBeamMetrics(beam.DoFn):
  """DoFn to add Element Processing Time metric to beam. Expects PCollection
  to contain a time_dict."""
 
@@ -170,8 +178,90 @@ def process(self, element):
  logger.info(
  f"{uri}: Time from {current_step} -> {next_step}: {step_time} seconds."
  )
-
+ yield ("custom_metrics", (data_latency_ms / 1000, element_processing_time / 1000))
  except Exception as e:
  logger.warning(
  f"Some error occured while adding metrics. Error {e}"
  )
+
+
+@dataclasses.dataclass
+class CreateTimeSeries(beam.DoFn):
+ """DoFn to write metrics TimeSeries data in Google Cloud Monitoring."""
+ job_name: str
+ project: str
+ region: str
+
+ def create_time_series(self, metric_name: str, metric_value: float) -> None:
+ """Writes data to a Metrics TimeSeries."""
+ client = monitoring_v3.MetricServiceClient()
+ series = monitoring_v3.TimeSeries()
+ series.metric.type = f"custom.googleapis.com/{metric_name}"
+ series.metric.labels["description"] = metric_name
+ series.resource.type = "dataflow_job"
+ series.resource.labels["job_name"] = self.job_name
+ series.resource.labels["project_id"] = self.project
+ series.resource.labels["region"] = self.region
+
+ now = time.time()
+ seconds = int(now)
+ nanos = int((now - seconds) * 10**9)
+ interval = monitoring_v3.TimeInterval(
+ {"end_time": {"seconds": seconds, "nanos": nanos}}
+ )
+
+ point = monitoring_v3.Point(
+ {"interval": interval, "value": {"double_value": metric_value}}
+ )
+ series.points = [point]
+ client.create_time_series(
+ name=f"projects/{self.project}", time_series=[series]
+ )
+ logger.info(
+ f"Successfully created time series for {metric_name}. Metric value: {metric_value}."
+ )
+
+ def process(self, element: t.Any):
+ _, metric_values = element
+ data_latency_times = [x[0] for x in metric_values]
+ element_processing_times = [x[1] for x in metric_values]
+
+ logger.info(f"data_latency_time values: {data_latency_times}")
+ self.create_time_series("data_latency_time_max", max(data_latency_times))
+ self.create_time_series(
+ "data_latency_time_mean", sum(data_latency_times) / len(data_latency_times)
+ )
+
+ logger.info(f"element_processing_time values: {element_processing_times}")
+ self.create_time_series("element_processing_time_max", max(element_processing_times))
+ self.create_time_series(
+ "element_processing_time_mean", sum(element_processing_times) / len(element_processing_times)
+ )
+
+
+@dataclasses.dataclass
+class AddMetrics(beam.PTransform, KwargsFactoryMixin):
+ job_name: str
+ project: str
+ region: str
+
+ def expand(self, pcoll: beam.PCollection):
+ return (
+ pcoll
+ | "AddBeamMetrics" >> beam.ParDo(AddBeamMetrics())
+ | "AddTimestamps"
+ >> beam.Map(
+ lambda element: window.TimestampedValue(element, time.time())
+ )
+ | "Window"
+ >> beam.WindowInto(
+ window.GlobalWindows(),
+ trigger=trigger.Repeatedly(trigger.AfterProcessingTime(5)),
+ accumulation_mode=trigger.AccumulationMode.DISCARDING,
+ )
+ | "GroupByKeyAndWindow" >> beam.GroupByKey(lambda element: element)
+ | "CreateTimeSeries"
+ >> beam.ParDo(
+ CreateTimeSeries(self.job_name, self.project, self.region)
+ )
+ )

weather_mv/loader_pipeline/pipeline.py

@@ -47,6 +47,16 @@ def pattern_to_uris(match_pattern: str, is_zarr: bool = False) -> t.Iterable[str
  yield from [x.path for x in match.metadata_list]
 
 
+def arguments_to_dict(args: t.List[str]) -> t.Dict[str, str]:
+ """Converts a list of arguments to a dictionary."""
+ result = {}
+ for i in range(0, len(args), 2):
+ key = args[i].lstrip("-")
+ value = args[i + 1]
+ result[key] = value
+ return result
+
+
 def pipeline(known_args: argparse.Namespace, pipeline_args: t.List[str]) -> None:
  all_uris = list(pattern_to_uris(known_args.uris, known_args.zarr))
  if not all_uris:
@@ -75,7 +85,11 @@ def pipeline(known_args: argparse.Namespace, pipeline_args: t.List[str]) -> None
  elif known_args.subcommand == 'regrid' or known_args.subcommand == 'rg':
  paths | "Regrid" >> Regrid.from_kwargs(**vars(known_args))
  elif known_args.subcommand == 'earthengine' or known_args.subcommand == 'ee':
- paths | "MoveToEarthEngine" >> ToEarthEngine.from_kwargs(**vars(known_args))
+ # all_args will contain all the arguments passed to the pipeline.
+ all_args = {}
+ all_args.update(arguments_to_dict(pipeline_args))
+ all_args.update(**vars(known_args))
+ paths | "MoveToEarthEngine" >> ToEarthEngine.from_kwargs(**all_args)
  else:
  raise ValueError('invalid subcommand!')