Importing dag module - fix reference before assignment graph

TileDB-Inc · Oct 12, 2023 · d31e9ff · d31e9ff
1 parent 36111b1
commit d31e9ff
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/tiledb/cloud/bioimg/ingestion.py b/src/tiledb/cloud/bioimg/ingestion.py
@@ -1,6 +1,7 @@
 from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, Union
 
 import tiledb
+from tiledb.cloud import dag
 from tiledb.cloud._common.utils import logger
 from tiledb.cloud.bioimg.helpers import batch
 from tiledb.cloud.bioimg.helpers import get_uris
@@ -81,8 +82,6 @@ def ingest_tiff_udf(
                 with vfs.open(input) as src:
                     OMETiffConverter.to_tiledb(src, output, **kwargs)
 
-    submit = graph.submit_local if local else graph.submit
-
     if isinstance(source, str):
         # Handle only lists
         source = [source]
@@ -96,12 +95,13 @@ def ingest_tiff_udf(
     task_prefix = f"{dag_name} - Task"
 
     logger.info("Building graph")
-    graph = tiledb.cloud.dag.DAG(
+    graph = dag.DAG(
         name=dag_name,
-        mode=tiledb.cloud.dag.Mode.BATCH,
+        mode=dag.Mode.BATCH,
         max_workers=max_workers,
         namespace=namespace,
     )
+    submit = graph.submit_local if local else graph.submit
 
     # serialize udf arguments
     compressor = kwargs.pop("compressor", None)
@@ -116,7 +116,7 @@ def ingest_tiff_udf(
             threads,
             *args,
             name=f"{task_prefix} - {i}/{num_batches}",
-            mode=tiledb.cloud.dag.Mode.BATCH,
+            mode=dag.Mode.BATCH,
             resources=DEFAULT_RESOURCES if resources is None else resources,
             image_name=DEFAULT_IMG_NAME,
             compressor=compressor_serial,