Add progress bar to remote write-back

Author: FlorentinD

graphdatascience/arrow_client/v2/remote_write_back_client.py

@@ -20,14 +20,14 @@ def __init__(self, arrow_client: AuthenticatedArrowClient, query_runner: QueryRu
         protocol_version = ProtocolVersionResolver(query_runner).resolve()
         self._write_protocol = WriteProtocol.select(protocol_version)
 
-    # TODO: Add progress logging
     def write(
         self,
         graph_name: str,
         job_id: str,
         concurrency: Optional[int] = None,
         property_overwrites: Optional[dict[str, str]] = None,
         relationship_type_overwrite: Optional[str] = None,
+        log_progress: bool = True,
     ) -> WriteBackResult:
         arrow_config = self._arrow_configuration()
 
@@ -49,7 +49,11 @@ def write(
         start_time = time.time()
 
         result = self._write_protocol.run_write_back(
-            self._query_runner, write_back_params, None, TerminationFlagNoop()
+            self._query_runner,
+            write_back_params,
+            None,
+            log_progress=log_progress,
+            terminationFlag=TerminationFlagNoop(),
         ).squeeze()
         write_millis = int((time.time() - start_time) * 1000)
 

graphdatascience/procedure_surface/arrow/catalog/node_properties_arrow_endpoints.py

@@ -34,6 +34,7 @@ def __init__(
         self._node_property_endpoints = NodePropertyEndpoints(
             arrow_client, self._write_back_client, show_progress=show_progress
         )
+        self._show_progress = show_progress
 
     def stream(
         self,
@@ -112,6 +113,7 @@ def write(
             job_id,
             concurrency=write_concurrency if write_concurrency is not None else concurrency,
             property_overwrites=node_property_spec.to_dict(),
+            log_progress=self._show_progress and log_progress,
         )
 
         return NodePropertiesWriteResult(

graphdatascience/procedure_surface/arrow/catalog/relationship_arrow_endpoints.py

@@ -101,6 +101,7 @@ def write(
             job_id,
             concurrency=write_concurrency if write_concurrency is not None else concurrency,
             relationship_type_overwrite=relationship_type,
+            log_progress=log_progress and self._show_progress,
         )
 
         written_relationships = (

graphdatascience/procedure_surface/arrow/catalog_arrow_endpoints.py

@@ -128,7 +128,7 @@ def project(
             TerminationFlag.create(),
             None,
             None,
-            logging,
+            self._show_progress and logging,
         )
 
         job_result = ProjectionResult(**JobClient.get_summary(self._arrow_client, job_id))

graphdatascience/procedure_surface/arrow/node_property_endpoints.py

@@ -90,6 +90,7 @@ def run_job_and_write(
             job_id,
             concurrency=write_concurrency if write_concurrency is not None else concurrency,
             property_overwrites=property_overwrites,
+            log_progress=show_progress,
         )
 
         # modify computation result to include write details

graphdatascience/query_runner/progress/progress_bar.py

@@ -1,12 +1,14 @@
-from typing import Any, Optional
+from __future__ import annotations
+
+from types import TracebackType
+from typing import Any, Optional, Type
 
 from tqdm.auto import tqdm
 
 from graphdatascience.query_runner.progress.progress_provider import TaskWithProgress
 
 
 class TqdmProgressBar:
-    # TODO helper method for creating for a test with obserable progress
     def __init__(self, task_name: str, relative_progress: Optional[float], bar_options: dict[str, Any] = {}):
         root_task_name = task_name
         if relative_progress is None:  # Qualitative progress report
@@ -26,6 +28,17 @@ def __init__(self, task_name: str, relative_progress: Optional[float], bar_optio
                 **bar_options,
             )
 
+    def __enter__(self: TqdmProgressBar) -> TqdmProgressBar:
+        return self
+
+    def __exit__(
+        self,
+        exception_type: Optional[Type[BaseException]],
+        exception_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        self.finish(success=exception_value is None)
+
     def update(
         self,
         status: str,
@@ -43,11 +56,11 @@ def update(
     def finish(self, success: bool) -> None:
         if not success:
             self._tqdm_bar.set_postfix_str("status: FAILED", refresh=True)
-            return
-
-        if self._tqdm_bar.total is not None:
-            self._tqdm_bar.update(self._tqdm_bar.total - self._tqdm_bar.n)
-        self._tqdm_bar.set_postfix_str("status: FINISHED", refresh=True)
+        else:
+            if self._tqdm_bar.total is not None:
+                self._tqdm_bar.update(self._tqdm_bar.total - self._tqdm_bar.n)
+            self._tqdm_bar.set_postfix_str("status: FINISHED", refresh=True)
+        self._tqdm_bar.close()
 
     @staticmethod
     def _relative_progress(task: TaskWithProgress) -> Optional[float]:

graphdatascience/query_runner/protocol/write_protocols.py

@@ -6,6 +6,7 @@
 from tenacity import retry, retry_if_result, wait_incrementing
 
 from graphdatascience.call_parameters import CallParameters
+from graphdatascience.query_runner.progress.progress_bar import TqdmProgressBar
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.query_runner.query_mode import QueryMode
 from graphdatascience.query_runner.query_runner import QueryRunner
@@ -33,6 +34,7 @@ def run_write_back(
         query_runner: QueryRunner,
         parameters: CallParameters,
         yields: Optional[list[str]],
+        log_progress: bool,
         terminationFlag: TerminationFlag,
     ) -> DataFrame:
         """Executes the write-back procedure"""
@@ -68,6 +70,7 @@ def run_write_back(
         query_runner: QueryRunner,
         parameters: CallParameters,
         yields: Optional[list[str]],
+        log_progress: bool,
         terminationFlag: TerminationFlag,
     ) -> DataFrame:
         return query_runner.call_procedure(
@@ -108,6 +111,7 @@ def run_write_back(
         query_runner: QueryRunner,
         parameters: CallParameters,
         yields: Optional[list[str]],
+        log_progress: bool,
         terminationFlag: TerminationFlag,
     ) -> DataFrame:
         return query_runner.call_procedure(
@@ -123,6 +127,9 @@ def run_write_back(
 
 
 class RemoteWriteBackV3(WriteProtocol):
+    def __init__(self, progress_bar_options: dict[str, Any] | None = None):
+        self._progress_bar_options = progress_bar_options or {}
+
     def write_back_params(
         self,
         graph_name: str,
@@ -138,6 +145,7 @@ def run_write_back(
         query_runner: QueryRunner,
         parameters: CallParameters,
         yields: Optional[list[str]],
+        log_progress: bool,
         terminationFlag: TerminationFlag,
     ) -> DataFrame:
         def is_not_completed(result: DataFrame) -> bool:
@@ -156,9 +164,9 @@ def is_not_completed(result: DataFrame) -> bool:
                 logging.DEBUG,
             ),
         )
-        def write_fn() -> DataFrame:
+        def write_fn(progress_bar: Optional[TqdmProgressBar]) -> DataFrame:
             terminationFlag.assert_running()
-            return query_runner.call_procedure(
+            result = query_runner.call_procedure(
                 ProtocolVersion.V3.versioned_procedure_name("gds.arrow.write"),
                 parameters,
                 yields,
@@ -168,4 +176,17 @@ def write_fn() -> DataFrame:
                 custom_error=False,
             )
 
-        return write_fn()
+            if progress_bar:
+                progress_bar.update(status=result.squeeze()["status"], progress=result.squeeze()["progress"] * 100)
+
+            return result
+
+        if log_progress:
+            with TqdmProgressBar(
+                task_name=f"Write-Back (graph: {parameters['graphName']})",
+                relative_progress=0.0,
+                bar_options=self._progress_bar_options,
+            ) as progress_bar:
+                return write_fn(progress_bar)
+        else:
+            return write_fn(None)

graphdatascience/query_runner/session_query_runner.py

@@ -226,7 +226,9 @@ def _remote_write_back(
         write_back_start = time.time()
 
         def run_write_back() -> DataFrame:
-            return write_protocol.run_write_back(self._db_query_runner, write_back_params, yields, terminationFlag)
+            return write_protocol.run_write_back(
+                self._db_query_runner, write_back_params, yields, log_progress=logging, terminationFlag=terminationFlag
+            )
 
         try:
             # Skipping progress for now as export has a different jobId

graphdatascience/tests/integrationV2/procedure_surface/arrow/test_pagerank_arrow_endpoints.py

@@ -94,7 +94,9 @@ def test_pagerank_mutate(pagerank_endpoints: PageRankArrowEndpoints, sample_grap
 @pytest.mark.db_integration
 def test_pagerank_write(arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner, db_graph: GraphV2) -> None:
     """Test PageRank write operation."""
-    endpoints = PageRankArrowEndpoints(arrow_client, RemoteWriteBackClient(arrow_client, query_runner))
+    endpoints = PageRankArrowEndpoints(
+        arrow_client, RemoteWriteBackClient(arrow_client, query_runner), show_progress=True
+    )
     result = endpoints.write(G=db_graph, write_property="pagerank")
 
     assert isinstance(result, PageRankWriteResult)

graphdatascience/tests/unit/conftest.py

@@ -169,11 +169,11 @@ def create_graph_constructor(
     def cloneWithoutRouting(self, host: str, port: int) -> QueryRunner:
         return self
 
-    def set__mock_result(self, result: DataFrame) -> None:
+    def set__mock_result(self, result: QueryResult) -> None:
         self._result_map.clear()
         self._result_map[""] = result
 
-    def add__mock_result(self, query_sub_string: str, result: DataFrame) -> None:
+    def add__mock_result(self, query_sub_string: str, result: QueryResult) -> None:
         self._result_map[query_sub_string] = result
 
     def get_mock_result(self, query: str) -> QueryResult:
@@ -190,6 +190,9 @@ def get_mock_result(self, query: str) -> QueryResult:
             )
         if len(matched_results) == 0:
             return DataFrame()
+
+        if isinstance(matched_results[0][1], Exception):
+            raise matched_results[0][1]
         return matched_results[0][1]