Skip to content

Commit

Permalink
Support using passed config for VFS listing in bioimage ingestion (#544)
Browse files Browse the repository at this point in the history
The bioimage ingestion needs to use the config for both ingestion/read
(currently supported) and for the bucket listing (what this PR adds).
  • Loading branch information
Shelnutt2 authored Jul 12, 2024
1 parent 7f20a6b commit cc8d879
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions src/tiledb/cloud/bioimg/ingestion.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ def build_io_uris_ingestion(
output_ext: str,
supported_exts: Tuple[str],
logger: logging.Logger,
config: Optional[Mapping[str, Any]] = None,
):
"""Match input uri/s with output destinations
:param source: A sequence of paths or path to input
Expand All @@ -110,7 +111,8 @@ def build_io_uris_ingestion(

import tiledb

vfs = tiledb.VFS()
ctx = tiledb.Ctx(config)
vfs = tiledb.VFS(ctx=ctx)
# Even though a tuple by definition when passed through submit becomes list
supported_exts = tuple(supported_exts)

Expand Down Expand Up @@ -155,12 +157,13 @@ def build_input_batches(
supported_exts: Tuple,
*,
verbose: bool,
config: Optional[Mapping[str, Any]] = None,
):
logger = get_logger_wrapper(verbose)

"""Groups input URIs into batches."""
uri_pairs = build_io_uris_ingestion(
source, output, out_ext, supported_exts, logger
source, output, out_ext, supported_exts, logger, config=config
)
logger.debug(f"Input batches:{uri_pairs}")
# If the user didn't specify a number of batches, run every import
Expand Down Expand Up @@ -215,7 +218,7 @@ def ingest_tiff_udf(
raise ValueError

write_context = tiledb.Ctx(config)
vfs = tiledb.VFS()
vfs = tiledb.VFS(ctx=write_context)

for input, output in io_uris:
with vfs.open(input) as src:
Expand Down Expand Up @@ -340,7 +343,8 @@ def register_dataset_udf(
_SUPPORTED_EXTENSIONS,
*args,
verbose=verbose,
access_credentials_name=acn,
config=config,
access_credentials_name=access_credentials_name,
name=f"{dag_name} input collector",
result_format="json",
)
Expand Down

0 comments on commit cc8d879

Please sign in to comment.