Support access_credentials_name for BATCH tasks (#308)

TileDB Cloud supports using certain types of credentials like AWS IAM
roles, Azure SAS tokens as part of a batch task graph. These credentials
can provide direct access to object stores which is useful in some
cases, especially during ingestion. This PR extends the existing support
for ingestion by adding the new parameter `acn` which is passed to batch
tasks.

Author: Shelnutt2

apis/python/src/tiledb/vector_search/ingestion.py

@@ -51,6 +51,7 @@ def ingest(
     trace_id: Optional[str] = None,
     use_sklearn: bool = True,
     mode: Mode = Mode.LOCAL,
+    acn: Optional[str] = None,
     **kwargs,
 ):
     """
@@ -134,6 +135,8 @@ def ingest(
         tiledb.vector_search's. Defaults to true.
     mode: Mode
         execution mode, defaults to LOCAL use BATCH for distributed execution
+    acn: Optional[str]
+        access credential name to be used when running in BATCH mode for object store access
     """
     import enum
     import json
@@ -1891,8 +1894,11 @@ def create_ingestion_dag(
         trace_id: Optional[str] = None,
         use_sklearn: bool = True,
         mode: Mode = Mode.LOCAL,
+        acn: Optional[str] = None,
         namespace: Optional[str] = None
     ) -> dag.DAG:
+
+        kwargs = {}
         if mode == Mode.BATCH:
             d = dag.DAG(
                 name="vector-ingestion",
@@ -1905,6 +1911,8 @@ def create_ingestion_dag(
                 namespace=namespace
             )
             threads = 16
+            if acn:
+                kwargs["access_credentials_name"] = acn
         else:
             d = dag.DAG(
                 name="vector-ingestion",
@@ -1950,6 +1958,7 @@ def create_ingestion_dag(
                 name="ingest",
                 resources={"cpu": str(threads), "memory": "16Gi"},
                 image_name=DEFAULT_IMG_NAME,
+                **kwargs
             )
             return d
         elif index_type == "IVF_FLAT":
@@ -1966,6 +1975,7 @@ def create_ingestion_dag(
                     name="copy-centroids",
                     resources={"cpu": "1", "memory": "2Gi"},
                     image_name=DEFAULT_IMG_NAME,
+                    **kwargs
                 )
             else:
                 random_sample_nodes = []
@@ -2020,6 +2030,7 @@ def create_ingestion_dag(
                                 name="read-random-sample-" + str(idx),
                                 resources={"cpu": "2", "memory": "6Gi"},
                                 image_name=DEFAULT_IMG_NAME,
+                                **kwargs
                             )
                         )
                         num_sampled += num_to_sample
@@ -2048,6 +2059,7 @@ def create_ingestion_dag(
                         name="kmeans",
                         resources={"cpu": "8", "memory": "32Gi"},
                         image_name=DEFAULT_IMG_NAME,
+                        **kwargs
                     )
 
                     for random_sample_node in random_sample_nodes:
@@ -2076,6 +2088,7 @@ def create_ingestion_dag(
                         name="init-centroids",
                         resources={"cpu": "1", "memory": "1Gi"},
                         image_name=DEFAULT_IMG_NAME,
+                        **kwargs
                     )
 
                     for random_sample_node in random_sample_nodes:
@@ -2110,6 +2123,7 @@ def create_ingestion_dag(
                                     name="k-means-part-" + str(task_id),
                                     resources={"cpu": str(threads), "memory": "12Gi"},
                                     image_name=DEFAULT_IMG_NAME,
+                                    **kwargs
                                 )
                             )
                             task_id += 1
@@ -2122,6 +2136,7 @@ def create_ingestion_dag(
                                     name="update-centroids-" + str(i),
                                     resources={"cpu": "1", "memory": "8Gi"},
                                     image_name=DEFAULT_IMG_NAME,
+                                    **kwargs
                                 )
                             )
                         internal_centroids_node = submit(
@@ -2130,6 +2145,7 @@ def create_ingestion_dag(
                             name="update-centroids",
                             resources={"cpu": "1", "memory": "8Gi"},
                             image_name=DEFAULT_IMG_NAME,
+                            **kwargs
                         )
                     centroids_node = submit(
                         write_centroids,
@@ -2143,6 +2159,7 @@ def create_ingestion_dag(
                         name="write-centroids",
                         resources={"cpu": "1", "memory": "2Gi"},
                         image_name=DEFAULT_IMG_NAME,
+                        **kwargs
                     )
 
             compute_indexes_node = submit(
@@ -2155,6 +2172,7 @@ def create_ingestion_dag(
                 name="compute-indexes",
                 resources={"cpu": "1", "memory": "2Gi"},
                 image_name=DEFAULT_IMG_NAME,
+                **kwargs
             )
 
             task_id = 0
@@ -2184,6 +2202,7 @@ def create_ingestion_dag(
                     name="ingest-" + str(task_id),
                     resources={"cpu": str(threads), "memory": "16Gi"},
                     image_name=DEFAULT_IMG_NAME,
+                    **kwargs
                 )
                 ingest_node.depends_on(centroids_node)
                 compute_indexes_node.depends_on(ingest_node)
@@ -2203,6 +2222,7 @@ def create_ingestion_dag(
                     name="ingest-" + str(task_id),
                     resources={"cpu": str(threads), "memory": "16Gi"},
                     image_name=DEFAULT_IMG_NAME,
+                    **kwargs
                 )
                 ingest_additions_node.depends_on(centroids_node)
                 compute_indexes_node.depends_on(ingest_additions_node)
@@ -2229,6 +2249,7 @@ def create_ingestion_dag(
                     name="consolidate-partition-" + str(task_id),
                     resources={"cpu": str(threads), "memory": "16Gi"},
                     image_name=DEFAULT_IMG_NAME,
+                    **kwargs
                 )
                 consolidate_partition_node.depends_on(compute_indexes_node)
                 task_id += 1
@@ -2513,6 +2534,7 @@ def consolidate_and_vacuum(
             trace_id=trace_id,
             use_sklearn=use_sklearn,
             mode=mode,
+            acn=acn,
             namespace=namespace,
         )
         logger.debug("Submitting ingestion graph")