Add submit_taskgraph for a common UX

src/tiledb/cloud/utilities/__init__.py

@@ -10,6 +10,7 @@
 from ._common import run_dag
 from ._common import serialize_filter
 from ._common import set_aws_context
+from ._common import submit_taskgraph
 from .consolidate import consolidate_and_vacuum
 from .consolidate import consolidate_fragments
 from .consolidate import group_fragments
@@ -31,6 +32,7 @@
     "read_file",
     "run_dag",
     "set_aws_context",
+    "submit_taskgraph",
     "consolidate_fragments",
     "consolidate_and_vacuum",
     "group_fragments",

src/tiledb/cloud/utilities/_common.py

@@ -8,6 +8,7 @@
 import sys
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import wait
+from concurrent.futures._base import TimeoutError
 from fnmatch import fnmatch
 from typing import (
     Any,
@@ -195,6 +196,74 @@ def run_dag(
     _print_logs(graph, debug=debug)
 
 
+def submit_taskgraph(
+    graph: dag.DAG,
+    *,
+    wait: bool = True,
+    update_sec: int = 10,
+) -> Mapping[str, str]:
+    """
+    Submit a taskgraph and optionally wait for completion.
+
+    :param graph: taskgraph to submit
+    :param wait: wait for the taskgraph to complete, defaults to True
+    :param update_sec: interval in seconds to update the progress bar, defaults to 10
+    :return: dictionary with status and graph_id
+    """
+
+    # Try to import tqdm, if not available, stub it out
+    try:
+        from tqdm import tqdm
+    except ImportError:
+
+        class tqdm:
+            def __init__(self, *args, **kwargs):
+                print("Please install tqdm to display a progress timer.")
+
+            def set_description_str(self, *args, **kwargs):
+                pass
+
+    graph.compute()
+
+    # Display a link to the log in TileDB, if available
+    if graph.server_graph_uuid:
+        print(
+            f"Taskgraph submitted to TileDB - https://cloud.tiledb.com/activity/taskgraphs/{graph.namespace}/{graph.server_graph_uuid}",
+        )
+
+    if wait:
+        # Display the progress bar in stdout, so the output is not red in a notebook
+        pbar = tqdm(
+            desc="Not Started",
+            bar_format="{desc}: {elapsed} (min:sec)",
+            file=sys.stdout,
+        )
+        cancelled = False
+
+        while True:
+            try:
+                graph.wait(update_sec)
+            except TimeoutError:
+                # Still running
+                pbar.set_description_str(str(graph.status))
+                continue
+            except RuntimeError as e:
+                if "No executions found for done Node" in str(e):
+                    cancelled = True
+            except StopIteration:
+                cancelled = True
+
+            if cancelled:
+                pbar.set_description("Cancelled")
+                return {"status": "Cancelled", "graph_id": str(graph.server_graph_uuid)}
+
+            # The taskgraph is done, break the loop
+            pbar.set_description_str(str(graph.status))
+            break
+
+    return {"status": str(graph.status), "graph_id": str(graph.server_graph_uuid)}
+
+
 def _print_logs(
     graph: dag.DAG,
     *,

tests/common/test_utils.py

@@ -3,12 +3,15 @@
 import pathlib
 import pickle
 import tempfile
+import time
 import unittest
 
 import pytz  # Test-only dependency.
 
+from tiledb.cloud import dag
 from tiledb.cloud._common import utils
 from tiledb.cloud.utilities import find
+from tiledb.cloud.utilities import submit_taskgraph
 
 
 class UtilsTest(unittest.TestCase):
@@ -69,6 +72,24 @@ def test_find(self):
                 len(list(find(tmp, include=lambda f: f.endswith(".dat")))), 1
             )
 
+    def test_submit_taskgraph(self):
+        # Submit and wait for completion.
+        tg = dag.DAG()
+
+        result = tg.submit_local(lambda: (time.sleep(3), 42)[1])
+        status = submit_taskgraph(tg, update_sec=1)
+
+        assert status == {"status": "Completed", "graph_id": "None"}
+        assert result.result() == 42
+
+        # Submit and continue.
+        tg = dag.DAG()
+        result = tg.submit_local(lambda: (time.sleep(3), 42)[1])
+        status = submit_taskgraph(tg, wait=False)
+
+        assert status == {"status": "Running", "graph_id": "None"}
+        assert result.result() == 42
+
 
 def _b64_unpickle(x):
     raw = base64.b64decode(x)