From 4727968e27da56da4e2fd3c1e24837233758ba49 Mon Sep 17 00:00:00 2001
From: j9sh264 <jash@infocusp.com>
Date: Mon, 23 Sep 2024 07:46:01 +0000
Subject: [PATCH] Now a single call will be made to write the 4 timeseries.

---
 weather_mv/loader_pipeline/metrics.py  | 41 ++++++++++++++++++--------
 weather_mv/loader_pipeline/pipeline.py | 17 +++--------
 2 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/weather_mv/loader_pipeline/metrics.py b/weather_mv/loader_pipeline/metrics.py
index 25258173..ab3c24ae 100644
--- a/weather_mv/loader_pipeline/metrics.py
+++ b/weather_mv/loader_pipeline/metrics.py
@@ -188,6 +188,7 @@ def process(self, element):
 @dataclasses.dataclass
 class CreateTimeSeries(beam.DoFn):
     """DoFn to write metrics TimeSeries data in Google Cloud Monitoring."""
+
     job_name: str
     project: str
     region: str
@@ -214,12 +215,7 @@ def create_time_series(self, metric_name: str, metric_value: float) -> None:
             {"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}."
-        )
+        return series
 
     def process(self, element: t.Any):
         _, metric_values = element
@@ -227,15 +223,34 @@ def process(self, element: t.Any):
         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)
+        data_latency_max_series = self.create_time_series_object(
+            "data_latency_time_max", max(data_latency_times)
+        )
+        data_latency_mean_series = self.create_time_series_object(
+            "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)
+        logger.info(
+            f"element_processing_time values: {element_processing_times}"
+        )
+        element_processing_max_series = self.create_time_series_object(
+            "element_processing_time_max", max(element_processing_times)
+        )
+        element_processing_mean_series = self.create_time_series_object(
+            "element_processing_time_mean",
+            sum(element_processing_times) / len(element_processing_times),
+        )
+
+        client = monitoring_v3.MetricServiceClient()
+        client.create_time_series(
+            name=f"projects/{self.project}",
+            time_series=[
+                data_latency_max_series,
+                data_latency_mean_series,
+                element_processing_max_series,
+                element_processing_mean_series,
+            ],
         )
 
 
diff --git a/weather_mv/loader_pipeline/pipeline.py b/weather_mv/loader_pipeline/pipeline.py
index d0a135ae..5daabd55 100644
--- a/weather_mv/loader_pipeline/pipeline.py
+++ b/weather_mv/loader_pipeline/pipeline.py
@@ -21,6 +21,7 @@
 
 import apache_beam as beam
 from apache_beam.io.filesystems import FileSystems
+from apache_beam.options.pipeline_options import PipelineOptions
 
 from .bq import ToBigQuery
 from .regrid import Regrid
@@ -47,16 +48,6 @@ 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:
@@ -85,10 +76,10 @@ 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':
+            pipeline_options = PipelineOptions(pipeline_args)
+            pipeline_options_dict = pipeline_options.get_all_options()
             # 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))
+            all_args = {**vars(known_args), **pipeline_options_dict}
             paths | "MoveToEarthEngine" >> ToEarthEngine.from_kwargs(**all_args)
         else:
             raise ValueError('invalid subcommand!')