Merge pull request #611 from FlorentinD/progress-bar-fix

Fix progress bars getting stuck at 0%

Author: FlorentinD

changelog.md

@@ -20,6 +20,7 @@
 * Fixed an issue where source and target IDs of relationships in heterogeneous OGBL graphs were not parsed correctly.
 * Fixed an issue where configuration parameters such as `aggregation` were ignored by `gds.graph.toUndirected`.
 * Fixed an issue where the `database` given for the `GraphDataScience` construction was not used for metadata retrieval, causing an exception to be raised if the default "neo4j" database was missing.
+* Fixed an issue where progress bars would not always complete.
 
 
 ## Improvements

graphdatascience/query_runner/arrow_graph_constructor.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import concurrent
 import json
 import math
 import warnings
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, NoReturn, Optional
 
 import numpy
 import pyarrow.flight as flight
@@ -93,7 +95,7 @@ def _send_action(self, action_type: str, meta_data: Dict[str, Any]) -> None:
 
         json.loads(collected_result[0].body.to_pybytes().decode())
 
-    def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm) -> None:
+    def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm[NoReturn]) -> None:
         table = Table.from_pandas(df)
         batches = table.to_batches(self._chunk_size)
         flight_descriptor = {"name": self._graph_name, "entity_type": entity_type}
@@ -108,6 +110,8 @@ def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm) -> None:
             for partition in batches:
                 writer.write_batch(partition)
                 pbar.update(partition.num_rows)
+        # Force a refresh to avoid the progress bar getting stuck at 0%
+        pbar.refresh()
 
     def _send_dfs(self, dfs: List[DataFrame], entity_type: str) -> None:
         desc = "Uploading Nodes" if entity_type == "node" else "Uploading Relationships"

graphdatascience/query_runner/neo4j_query_runner.py

@@ -5,7 +5,7 @@
 import time
 import warnings
 from concurrent.futures import Future, ThreadPoolExecutor, wait
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
 from uuid import uuid4
 
 import neo4j
@@ -226,14 +226,19 @@ def _forward_cypher_warnings(self, notification: Dict[str, Any]) -> None:
             self._logger.info(notification)
 
     def _log(self, job_id: str, future: "Future[Any]", database: Optional[str] = None) -> None:
-        pbar = None
+        pbar: Optional[tqdm[NoReturn]] = None
         warn_if_failure = True
 
         while wait([future], timeout=self._LOG_POLLING_INTERVAL).not_done:
             try:
                 tier = "beta." if self._server_version < ServerVersion(2, 5, 0) else ""
+                # we only retrieve the progress of the root task
                 progress = self.run_cypher(
-                    f"CALL gds.{tier}listProgress('{job_id}') YIELD taskName, progress", database=database
+                    f"CALL gds.{tier}listProgress('{job_id}')"
+                    + " YIELD taskName, progress"
+                    + " RETURN taskName, progress"
+                    + " LIMIT 1",
+                    database=database,
                 )
             except Exception as e:
                 # Do nothing if the procedure either:
@@ -248,17 +253,19 @@ def _log(self, job_id: str, future: "Future[Any]", database: Optional[str] = Non
                     continue
 
             progress_percent = progress["progress"][0]
-            if not progress_percent == "n/a":
-                task_name = progress["taskName"][0].split("|--")[-1][1:]
-                pbar = pbar or tqdm(total=100, unit="%", desc=task_name)
-            else:
+            if progress_percent == "n/a":
                 return
 
+            root_task_name = progress["taskName"][0].split("|--")[-1][1:]
+            if not pbar:
+                pbar = tqdm(total=100, unit="%", desc=root_task_name, maxinterval=self._LOG_POLLING_INTERVAL)
+
             parsed_progress = float(progress_percent[:-1])
             pbar.update(parsed_progress - pbar.n)
 
         if pbar:
             pbar.update(100 - pbar.n)
+            pbar.refresh()
 
     def set_database(self, database: str) -> None:
         self._database = database

mypy.ini

@@ -8,9 +8,6 @@ ignore_missing_imports = True
 [mypy-pyarrow.flight]
 ignore_missing_imports = True
 
-[mypy-tqdm.auto]
-ignore_missing_imports = True
-
 [mypy-textdistance]
 ignore_missing_imports = True
 

requirements/dev/dev.txt

@@ -10,3 +10,4 @@ tox == 4.11.3
 types-setuptools == 68.1.0.1
 sphinx == 7.2.6
 types-requests
+types-tqdm

scripts/checkstyle

@@ -4,16 +4,16 @@ set -o errexit
 set -o nounset
 set -o pipefail
 
-black --check .
-isort --check .
-flake8
-mypy .
+python -m black --check .
+python -m isort --check .
+python -m flake8
+python -m mypy .
 
 NOTEBOOKS="./examples/*.ipynb" # ./examples/dev/*.ipynb"
 for f in $NOTEBOOKS
 do
   NB=$(cat $f)
-  FORMATTED_NB=$(jupyter nbconvert \
+  FORMATTED_NB=$(python -m jupyter nbconvert \
     --clear-output \
     --stdout \
     --ClearOutputPreprocessor.enabled=True \

scripts/makestyle

@@ -5,9 +5,9 @@ set -o nounset
 set -o pipefail
 set -o xtrace
 
-black .
-isort .
-jupyter nbconvert \
+python -m black .
+python -m isort .
+python -m jupyter nbconvert \
   --clear-output \
   --inplace \
   --ClearOutputPreprocessor.enabled=True \