Display running subtask in progress bar

Author: FlorentinD

changelog.md

@@ -17,6 +17,7 @@
 ## Improvements
 
 * Display progress bar for remote projection and open-ended tasks.
+* Improve progress bar by showing currently running task.
 * Allow passing the optional graph filter also as type `str` to `gds.graph.list()` instead of only `Graph`.
 
 

graphdatascience/query_runner/progress/progress_provider.py

@@ -8,6 +8,7 @@ class TaskWithProgress:
     task_name: str
     progress_percent: str
     status: str
+    sub_tasks_description: Optional[str] = None
 
 
 class ProgressProvider(ABC):

graphdatascience/query_runner/progress/query_progress_logger.py

@@ -70,7 +70,6 @@ def _log(
                     if has_relative_progress:
                         pbar = tqdm(total=100, unit="%", desc=root_task_name, maxinterval=self._LOG_POLLING_INTERVAL)
                     else:
-                        # TODO add {n_fmt} once task_with_progress provides the absolute progress
                         pbar = tqdm(
                             total=None,
                             unit="",
@@ -79,7 +78,9 @@ def _log(
                             bar_format="{desc} [elapsed: {elapsed} {postfix}]",
                         )
 
-                pbar.set_postfix_str(f"status: {task_with_progress.status}")
+                pbar.set_postfix_str(
+                    f"status: {task_with_progress.status}, task: {task_with_progress.sub_tasks_description}"
+                )
                 if has_relative_progress:
                     parsed_progress = float(progress_percent[:-1])
                     new_progress = parsed_progress - pbar.n

graphdatascience/query_runner/progress/query_progress_provider.py

@@ -21,16 +21,33 @@ def __init__(
 
     def root_task_with_progress(self, job_id: str, database: Optional[str] = None) -> TaskWithProgress:
         tier = "beta." if self._server_version_func() < ServerVersion(2, 5, 0) else ""
-        # we only retrieve the progress of the root task
+
+        # expect at exactly one row (query will fail if not existing)
         progress = self._run_cypher_func(
             f"CALL gds.{tier}listProgress('{job_id}')"
             + " YIELD taskName, progress, status"
-            + " RETURN taskName, progress, status"
-            + " LIMIT 1",
+            + " RETURN taskName, progress, status",
             database,
-        ).squeeze()  # expect at exactly one row (query will fail if not existing)
-
-        progress_percent = progress["progress"]
-        root_task_name = progress["taskName"].split("|--")[-1][1:]
-
-        return TaskWithProgress(root_task_name, progress_percent, progress["status"])
+        )
+
+        # compute depth of each subtask
+        progress["trimmedName"] = progress["taskName"].str.lstrip()
+        progress["depth"] = progress["taskName"].str.len() - progress["trimmedName"].str.len()
+        progress.sort_values("depth", ascending=True, inplace=True)
+
+        root_task = progress.iloc[0]
+        root_progress_percent = root_task["progress"]
+        root_task_name = root_task["trimmedName"].replace("|--", "")
+        root_status = root_task["status"]
+
+        subtask_descriptions = None
+        running_tasks = progress[progress["status"] == "RUNNING"]
+        if running_tasks["taskName"].size > 1:  # at least one subtask
+            subtasks = running_tasks[1:]  # remove root task
+            subtask_descriptions = "::".join(
+                list(subtasks["taskName"].apply(lambda name: name.split("|--")[-1].strip()))
+            )
+
+        return TaskWithProgress(
+            root_task_name, root_progress_percent, root_status, sub_tasks_description=subtask_descriptions
+        )

graphdatascience/tests/unit/query_runner/progress/test_query_progress_logger.py

@@ -12,10 +12,7 @@
 
 def test_call_through_functions() -> None:
     def fake_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
-        assert (
-            query
-            == "CALL gds.listProgress('foo') YIELD taskName, progress, status RETURN taskName, progress, status LIMIT 1"
-        )
+        assert "CALL gds.listProgress('foo')" in query
         assert database == "database"
 
         return DataFrame([{"progress": "n/a", "taskName": "Test task", "status": "RUNNING"}])
@@ -46,10 +43,7 @@ def fake_query() -> DataFrame:
 
 def test_uses_beta_endpoint() -> None:
     def fake_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
-        assert (
-            query
-            == "CALL gds.beta.listProgress('foo') YIELD taskName, progress, status RETURN taskName, progress, status LIMIT 1"
-        )
+        assert "CALL gds.beta.listProgress('foo')" in query
         assert database == "database"
 
         return DataFrame([{"progress": "n/a", "taskName": "Test task", "status": "RUNNING"}])
@@ -76,6 +70,33 @@ def simple_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
     assert isinstance(progress_provider, QueryProgressProvider)
 
 
+def test_uses_query_provider_with_task_description() -> None:
+    server_version = ServerVersion(3, 0, 0)
+    detailed_progress = DataFrame(
+        [
+            {"progress": "n/a", "taskName": "Test task", "status": "RUNNING"},
+            {"progress": "n/a", "taskName": "    |-- root 1/1", "status": "RUNNING"},
+            {"progress": "n/a", "taskName": "       |-- leaf", "status": "RUNNING"},
+            {"progress": "n/a", "taskName": "finished task", "status": "FINISHED"},
+            {"progress": "n/a", "taskName": "pending task", "status": "PENDING"},
+        ]
+    )
+
+    query_runner = CollectingQueryRunner(server_version, result_mock={"gds.listProgress": detailed_progress})
+
+    def simple_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
+        return query_runner.run_cypher(query, db=database)
+
+    qpl = QueryProgressLogger(simple_run_cypher, lambda: server_version)
+    progress_provider = qpl._select_progress_provider("test-job")
+    assert isinstance(progress_provider, QueryProgressProvider)
+
+    progress = progress_provider.root_task_with_progress("test-job", "database")
+
+    assert progress.sub_tasks_description == "root 1/1::leaf"
+    assert progress.task_name == "Test task"
+
+
 def test_uses_static_store() -> None:
     def fake_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
         return DataFrame([{"progress": "n/a", "taskName": "Test task", "status": "RUNNING"}])
@@ -88,3 +109,4 @@ def fake_run_cypher(query: str, database: Optional[str] = None) -> DataFrame:
     task_with_volume = progress_provider.root_task_with_progress("test-job")
     assert task_with_volume.task_name == "Test task"
     assert task_with_volume.progress_percent == "n/a"
+    assert task_with_volume.sub_tasks_description is None