Skip to content

biopb.tensor.client

biopb.tensor.client

Python client for TensorFlight server.

This module provides a lazy numpy-like array interface using dask.array for accessing tensors stored in a Flight server.

Features: - Lazy chunk loading via dask.array - LRU caching via cachey - Numpy-compatible slicing and operations

ResolveCancelled

Bases: Exception

Raised by :meth:TensorFlightClient.resolve when its should_cancel callback asks it to stop.

The client stops consuming the resolve stream and unwinds; the server's recall daemon thread runs to completion and caches its result, so a later :meth:resolve coalesces onto the finished work rather than re-downloading.

TensorFlightClient

TensorFlightClient(
    location: str = "grpc://localhost:8815",
    cache_bytes: int = 1000000000,
    token: Optional[str] = None,
)

Client for accessing tensors from a TensorFlightServer.

This client provides lazy, cached access to multi-dimensional arrays stored in a Flight server, with support for multifield acquisitions where tensors within a source have different shapes.

Usage

client = TensorFlightClient('grpc://localhost:8815')

List data sources (each may contain multiple tensors)

sources = client.list_sources()

Get source-level metadata

metadata = client.get_source_metadata('my-source')

Access a tensor by its globally-unique array_id (identity policy):

'source_id/field' for a multi-tensor source, or 'source_id' for a

single-tensor one. See proto/biopb/tensor/descriptor.proto.

arr = client.get_tensor('my-source/tensor-0') # Returns dask.array data = arr[0:100, 0:100].compute() # Load slice

Note

The dask arrays returned by get_tensor() are picklable and work with dask.distributed: each worker fetches chunks over its own connection, so you can scatter an array across a cluster and compute on it.

Initialize the Flight client.

Parameters:

Name Type Description Default
location str

Flight server location

'grpc://localhost:8815'
cache_bytes int

Maximum bytes for chunk cache (default 1GB)

1000000000
token Optional[str]

Bearer token for server authentication. None disables auth.

None
Source code in src/main/python/biopb/tensor/client.py
def __init__(
    self,
    location: str = "grpc://localhost:8815",
    cache_bytes: int = 1_000_000_000,  # 1GB default
    token: Optional[str] = None,
):
    """Initialize the Flight client.

    Args:
        location: Flight server location
        cache_bytes: Maximum bytes for chunk cache (default 1GB)
        token: Bearer token for server authentication.  ``None`` disables auth.
    """
    logger.info(
        f"Connecting to Flight server at {location}, cache={cache_bytes}B, auth={token is not None}"
    )
    # Normalize location for Arrow Flight (grpcs:// -> grpc+tls://)
    normalized = _normalize_location(location)
    # Store pickle-safe connection parameters
    self._location = normalized
    self._token = token
    self._cache_bytes = cache_bytes
    # Create FlightClient for direct API calls (list_flights, get_flight_info, uploads)
    self._client = flight.FlightClient(normalized)
    self._call_options = (
        flight.FlightCallOptions(
            headers=[(b"authorization", f"Bearer {token}".encode())]
        )
        if token
        else flight.FlightCallOptions()
    )
    # Cache descriptors for metadata. Keyed by (source_id, bare array_id):
    # array_id alone is NOT unique across sources -- e.g. aicsimageio names
    # every single-scene file's tensor "Image:0" -- so a bare-array_id key
    # collides and silently returns another source's descriptor (issue #45).
    self._sources: Dict[str, DataSourceDescriptor] = {}
    self._descriptors: Dict[Tuple[str, str], TensorDescriptor] = {}

list_sources

list_sources() -> Dict[str, DataSourceDescriptor]

List available data sources.

Returns:

Type Description
Dict[str, DataSourceDescriptor]

Dictionary mapping source_id to DataSourceDescriptor.

Dict[str, DataSourceDescriptor]

Each DataSourceDescriptor.tensors contains TensorDescriptor info

Dict[str, DataSourceDescriptor]

with shape/dtype for all tensors in that source.

Note

Results may be truncated if server has max_list_flights_results configured. Check schema metadata for truncation info (truncated=True indicates more sources exist on server than were returned).

Source code in src/main/python/biopb/tensor/client.py
def list_sources(self) -> Dict[str, DataSourceDescriptor]:
    """List available data sources.

    Returns:
        Dictionary mapping source_id to DataSourceDescriptor.
        Each DataSourceDescriptor.tensors contains TensorDescriptor info
        with shape/dtype for all tensors in that source.

    Note:
        Results may be truncated if server has max_list_flights_results configured.
        Check schema metadata for truncation info (truncated=True indicates
        more sources exist on server than were returned).
    """
    source_descriptors = {}
    truncated = False
    total_sources = None

    for info in self._client.list_flights(options=self._call_options):
        source_desc = DataSourceDescriptor.FromString(info.descriptor.command)
        source_descriptors[source_desc.source_id] = source_desc
        # Cache tensor descriptors
        for tensor_desc in source_desc.tensors:
            self._descriptors[
                self._descriptor_key(source_desc.source_id, tensor_desc.array_id)
            ] = tensor_desc

        # Check schema metadata for truncation info
        if info.schema.metadata:
            truncated_bytes = info.schema.metadata.get(b"truncated")
            if truncated_bytes:
                truncated = truncated_bytes.decode() == "True"
            total_sources_bytes = info.schema.metadata.get(b"total_sources")
            if total_sources_bytes:
                total_sources = int(total_sources_bytes.decode())

    self._sources = source_descriptors

    if truncated and total_sources:
        logger.warning(
            f"list_sources: returned {len(source_descriptors)} of {total_sources} sources (truncated)"
        )
    else:
        logger.info(f"list_sources: returned {len(source_descriptors)} sources")

    return source_descriptors

query_sources

query_sources(sql: str, *, format: str = 'arrow') -> Any

Execute SQL query against server's source metadata database.

The server-side metadata database is mandatory (biopb/biopb#225), so any standard tensor-server supports this. Only an embedded server explicitly constructed without a metadata database rejects the query.

Parameters:

Name Type Description Default
sql str

SQL query (e.g., "SELECT source_id, source_type FROM sources WHERE dtype='uint16'")

required
format str

Shape of the returned result:

  • "arrow" (default) — a pyarrow.Table. This is the historical return type; the default is unchanged for backward compatibility. Zero-copy, and the only format that preserves the schema metadata described under Note.
  • "pandas" — a pandas.DataFrame (requires pandas).
  • "records" — a list[dict], one dict per row.
'arrow'

Returns:

Type Description
Any

The query result in the requested format; an empty query

Any

returns an empty object of that same type. For "pandas" and

Any

"records" the usual Arrow->Python coercion applies (list

Any

columns such as shape_summary become Python lists / object

Any

dtype, and nullable integer columns may widen to float). For

Any

"pandas", NULLs in string columns (e.g. metadata_json) are

Any

normalized to None rather than the truthy float NaN Arrow

Any

would otherwise produce, so if row.metadata_json: behaves as

Any

expected.

Note

The server reports truncation via schema metadata (total_sources / returned_sources). Those keys survive only on the "arrow" result; for every format truncation is also surfaced via a logged INFO line.

Raises:

Type Description
ValueError

If format is not one of the supported values. (SQL validation -- forbidden keywords / disallowed tables -- happens server-side and surfaces as a Flight error, below, not a client-side ValueError.)

ImportError

If format="pandas" but pandas is not installed.

FlightServerError

If the server has no metadata database enabled, or rejects the query (e.g. forbidden keywords / disallowed tables).

Example

client = TensorFlightClient('grpc://localhost:8815') table = client.query_sources("SELECT source_id FROM sources WHERE source_type='ome-zarr'") table.to_pandas() # or pass format="pandas" to get a DataFrame

Source code in src/main/python/biopb/tensor/client.py
def query_sources(self, sql: str, *, format: str = "arrow") -> Any:
    """Execute SQL query against server's source metadata database.

    The server-side metadata database is mandatory (biopb/biopb#225), so any
    standard tensor-server supports this. Only an embedded server explicitly
    constructed without a metadata database rejects the query.

    Args:
        sql: SQL query (e.g., "SELECT source_id, source_type FROM sources WHERE dtype='uint16'")
        format: Shape of the returned result:

            - ``"arrow"`` (default) — a ``pyarrow.Table``. This is the
              historical return type; the default is unchanged for backward
              compatibility. Zero-copy, and the only format that preserves
              the schema metadata described under *Note*.
            - ``"pandas"`` — a ``pandas.DataFrame`` (requires pandas).
            - ``"records"`` — a ``list[dict]``, one dict per row.

    Returns:
        The query result in the requested ``format``; an empty query
        returns an empty object of that same type. For ``"pandas"`` and
        ``"records"`` the usual Arrow->Python coercion applies (list
        columns such as ``shape_summary`` become Python lists / object
        dtype, and nullable integer columns may widen to float). For
        ``"pandas"``, NULLs in string columns (e.g. ``metadata_json``) are
        normalized to ``None`` rather than the truthy float ``NaN`` Arrow
        would otherwise produce, so ``if row.metadata_json:`` behaves as
        expected.

    Note:
        The server reports truncation via schema metadata
        (``total_sources`` / ``returned_sources``). Those keys survive only
        on the ``"arrow"`` result; for every format truncation is also
        surfaced via a logged INFO line.

    Raises:
        ValueError: If *format* is not one of the supported values. (SQL
            validation -- forbidden keywords / disallowed tables -- happens
            server-side and surfaces as a Flight error, below, not a
            client-side ValueError.)
        ImportError: If ``format="pandas"`` but pandas is not installed.
        FlightServerError: If the server has no metadata database enabled,
            or rejects the query (e.g. forbidden keywords / disallowed
            tables).

    Example:
        >>> client = TensorFlightClient('grpc://localhost:8815')
        >>> table = client.query_sources("SELECT source_id FROM sources WHERE source_type='ome-zarr'")
        >>> table.to_pandas()  # or pass format="pandas" to get a DataFrame
    """
    if format not in ("pandas", "arrow", "records"):
        raise ValueError(
            f"query_sources: unknown format {format!r}; "
            "expected 'pandas', 'arrow', or 'records'"
        )

    cmd = FlightCmd(
        source_id="__metadata_query__",
        metadata_query=MetadataQueryOption(sql=sql),
    )
    descriptor = flight.FlightDescriptor.for_command(cmd.SerializeToString())
    info = self._client.get_flight_info(descriptor, options=self._call_options)

    # Check schema metadata for truncation info
    if info.schema.metadata:
        total_sources = info.schema.metadata.get(b"total_sources")
        if total_sources:
            total = int(total_sources.decode())
            returned = info.schema.metadata.get(b"returned_sources")
            if returned:
                returned_count = int(returned.decode())
                if returned_count < total:
                    logger.info(
                        f"query_sources: returned {returned_count} of {total} sources (truncated)"
                    )
                else:
                    logger.info(f"query_sources: returned {returned_count} sources")

    # Fetch results via DoGet
    if info.endpoints:
        reader = self._client.do_get(
            info.endpoints[0].ticket, options=self._call_options
        )
        table = reader.read_all()
    else:
        # Empty result
        table = info.schema.empty_table()

    return self._format_query_result(table, format)

get_source_metadata

get_source_metadata(source_id: str) -> dict

Get source-level OME/vendor metadata as a dict.

Parameters:

Name Type Description Default
source_id str

Source identifier

required

Returns:

Type Description
dict

The source's metadata dict (the format-specific OME/vendor metadata),

dict

or an empty dict if the source carries none.

Raises:

Type Description
ValueError

If the source is unknown, or unresolved (cloud / synced-folder) -- call :meth:resolve first.

Source code in src/main/python/biopb/tensor/client.py
def get_source_metadata(self, source_id: str) -> dict:
    """Get source-level OME/vendor metadata as a dict.

    Args:
        source_id: Source identifier

    Returns:
        The source's metadata dict (the format-specific OME/vendor metadata),
        or an empty dict if the source carries none.

    Raises:
        ValueError: If the source is unknown, or unresolved (cloud /
            synced-folder) -- call :meth:`resolve` first.
    """
    import json

    if source_id not in self._sources:
        self.list_sources()

    source_desc = self._sources.get(source_id)
    if source_desc is None:
        raise ValueError(f"Source not found: {source_id}")

    if not source_desc.tensors:
        # Unresolved (cloud / synced-folder) source: tensors are unknown
        # until resolve. Don't silently return {} -- that conflates
        # "unresolved" with "resolved, no metadata" (the line below). Steer
        # the caller to the explicit, consented resolve() instead, matching
        # get_physical_scale / get_tensor (#108). Crucially this stays a
        # cheap read: it must NOT silently recall the whole file the way a
        # resolve-on-serve probe (get_descriptor) would.
        raise _unresolved_source_error(source_id)

    # metadata_json is populated on the descriptor GetFlightInfo returns, so
    # we fetch it via the source's first tensor.
    first_tensor = source_desc.tensors[0]
    cmd = FlightCmd(
        source_id=source_id,
        tensor_read=TensorReadOption(
            tensor_id=first_tensor.array_id,
            with_metadata=True,
        ),
    )
    flight_desc = flight.FlightDescriptor.for_command(cmd.SerializeToString())
    info = self._client.get_flight_info(flight_desc, options=self._call_options)
    response_desc = TensorDescriptor.FromString(info.descriptor.command)

    if response_desc.metadata_json:
        # The server wraps it as {"type": ..., "dim_label": [...],
        # "metadata": {...}}; return just the inner metadata dict.
        wrapped = json.loads(response_desc.metadata_json)
        return wrapped.get("metadata", {})
    return {}

get_physical_scale

get_physical_scale(
    array_id: Optional[str] = None,
    tensor_id: Optional[str] = None,
    *,
    source_id: Optional[str] = None
) -> Optional[Tuple[List[float], List[str]]]

Per-dimension physical pixel size + unit for a tensor.

Returns (scale, unit): two lists aligned with the tensor's dim_labels (source axis order), or None when no physical sizes are known (an older server, or a format that carries none).

physical_scale/physical_unit are TensorDescriptor fields the server fills on every GetFlightInfo (issue #31), so this reads the descriptor a prior :meth:get_tensor already cached -- no extra RPC when it is cached, and it never requests the opt-in metadata_json field on that same descriptor. (Contrast :meth:get_source_metadata, which forces with_metadata to ship the whole OME tree; do not dig physical sizes out of that -- this is the compact projection meant for display scale.)

Parameters:

Name Type Description Default
array_id Optional[str]

Globally-unique tensor id (identity policy) -- e.g. "zarr_a3f2" or "aics_7f3/Image:0". A bare single-tensor source id resolves to its sole tensor. A bare multi-tensor source id anchors on the source's default (first) tensor -- unlike get_tensor, which requires the field be named; pass the qualified source_id/field to target a specific scene.

None
tensor_id Optional[str]

DEPRECATED. The legacy (source_id, tensor_id) form; pass the array_id as the single first argument instead.

None

Returns:

Type Description
Optional[Tuple[List[float], List[str]]]

(scale, unit) lists, or None if no physical scale is known.

Source code in src/main/python/biopb/tensor/client.py
def get_physical_scale(
    self,
    array_id: Optional[str] = None,
    tensor_id: Optional[str] = None,
    *,
    source_id: Optional[str] = None,
) -> Optional[Tuple[List[float], List[str]]]:
    """Per-dimension physical pixel size + unit for a tensor.

    Returns ``(scale, unit)``: two lists aligned with the tensor's
    ``dim_labels`` (source axis order), or ``None`` when no physical sizes
    are known (an older server, or a format that carries none).

    ``physical_scale``/``physical_unit`` are ``TensorDescriptor`` fields the
    server fills on every ``GetFlightInfo`` (issue #31), so this reads the
    descriptor a prior :meth:`get_tensor` already cached -- no extra RPC when
    it is cached, and it never requests the opt-in ``metadata_json`` field on
    that same descriptor. (Contrast :meth:`get_source_metadata`, which forces
    ``with_metadata`` to ship the whole OME tree; do not dig physical sizes
    out of that -- this is the compact projection meant for display scale.)

    Args:
        array_id: Globally-unique tensor id (identity policy) -- e.g.
            ``"zarr_a3f2"`` or ``"aics_7f3/Image:0"``. A bare single-tensor
            source id resolves to its sole tensor. A bare *multi*-tensor
            source id anchors on the source's default (first) tensor --
            unlike ``get_tensor``, which requires the field be named; pass the
            qualified ``source_id/field`` to target a specific scene.
        tensor_id: DEPRECATED. The legacy ``(source_id, tensor_id)`` form;
            pass the array_id as the single first argument instead.

    Returns:
        ``(scale, unit)`` lists, or ``None`` if no physical scale is known.
    """
    source_id, tensor_id = self._resolve_array_id(
        array_id, tensor_id, "get_physical_scale", source_id=source_id
    )
    desc = (
        self._descriptors.get(self._descriptor_key(source_id, tensor_id))
        if tensor_id
        else None
    )
    if desc is None:
        # Don't silently recall (download) a whole cloud file just to read its
        # pixel size: if the source is known-unresolved, steer the caller to
        # resolve() explicitly -- consistent with get_tensor, and faithful to
        # resolution being a consented act, not a side effect of a metadata
        # probe. (Only catches sources already in the catalog cache; a
        # never-listed id still falls through to the fetch below, same as
        # every other entry point.)
        cached = self._sources.get(source_id)
        if cached is not None and not cached.tensors:
            raise _unresolved_source_error(source_id)
        # tensor_id None -> the source's default (first) tensor. A real fetch
        # error (server unreachable, source not found) propagates to the
        # caller -- it must stay distinguishable from "no physical scale
        # recorded", which is the only case that yields None. This matches
        # the pre-#75 contract, where the get_source() fetch was unguarded.
        desc = self._fetch_tensor_descriptor(source_id, tensor_id)
    if not desc.physical_scale:
        return None
    return list(desc.physical_scale), list(desc.physical_unit)

get_descriptor

get_descriptor(array_id: str) -> TensorDescriptor

Fetch one tensor's TensorDescriptor by its globally-unique array_id.

A tensor is identified by its array_id alone (see the tensor identity policy at the top of proto/biopb/tensor/descriptor.proto), so this takes that one identifier rather than a (source_id, tensor_id) pair. Works even when the source is beyond the (truncatable) list_sources() cap, and the result is cached. Passing a bare source_id (single-tensor source, or to anchor on a multi-tensor source's default/first tensor) is accepted. To enumerate ALL tensors/scenes of a source, use list_sources()[source_id].tensors -- NOT this method.

This is a cheap probe -- it does NOT resolve. On an unresolved (cloud / synced-folder) source it raises an error pointing at :meth:resolve, never triggering a download. Call :meth:resolve first to read such a source.

Parameters:

Name Type Description Default
array_id str

Globally-unique tensor id, e.g. "zarr_a3f2" (single- tensor source) or "aics_7f3/Image:0" (multi-tensor source).

required

Returns:

Type Description
TensorDescriptor

The TensorDescriptor for that tensor.

Source code in src/main/python/biopb/tensor/client.py
def get_descriptor(self, array_id: str) -> "TensorDescriptor":
    """Fetch one tensor's ``TensorDescriptor`` by its globally-unique array_id.

    A tensor is identified by its ``array_id`` alone (see the tensor identity
    policy at the top of ``proto/biopb/tensor/descriptor.proto``), so this
    takes that one identifier rather than a ``(source_id, tensor_id)`` pair.
    Works even when the source is beyond the (truncatable) ``list_sources()``
    cap, and the result is cached. Passing a bare ``source_id`` (single-tensor
    source, or to anchor on a multi-tensor source's default/first tensor) is
    accepted. To enumerate ALL tensors/scenes of a source, use
    ``list_sources()[source_id].tensors`` -- NOT this method.

    This is a cheap probe -- it does NOT resolve. On an unresolved (cloud /
    synced-folder) source it raises an error pointing at :meth:`resolve`,
    never triggering a download. Call :meth:`resolve` first to read such a
    source.

    Args:
        array_id: Globally-unique tensor id, e.g. ``"zarr_a3f2"`` (single-
            tensor source) or ``"aics_7f3/Image:0"`` (multi-tensor source).

    Returns:
        The ``TensorDescriptor`` for that tensor.
    """
    # source_id is the slash-free prefix; the full array_id is the tensor_id.
    source_id = array_id.split("/", 1)[0]
    return self._fetch_tensor_descriptor(source_id, array_id)

resolve

resolve(
    source_id: str,
    *,
    on_progress: Optional[
        Callable[[ResolveProgress], None]
    ] = None,
    should_cancel: Optional[Callable[[], bool]] = None
) -> DataSourceDescriptor

Resolve an unresolved source and return its full DataSourceDescriptor.

.. note:: Experimental. Cloud / remote source support (unresolved sources, resolve, and :meth:warm) is experimental and its behavior may change.

An unresolved source is catalogued by URL only -- its shape/dtype/field list are unknown until first access (it lists with data_resident False and an empty list_sources()[source_id].tensors). The canonical case is a cloud / synced-folder ("Files-On-Demand") source.

Resolving asks the server to hydrate it. For a dehydrated placeholder this downloads the whole file -- a recall that can take minutes, consume local disk, and fail when offline -- then reads its real shape, dtype, and field list. This is the heavyweight, consenting operation that catalog browsing (:meth:list_sources / :meth:query_sources) deliberately avoids; call it only when you intend to read the data. After it returns, :meth:get_tensor and friends work normally.

Idempotent: resolving an already-resolved source just re-fetches it.

Parameters:

Name Type Description Default
source_id str

The source to resolve (e.g. "onedrive_a3f2").

required
on_progress Optional[Callable[[ResolveProgress], None]]

Optional callback invoked with a ResolveProgress (elapsed seconds, target name, target size in bytes) on each server heartbeat, so a caller can display progress. Called on the calling thread; keep it cheap and non-blocking.

None
should_cancel Optional[Callable[[], bool]]

Optional predicate polled on each heartbeat; when it returns True the client stops consuming the stream and raises :class:ResolveCancelled. The server-side recall continues to completion and is cached, so a later resolve reuses it.

None

Returns:

Type Description
DataSourceDescriptor

The full DataSourceDescriptor with every tensor/field enumerated

DataSourceDescriptor

-- the complete field set in one call, regardless of catalog size.

Raises:

Type Description
ResolveCancelled

if should_cancel asked to stop mid-resolve.

Source code in src/main/python/biopb/tensor/client.py
def resolve(
    self,
    source_id: str,
    *,
    on_progress: Optional[Callable[["ResolveProgress"], None]] = None,
    should_cancel: Optional[Callable[[], bool]] = None,
) -> "DataSourceDescriptor":
    """Resolve an unresolved source and return its full ``DataSourceDescriptor``.

    .. note:: Experimental. Cloud / remote source support (unresolved sources,
       resolve, and :meth:`warm`) is experimental and its behavior may change.

    An *unresolved* source is catalogued by URL only -- its shape/dtype/field
    list are unknown until first access (it lists with ``data_resident`` False
    and an empty ``list_sources()[source_id].tensors``). The canonical case is
    a cloud / synced-folder ("Files-On-Demand") source.

    Resolving asks the server to hydrate it. For a dehydrated placeholder this
    **downloads the whole file** -- a recall that can take minutes, consume
    local disk, and fail when offline -- then reads its real shape, dtype, and
    field list. This is the heavyweight, *consenting* operation that catalog
    browsing (:meth:`list_sources` / :meth:`query_sources`) deliberately
    avoids; call it only when you intend to read the data. After it returns,
    :meth:`get_tensor` and friends work normally.

    Idempotent: resolving an already-resolved source just re-fetches it.

    Args:
        source_id: The source to resolve (e.g. ``"onedrive_a3f2"``).
        on_progress: Optional callback invoked with a ``ResolveProgress``
            (elapsed seconds, target name, target size in bytes) on each
            server heartbeat, so a caller can display progress. Called on the
            calling thread; keep it cheap and non-blocking.
        should_cancel: Optional predicate polled on each heartbeat; when it
            returns True the client stops consuming the stream and raises
            :class:`ResolveCancelled`. The server-side recall continues to
            completion and is cached, so a later ``resolve`` reuses it.

    Returns:
        The full ``DataSourceDescriptor`` with every tensor/field enumerated
        -- the complete field set in one call, regardless of catalog size.

    Raises:
        ResolveCancelled: if ``should_cancel`` asked to stop mid-resolve.
    """
    # One dedicated, streaming ``resolve`` action: it is the SINGLE server
    # entry point that performs the (possibly minutes-long) recall, and it
    # returns the full DataSourceDescriptor directly -- no GetFlightInfo +
    # list_sources two-step, so no truncation hole for multi-field sources
    # beyond the list cap. The action streams ``ResolveStreamMessage``
    # heartbeats (a ``progress`` arm) to keep the connection warm under proxy
    # idle timeouts; the single terminal message carries the descriptor in
    # its ``result`` arm. ``should_cancel`` / ``on_progress`` are polled once
    # per received message, i.e. roughly once per server heartbeat.
    action = flight.Action("resolve", source_id.encode("utf-8"))
    desc: Optional[DataSourceDescriptor] = None
    for which, msg, body in self._iter_action_messages(
        action, ResolveStreamMessage
    ):
        if should_cancel is not None and should_cancel():
            raise ResolveCancelled(f"resolve('{source_id}') cancelled by caller")
        if which == "progress":
            if on_progress is not None:
                on_progress(msg.progress)
        elif which == "result":
            desc = DataSourceDescriptor()
            desc.CopyFrom(msg.result)
        else:
            # Legacy server: a non-empty body IS a bare serialized
            # DataSourceDescriptor (pre-envelope protocol).
            desc = DataSourceDescriptor.FromString(body)
    if desc is None:
        raise RuntimeError(
            f"resolve('{source_id}') returned no descriptor "
            "(server closed the stream without a result)"
        )
    self._sources[source_id] = desc
    return desc

warm

warm(
    source_id: str,
    *,
    on_progress: Optional[
        Callable[[WarmProgress], None]
    ] = None,
    should_cancel: Optional[Callable[[], bool]] = None
) -> WarmProgress

Hydrate-ahead: recall a resolved source's member files on the server.

.. note:: Experimental. Cloud / remote source support (:meth:resolve and this hydrate-ahead path) is experimental and its behavior may change.

:meth:resolve populates a source's metadata but, for a multi-file cloud source (zarr / ome-zarr / ndtiff / tiff-sequence / micromanager), leaves the bulk pixel data dehydrated -- each member file then recalls one-at-a-time, slowly, the first time a read touches it (the viewer scrubbing planes is the worst case). warm opts into pulling them all resident up front so later reads never stall.

The recall happens entirely server-side (the server walks the source directory and reads each file to force the sync engine's recall); no pixels cross the wire, only progress. It is idempotent -- already-resident files are cheap local reads -- so a warm re-run after a cancel simply finishes the remainder. Only meaningful for multi-file sources; a single-file source returns immediately (resolve already recalled it).

Parameters:

Name Type Description Default
source_id str

The (already-resolved) source to warm.

required
on_progress Optional[Callable[[WarmProgress], None]]

Optional callback invoked with a WarmProgress (files/bytes done vs total, current file name, elapsed) on each progress message. Called on the calling thread; keep it cheap.

None
should_cancel Optional[Callable[[], bool]]

Optional predicate polled per message; when it returns True the client closes the stream -- which the server observes and stops the recall promptly -- and this raises :class:ResolveCancelled. Files already recalled stay resident.

None

Returns:

Type Description
WarmProgress

The terminal WarmProgress snapshot (files_done /

WarmProgress

bytes_done reflect what was made resident; on a no-op source

WarmProgress

files_total == 0).

Raises:

Type Description
ResolveCancelled

if should_cancel asked to stop mid-warm.

RuntimeError

if the server predates the warm action (too old for hydrate-ahead), or closes the stream without a terminal status.

Source code in src/main/python/biopb/tensor/client.py
def warm(
    self,
    source_id: str,
    *,
    on_progress: Optional[Callable[["WarmProgress"], None]] = None,
    should_cancel: Optional[Callable[[], bool]] = None,
) -> "WarmProgress":
    """Hydrate-ahead: recall a resolved source's member files on the server.

    .. note:: Experimental. Cloud / remote source support (:meth:`resolve` and
       this hydrate-ahead path) is experimental and its behavior may change.

    :meth:`resolve` populates a source's *metadata* but, for a multi-file
    cloud source (zarr / ome-zarr / ndtiff / tiff-sequence / micromanager),
    leaves the bulk pixel data dehydrated -- each member file then recalls
    one-at-a-time, slowly, the first time a read touches it (the viewer
    scrubbing planes is the worst case). ``warm`` opts into pulling them all
    resident up front so later reads never stall.

    The recall happens **entirely server-side** (the server walks the source
    directory and reads each file to force the sync engine's recall); no
    pixels cross the wire, only progress. It is idempotent -- already-resident
    files are cheap local reads -- so a ``warm`` re-run after a cancel simply
    finishes the remainder. Only meaningful for multi-file sources; a
    single-file source returns immediately (resolve already recalled it).

    Args:
        source_id: The (already-resolved) source to warm.
        on_progress: Optional callback invoked with a ``WarmProgress``
            (files/bytes done vs total, current file name, elapsed) on each
            progress message. Called on the calling thread; keep it cheap.
        should_cancel: Optional predicate polled per message; when it returns
            True the client closes the stream -- which the server observes and
            stops the recall promptly -- and this raises
            :class:`ResolveCancelled`. Files already recalled stay resident.

    Returns:
        The terminal ``WarmProgress`` snapshot (``files_done`` /
        ``bytes_done`` reflect what was made resident; on a no-op source
        ``files_total == 0``).

    Raises:
        ResolveCancelled: if ``should_cancel`` asked to stop mid-warm.
        RuntimeError: if the server predates the ``warm`` action (too old for
            hydrate-ahead), or closes the stream without a terminal status.
    """
    action = flight.Action("warm", source_id.encode("utf-8"))
    done: Optional[WarmProgress] = None
    unknown = (
        "Hydrate-ahead is unavailable: the tensor server is too old "
        "to support the 'warm' action. Upgrade the server, or just "
        "read the data on demand (it will recall lazily)."
    )
    for which, msg, _ in self._iter_action_messages(
        action, WarmStreamMessage, unknown_action_msg=unknown
    ):
        if should_cancel is not None and should_cancel():
            raise ResolveCancelled(f"warm('{source_id}') cancelled by caller")
        if which == "progress":
            if on_progress is not None:
                on_progress(msg.progress)
        elif which == "done":
            done = WarmProgress()
            done.CopyFrom(msg.done)
    if done is None:
        raise RuntimeError(
            f"warm('{source_id}') returned no terminal status "
            "(server closed the stream without a 'done')"
        )
    return done

add_source

add_source(
    url: str,
    *,
    source_type: str = "",
    dim_labels: Optional[List[str]] = None,
    on_progress: Optional[
        Callable[[AddSourceProgress], None]
    ] = None,
    should_cancel: Optional[Callable[[], bool]] = None
) -> AddSourceResult

Register a local path on the SERVER as a served source at runtime.

This is the wire entrypoint behind the tensor-browser's drag-drop: it hands the server a filesystem path (or directory) that it interprets on its own filesystem, and the server routes it through the same claim -> adapter -> catalog pipeline the directory watcher uses. A dropped directory that is not itself a dataset is walked recursively and may register several sources, so the action streams progress and a final tally rather than returning a single descriptor.

The path must exist on the server. Because a dropped directory's walk has no known size up front, there is no percentage -- progress is a running count of sources registered so far.

Parameters:

Name Type Description Default
url str

Absolute path (or directory) on the server's filesystem.

required
source_type str

Explicit adapter type (e.g. "zarr", "ome-zarr"); empty means auto-detect via the adapters' claim protocol.

''
dim_labels Optional[List[str]]

Optional dimension labels for the registered tensor(s).

None
on_progress Optional[Callable[[AddSourceProgress], None]]

Optional callback invoked with an AddSourceProgress (count + current path + last descriptor) per source as it registers. Called on the calling thread; keep it cheap.

None
should_cancel Optional[Callable[[], bool]]

Optional predicate polled per message; when it returns True the client closes the stream, which the server observes and stops discovery -- sources already registered stay registered.

None

Returns:

Type Description
AddSourceResult

The terminal AddSourceResult (added descriptors,

AddSourceResult

already_present source_ids, failed (path, reason) pairs).

AddSourceResult

A directory dropped above the large-scan threshold comes back as a

AddSourceResult

failed entry, not a special flag.

Raises:

Type Description
FlightServerError

whole-request failure (path not found / unreadable on the server, or the server declines the request).

RuntimeError

the server predates the add_source action, or closed the stream without a terminal result.

Source code in src/main/python/biopb/tensor/client.py
def add_source(
    self,
    url: str,
    *,
    source_type: str = "",
    dim_labels: Optional[List[str]] = None,
    on_progress: Optional[Callable[["AddSourceProgress"], None]] = None,
    should_cancel: Optional[Callable[[], bool]] = None,
) -> "AddSourceResult":
    """Register a local path on the SERVER as a served source at runtime.

    This is the wire entrypoint behind the tensor-browser's drag-drop: it
    hands the server a filesystem path (or directory) that it interprets on
    *its own* filesystem, and the server routes it through the same claim ->
    adapter -> catalog pipeline the directory watcher uses. A dropped
    directory that is not itself a dataset is walked recursively and may
    register several sources, so the action streams progress and a final
    tally rather than returning a single descriptor.

    The path must exist on the server. Because a dropped directory's walk has
    no known size up front, there is no percentage -- progress is a running
    count of sources registered so far.

    Args:
        url: Absolute path (or directory) on the server's filesystem.
        source_type: Explicit adapter type (e.g. ``"zarr"``, ``"ome-zarr"``);
            empty means auto-detect via the adapters' claim protocol.
        dim_labels: Optional dimension labels for the registered tensor(s).
        on_progress: Optional callback invoked with an ``AddSourceProgress``
            (count + current path + last descriptor) per source as it
            registers. Called on the calling thread; keep it cheap.
        should_cancel: Optional predicate polled per message; when it returns
            True the client closes the stream, which the server observes and
            stops discovery -- sources already registered stay registered.

    Returns:
        The terminal ``AddSourceResult`` (``added`` descriptors,
        ``already_present`` source_ids, ``failed`` ``(path, reason)`` pairs).
        A directory dropped above the large-scan threshold comes back as a
        ``failed`` entry, not a special flag.

    Raises:
        flight.FlightServerError: whole-request failure (path not found /
            unreadable on the server, or the server declines the request).
        RuntimeError: the server predates the ``add_source`` action, or
            closed the stream without a terminal result.
    """
    req = AddSourceRequest(
        url=url,
        source_type=source_type,
        dim_labels=dim_labels or [],
    )
    action = flight.Action("add_source", req.SerializeToString())
    unknown = (
        "Runtime source registration is unavailable: the tensor "
        "server is too old to support the 'add_source' action. "
        "Upgrade the server, or add the source via its config file."
    )
    result: Optional[AddSourceResult] = None
    for which, msg, _ in self._iter_action_messages(
        action, AddSourceStreamMessage, unknown_action_msg=unknown
    ):
        if which == "progress":
            if on_progress is not None:
                on_progress(msg.progress)
        elif which == "result":
            result = AddSourceResult()
            result.CopyFrom(msg.result)
        # Poll AFTER consuming this message, not before: a cancel landing
        # exactly on the terminal ``result`` must not discard a completed
        # tally already captured above (issue #4). Closing the stream keeps
        # everything already registered server-side.
        if should_cancel is not None and should_cancel():
            break
    if result is None:
        # A caller-driven cancel breaks before the terminal result; report an
        # empty tally rather than an error (the cancel was intentional).
        if should_cancel is not None and should_cancel():
            return AddSourceResult()
        raise RuntimeError(
            f"add_source('{url}') returned no terminal result "
            "(server closed the stream without a result)"
        )
    return result

remove_source

remove_source(root_url: str) -> RemoveSourceResult

Deregister a drag-dropped source branch on the SERVER at runtime.

The narrow counterpart to :meth:add_source: it removes ONLY drag-dropped sources, which the server identifies by the dnd:// origin scheme on their catalog source_url. root_url is such a branch root (a dnd://... value); every source at or under it is removed as a unit. A non-dnd:// root_url is refused by the server.

Parameters:

Name Type Description Default
root_url str

The dnd:// branch root to remove (from the browser's dropped-root node).

required

Returns:

Type Description
RemoveSourceResult

A RemoveSourceResult with removed (source_ids) and failed

RemoveSourceResult

(AddSourceFailure whose path carries the source_id).

Raises:

Type Description
FlightServerError

the server refused the request (e.g. a non-dnd:// root, or removal not enabled).

RuntimeError

the server predates the remove_source action, or returned no result.

Source code in src/main/python/biopb/tensor/client.py
def remove_source(self, root_url: str) -> "RemoveSourceResult":
    """Deregister a drag-dropped source branch on the SERVER at runtime.

    The narrow counterpart to :meth:`add_source`: it removes ONLY
    drag-dropped sources, which the server identifies by the ``dnd://``
    origin scheme on their catalog ``source_url``. ``root_url`` is such a
    branch root (a ``dnd://...`` value); every source at or under it is
    removed as a unit. A non-``dnd://`` ``root_url`` is refused by the server.

    Args:
        root_url: The ``dnd://`` branch root to remove (from the browser's
            dropped-root node).

    Returns:
        A ``RemoveSourceResult`` with ``removed`` (source_ids) and ``failed``
        (``AddSourceFailure`` whose ``path`` carries the source_id).

    Raises:
        flight.FlightServerError: the server refused the request (e.g. a
            non-``dnd://`` root, or removal not enabled).
        RuntimeError: the server predates the ``remove_source`` action, or
            returned no result.
    """
    req = RemoveSourceRequest(root_url=root_url)
    action = flight.Action("remove_source", req.SerializeToString())
    try:
        results = self._client.do_action(action, options=self._call_options)
        result_bytes = next(results)
    except flight.FlightError as exc:
        if "Unknown action" in str(exc):
            raise RuntimeError(
                "Source removal is unavailable: the tensor server is too old "
                "to support the 'remove_source' action. Upgrade the server."
            ) from exc
        raise
    except StopIteration as exc:
        raise RuntimeError(
            f"remove_source('{root_url}') returned no result"
        ) from exc
    return RemoveSourceResult.FromString(result_bytes.body.to_pybytes())

get_source

get_source(
    source_id: str, tensor_id: Optional[str] = None
) -> DataSourceDescriptor

DEPRECATED -- use :meth:get_descriptor (one tensor) or :meth:list_sources (all tensors of a source) instead.

This method is inconsistent with the tensor identity policy (proto/biopb/tensor/descriptor.proto): it splits a tensor's identity into a (source_id, tensor_id) pair, whereas a tensor is identified by its globally-unique array_id alone. It is also misnamed -- despite returning a DataSourceDescriptor, it is a single-tensor probe: the returned .tensors holds ONLY the resolved tensor (or the source's default/first tensor when tensor_id is None), never the full scene list. For a multi-scene source, get_source(id).tensors is length-1 and scenes 2..N are NOT enumerated here. Use list_sources()[id].tensors for the complete enumeration.

Parameters:

Name Type Description Default
source_id str

Data source identifier.

required
tensor_id Optional[str]

Optional tensor to anchor the lookup; None -> the source's default (first) tensor.

None

Returns:

Type Description
DataSourceDescriptor

DataSourceDescriptor wrapping the single resolved TensorDescriptor.

Source code in src/main/python/biopb/tensor/client.py
def get_source(
    self,
    source_id: str,
    tensor_id: Optional[str] = None,
) -> "DataSourceDescriptor":
    """DEPRECATED -- use :meth:`get_descriptor` (one tensor) or
    :meth:`list_sources` (all tensors of a source) instead.

    This method is inconsistent with the tensor identity policy
    (``proto/biopb/tensor/descriptor.proto``): it splits a tensor's identity
    into a ``(source_id, tensor_id)`` pair, whereas a tensor is identified by
    its globally-unique ``array_id`` alone. It is also misnamed -- despite
    returning a ``DataSourceDescriptor``, it is a single-tensor probe: the
    returned ``.tensors`` holds ONLY the resolved tensor (or the source's
    default/first tensor when ``tensor_id`` is None), never the full scene
    list. For a multi-scene source, ``get_source(id).tensors`` is length-1
    and scenes 2..N are NOT enumerated here. Use
    ``list_sources()[id].tensors`` for the complete enumeration.

    Args:
        source_id: Data source identifier.
        tensor_id: Optional tensor to anchor the lookup; None -> the source's
            default (first) tensor.

    Returns:
        DataSourceDescriptor wrapping the single resolved TensorDescriptor.
    """
    warnings.warn(
        "TensorFlightClient.get_source() is deprecated and will be removed in "
        "a future release: it is inconsistent with the array_id identity "
        "policy and returns only a single tensor, not a source's full scene "
        "list. Use get_descriptor(array_id) for one tensor, or list_sources() "
        "to enumerate all tensors of a source.",
        DeprecationWarning,
        stacklevel=2,
    )
    tensor_desc = self._fetch_tensor_descriptor(source_id, tensor_id)
    source_desc = DataSourceDescriptor(source_id=source_id, tensors=[tensor_desc])
    # Do NOT clobber a full enumeration previously cached by list_sources();
    # only seed _sources when this source is otherwise unknown (issue #75).
    if source_id not in self._sources:
        self._sources[source_id] = source_desc
    return source_desc

get_tensor

get_tensor(
    array_id: Optional[str] = None,
    tensor_id: Optional[str] = None,
    slice_hint: Optional[Tuple[slice, ...]] = None,
    scale_hint: Optional[Sequence[int]] = None,
    reduction_method: Optional[str] = None,
    *,
    source_id: Optional[str] = None
) -> da.Array

Get a lazy dask array for a tensor, addressed by its array_id.

Parameters:

Name Type Description Default
array_id Optional[str]

Globally-unique tensor id (identity policy) -- e.g. "zarr_a3f2" for a single-tensor source or "aics_7f3/Image:0" for a multi-tensor source.

None
tensor_id Optional[str]

DEPRECATED. The legacy (source_id, tensor_id) form; pass the array_id as the single first argument instead.

None
slice_hint Optional[Tuple[slice, ...]]

Optional slice tuple to filter chunks

None
scale_hint Optional[Sequence[int]]

Optional per-dimension integer downsampling factors

None
reduction_method Optional[str]

Optional dynamic reduction method for scaled reads

None

Returns:

Type Description
Array

dask.array with lazy chunk loading

Raises:

Type Description
ValueError

If source not found, tensor not found, or a bare multi-tensor source id is given without a within-source field

Source code in src/main/python/biopb/tensor/client.py
def get_tensor(
    self,
    array_id: Optional[str] = None,
    tensor_id: Optional[str] = None,
    slice_hint: Optional[Tuple[slice, ...]] = None,
    scale_hint: Optional[Sequence[int]] = None,
    reduction_method: Optional[str] = None,
    *,
    source_id: Optional[str] = None,
) -> da.Array:
    """Get a lazy dask array for a tensor, addressed by its array_id.

    Args:
        array_id: Globally-unique tensor id (identity policy) -- e.g.
            ``"zarr_a3f2"`` for a single-tensor source or
            ``"aics_7f3/Image:0"`` for a multi-tensor source.
        tensor_id: DEPRECATED. The legacy ``(source_id, tensor_id)`` form;
            pass the array_id as the single first argument instead.
        slice_hint: Optional slice tuple to filter chunks
        scale_hint: Optional per-dimension integer downsampling factors
        reduction_method: Optional dynamic reduction method for scaled reads

    Returns:
        dask.array with lazy chunk loading

    Raises:
        ValueError: If source not found, tensor not found, or a bare
            multi-tensor source id is given without a within-source field
    """
    source_id, tensor_id = self._resolve_array_id(
        array_id, tensor_id, "get_tensor", source_id=source_id
    )
    logger.debug(f"get_tensor: source_id={source_id}, tensor_id={tensor_id}")

    # Get flight info context
    ctx = self._get_tensor_context(
        source_id=source_id,
        tensor_id=tensor_id,
        slice_hint=slice_hint,
        scale_hint=scale_hint,
        reduction_method=reduction_method,
    )

    # Build dask array from context
    chunks = [ep[0] for ep in ctx.endpoints]
    chunk_bounds_list = [ep[1] for ep in ctx.endpoints]
    dask_arr = self._build_dask_array(
        desc=ctx.descriptor,
        chunks=chunks,
        chunk_bounds=chunk_bounds_list,
        schema_metadata=ctx.schema_metadata,
    )

    # Crop to the originally requested region.
    # The server snaps slice_hint outward to lcm-aligned chunk boundaries, so
    # the returned descriptor.shape may be larger than what was requested.
    # We crop the dask array back to the exact requested region here.
    if ctx.original_slice_hint is not None and ctx.descriptor.HasField(
        "slice_hint"
    ):
        realized = ctx.descriptor.slice_hint
        ndim = len(ctx.descriptor.shape)
        scale = list(ctx.read_opt.scale_hint) if ctx.read_opt.scale_hint else None
        crop = []
        for ax in range(ndim):
            req_start = int(ctx.original_slice_hint.start[ax])
            req_stop = int(ctx.original_slice_hint.stop[ax])
            ret_start = int(realized.start[ax])
            s = int(scale[ax]) if scale and ax < len(scale) else 1
            if s > 1:
                logical_start = (req_start - ret_start) // s
                logical_stop = (req_stop - ret_start + s - 1) // s
            else:
                logical_start = req_start - ret_start
                logical_stop = req_stop - ret_start
            crop.append(slice(logical_start, logical_stop))
        dask_arr = dask_arr[tuple(crop)]

    return dask_arr

get_tensor_pb

get_tensor_pb(
    array_id: Optional[str] = None,
    tensor_id: Optional[str] = None,
    slice_hint: Optional[Tuple[slice, ...]] = None,
    scale_hint: Optional[Sequence[int]] = None,
    reduction_method: Optional[str] = None,
    *,
    source_id: Optional[str] = None
) -> SerializedTensor

Get a SerializedTensor protobuf for cross-process transfer.

Returns a protobuf containing connection info and chunk tickets for lazy reconstruction. The protobuf can be serialized to bytes and broadcast to worker processes, where each worker can call tensor_from_pb() to reconstruct a lazy dask array.

Parameters:

Name Type Description Default
array_id Optional[str]

Globally-unique tensor id (identity policy) -- e.g. "zarr_a3f2" or "aics_7f3/Image:0".

None
tensor_id Optional[str]

DEPRECATED. The legacy (source_id, tensor_id) form; pass the array_id as the single first argument instead.

None
slice_hint Optional[Tuple[slice, ...]]

Optional slice tuple to filter chunks

None
scale_hint Optional[Sequence[int]]

Optional per-dimension integer downsampling factors

None
reduction_method Optional[str]

Optional dynamic reduction method for scaled reads

None

Returns:

Type Description
SerializedTensor

SerializedTensor protobuf object

Source code in src/main/python/biopb/tensor/client.py
def get_tensor_pb(
    self,
    array_id: Optional[str] = None,
    tensor_id: Optional[str] = None,
    slice_hint: Optional[Tuple[slice, ...]] = None,
    scale_hint: Optional[Sequence[int]] = None,
    reduction_method: Optional[str] = None,
    *,
    source_id: Optional[str] = None,
) -> SerializedTensor:
    """Get a SerializedTensor protobuf for cross-process transfer.

    Returns a protobuf containing connection info and chunk tickets
    for lazy reconstruction. The protobuf can be serialized to bytes
    and broadcast to worker processes, where each worker can call
    tensor_from_pb() to reconstruct a lazy dask array.

    Args:
        array_id: Globally-unique tensor id (identity policy) -- e.g.
            ``"zarr_a3f2"`` or ``"aics_7f3/Image:0"``.
        tensor_id: DEPRECATED. The legacy ``(source_id, tensor_id)`` form;
            pass the array_id as the single first argument instead.
        slice_hint: Optional slice tuple to filter chunks
        scale_hint: Optional per-dimension integer downsampling factors
        reduction_method: Optional dynamic reduction method for scaled reads

    Returns:
        SerializedTensor protobuf object
    """
    source_id, tensor_id = self._resolve_array_id(
        array_id, tensor_id, "get_tensor_pb", source_id=source_id
    )
    logger.debug(f"get_tensor_pb: source_id={source_id}, tensor_id={tensor_id}")

    # Get flight info context
    ctx = self._get_tensor_context(
        source_id=source_id,
        tensor_id=tensor_id,
        slice_hint=slice_hint,
        scale_hint=scale_hint,
        reduction_method=reduction_method,
    )

    # Serialize endpoints
    serialized_endpoints = []
    for chunk_id, bounds in ctx.endpoints:
        ticket = TensorTicket(chunk_id=chunk_id)
        serialized_ep = SerializedEndpoint(
            ticket=ticket,
            chunk_bounds=bounds,
        )
        serialized_endpoints.append(serialized_ep)

    # Build SerializedTensor
    serialized_tensor = SerializedTensor(
        tensor_descriptor=ctx.descriptor,
        location=self._location,
        auth_token=self._token or "",
        endpoints=serialized_endpoints,
    )
    if ctx.original_slice_hint is not None:
        serialized_tensor.original_slice_hint.CopyFrom(ctx.original_slice_hint)

    # Add schema metadata for SHM transfer feature detection
    if ctx.schema_metadata is not None:
        serialized_tensor.schema_metadata.update(ctx.schema_metadata)

    return serialized_tensor

tensor_from_pb staticmethod

tensor_from_pb(
    pb: SerializedTensor, cache_bytes: int = 1000000000
) -> da.Array

Reconstruct a lazy dask array from SerializedTensor protobuf.

Creates a dask array that fetches chunks from the Flight server independently. Each worker process maintains its own connection pool and LRU cache keyed by (location, auth_token).

If endpoints field is empty, calls GetFlightInfo on the server to rebuild the endpoint list.

If debug_pickled_array is populated, unpickles directly (bypasses server).

Parameters:

Name Type Description Default
pb SerializedTensor

SerializedTensor protobuf object

required
cache_bytes int

Maximum bytes for chunk cache (default 1GB). Only effective for the first tensor created in a process for a given (location, auth_token) pair.

1000000000

Returns:

Type Description
Array

dask.array with lazy chunk loading

Source code in src/main/python/biopb/tensor/client.py
@staticmethod
def tensor_from_pb(
    pb: SerializedTensor,
    cache_bytes: int = 1_000_000_000,
) -> da.Array:
    """Reconstruct a lazy dask array from SerializedTensor protobuf.

    Creates a dask array that fetches chunks from the Flight server
    independently. Each worker process maintains its own connection
    pool and LRU cache keyed by (location, auth_token).

    If endpoints field is empty, calls GetFlightInfo on the server
    to rebuild the endpoint list.

    If debug_pickled_array is populated, unpickles directly (bypasses server).

    Args:
        pb: SerializedTensor protobuf object
        cache_bytes: Maximum bytes for chunk cache (default 1GB).
            Only effective for the first tensor created in a process
            for a given (location, auth_token) pair.

    Returns:
        dask.array with lazy chunk loading
    """
    import pickle

    # Debug path: unpickle directly if debug_pickled_array is present
    if pb.debug_pickled_array:
        return pickle.loads(pb.debug_pickled_array)

    descriptor = pb.tensor_descriptor
    shape = tuple(descriptor.shape)
    dtype = np.dtype(descriptor.dtype)

    # Parse endpoints - if empty, fetch from GetFlightInfo
    chunks = []
    chunk_bounds_list = []

    if pb.endpoints:
        # Use serialized endpoints directly
        for ep in pb.endpoints:
            chunks.append(ep.ticket.chunk_id)
            chunk_bounds_list.append(ep.chunk_bounds)
    else:
        # Endpoints not provided - call GetFlightInfo to rebuild
        logger.debug("tensor_from_pb: endpoints empty, calling GetFlightInfo")
        chunks, chunk_bounds_list = _fetch_endpoints_via_get_flight_info(pb)

    # Build chunk map
    chunk_map = {}
    axis_starts = [
        sorted({int(bounds.start[axis]) for bounds in chunk_bounds_list})
        for axis in range(len(shape))
    ]
    axis_index_maps = [
        {start: index for index, start in enumerate(starts)}
        for starts in axis_starts
    ]
    for chunk_id, bounds in zip(chunks, chunk_bounds_list):
        chunk_idx = tuple(
            axis_index_maps[d][int(bounds.start[d])] for d in range(len(shape))
        )
        chunk_map[chunk_idx] = (chunk_id, bounds)

    # Build dask array with lazy chunk fetching
    ndim = len(shape)
    grid_shape = tuple(max(idx[d] + 1 for idx in chunk_map) for d in range(ndim))

    # Extract schema_metadata from pb for SHM transfer
    schema_metadata = dict(pb.schema_metadata) if pb.schema_metadata else None

    dask_arr = _build_dask_array_from_chunk_map(
        chunk_map,
        grid_shape,
        shape,
        dtype,
        pb.location,
        pb.auth_token if pb.auth_token else None,
        cache_bytes,
        schema_metadata,
    )

    # Crop to the originally requested region if original_slice_hint present
    if pb.HasField("original_slice_hint") and descriptor.HasField("slice_hint"):
        realized = descriptor.slice_hint
        original = pb.original_slice_hint
        ndim = len(descriptor.shape)
        scale = list(descriptor.scale_hint) if descriptor.scale_hint else None
        crop = []
        for ax in range(ndim):
            req_start = int(original.start[ax])
            req_stop = int(original.stop[ax])
            ret_start = int(realized.start[ax])
            s = int(scale[ax]) if scale and ax < len(scale) else 1
            if s > 1:
                logical_start = (req_start - ret_start) // s
                logical_stop = (req_stop - ret_start + s - 1) // s
            else:
                logical_start = req_start - ret_start
                logical_stop = req_stop - ret_start
            crop.append(slice(logical_start, logical_stop))
        dask_arr = dask_arr[tuple(crop)]

    return dask_arr

upload_array

upload_array(
    arr: Array,
    source_name: str,
    chunk_shape: Optional[Sequence[int]] = None,
    dim_labels: Optional[Sequence[str]] = None,
    ome_metadata: Optional[dict] = None,
) -> str

Upload dask array to server.

Parameters:

Name Type Description Default
arr Array

Dask array to upload

required
source_name str

Source identifier format: - "cache:my-name" → cache-backed (ephemeral) - "cache:" → cache-backed with server-generated name - "ome_zarr:my-name" → zarr-backed (persistent) - "ome_zarr:" → zarr-backed with server-generated name

required
chunk_shape Optional[Sequence[int]]

Override chunk shape. If None, uses arr.chunksize with automatic rechunking if chunks are non-uniform.

None
dim_labels Optional[Sequence[str]]

Optional dimension labels

None
ome_metadata Optional[dict]

Optional OME metadata dict

None

Returns:

Type Description
str

source_id of created source (e.g., "cache_abc123" or "ome_zarr_def456")

Source code in src/main/python/biopb/tensor/client.py
def upload_array(
    self,
    arr: da.Array,
    source_name: str,
    chunk_shape: Optional[Sequence[int]] = None,
    dim_labels: Optional[Sequence[str]] = None,
    ome_metadata: Optional[dict] = None,
) -> str:
    """Upload dask array to server.

    Args:
        arr: Dask array to upload
        source_name: Source identifier format:
            - "cache:my-name" → cache-backed (ephemeral)
            - "cache:" → cache-backed with server-generated name
            - "ome_zarr:my-name" → zarr-backed (persistent)
            - "ome_zarr:" → zarr-backed with server-generated name
        chunk_shape: Override chunk shape. If None, uses arr.chunksize with
                     automatic rechunking if chunks are non-uniform.
        dim_labels: Optional dimension labels
        ome_metadata: Optional OME metadata dict

    Returns:
        source_id of created source (e.g., "cache_abc123" or "ome_zarr_def456")
    """
    # Determine target chunk shape
    if chunk_shape is None:
        chunk_shape = arr.chunksize

        # Check if dask chunks are non-uniform
        needs_rechunk = not all(
            len(set(arr.chunks[d])) == 1 for d in range(arr.ndim)
        )

        if needs_rechunk:
            uniform_chunks = tuple(
                max(arr.chunks[d]) if arr.chunks[d] else arr.shape[d]
                for d in range(arr.ndim)
            )
            arr = arr.rechunk(uniform_chunks)
            chunk_shape = uniform_chunks
    else:
        if tuple(chunk_shape) != tuple(arr.chunksize):
            arr = arr.rechunk(tuple(chunk_shape))

    # Create source
    source_id = self.create_source(
        source_name=source_name,
        shape=arr.shape,
        dtype=arr.dtype.str,
        chunk_shape=chunk_shape,
        dim_labels=dim_labels,
        ome_metadata=ome_metadata,
    )

    # Upload chunks
    ndim = arr.ndim
    chunk_shape_tuple = tuple(chunk_shape)
    chunks_per_dim = [
        (arr.shape[d] + chunk_shape_tuple[d] - 1) // chunk_shape_tuple[d]
        for d in range(ndim)
    ]

    from itertools import product

    for chunk_idx in product(*(range(n) for n in chunks_per_dim)):
        chunk_start = [
            idx * chunk_shape_tuple[d] for d, idx in enumerate(chunk_idx)
        ]
        chunk_stop = [
            min((idx + 1) * chunk_shape_tuple[d], arr.shape[d])
            for d, idx in enumerate(chunk_idx)
        ]

        bounds = ChunkBounds(start=chunk_start, stop=chunk_stop)

        slices = tuple(
            slice(chunk_start[d], chunk_stop[d]) for d in range(arr.ndim)
        )
        chunk_data = arr[slices].compute()
        self.upload_chunk(source_id, bounds, chunk_data)

    return source_id

upload_zarr

upload_zarr(
    zarr_path: str,
    source_name: str,
    chunk_shape: Optional[Sequence[int]] = None,
    dim_labels: Optional[Sequence[str]] = None,
    ome_metadata: Optional[dict] = None,
) -> str

Upload local zarr to server.

Parameters:

Name Type Description Default
zarr_path str

Path to local zarr directory

required
source_name str

Source identifier format: - "cache:my-name" → cache-backed (ephemeral) - "cache:" → cache-backed with server-generated name - "ome_zarr:my-name" → zarr-backed (persistent) - "ome_zarr:" → zarr-backed with server-generated name

required
chunk_shape Optional[Sequence[int]]

Override chunk shape. If None, uses zarr's chunk shape.

None
dim_labels Optional[Sequence[str]]

Optional dimension labels (read from zarr if not provided)

None
ome_metadata Optional[dict]

Optional OME metadata (read from zarr if not provided)

None

Returns:

Type Description
str

source_id of created source (e.g., "cache_abc123" or "ome_zarr_def456")

Source code in src/main/python/biopb/tensor/client.py
def upload_zarr(
    self,
    zarr_path: str,
    source_name: str,
    chunk_shape: Optional[Sequence[int]] = None,
    dim_labels: Optional[Sequence[str]] = None,
    ome_metadata: Optional[dict] = None,
) -> str:
    """Upload local zarr to server.

    Args:
        zarr_path: Path to local zarr directory
        source_name: Source identifier format:
            - "cache:my-name" → cache-backed (ephemeral)
            - "cache:" → cache-backed with server-generated name
            - "ome_zarr:my-name" → zarr-backed (persistent)
            - "ome_zarr:" → zarr-backed with server-generated name
        chunk_shape: Override chunk shape. If None, uses zarr's chunk shape.
        dim_labels: Optional dimension labels (read from zarr if not provided)
        ome_metadata: Optional OME metadata (read from zarr if not provided)

    Returns:
        source_id of created source (e.g., "cache_abc123" or "ome_zarr_def456")
    """
    import zarr

    arr = zarr.open_array(zarr_path, mode="r")

    # Read metadata from local zarr if not provided
    zattrs_path = Path(zarr_path) / ".zattrs"
    if zattrs_path.exists():
        with open(zattrs_path) as f:
            zattrs = json.load(f)
        if ome_metadata is None and "multiscales" in zattrs:
            ome_metadata = zattrs
        if dim_labels is None and "multiscales" in zattrs:
            axes = zattrs["multiscales"][0].get("axes", [])
            dim_labels = [
                ax.get("name") if isinstance(ax, dict) else str(ax) for ax in axes
            ]

    dask_arr = da.from_zarr(zarr_path)
    effective_chunk_shape = chunk_shape or arr.chunks

    return self.upload_array(
        dask_arr,
        source_name=source_name,
        chunk_shape=effective_chunk_shape,
        dim_labels=dim_labels,
        ome_metadata=ome_metadata,
    )

create_source

create_source(
    source_name: str,
    shape: Sequence[int],
    dtype: str,
    chunk_shape: Sequence[int],
    dim_labels: Optional[Sequence[str]] = None,
    ome_metadata: Optional[dict] = None,
) -> str

Create source on server (internal).

Parameters:

Name Type Description Default
source_name str

"cache:name" → cache-backed; "ome_zarr:name" → zarr-backed "cache:" or "ome_zarr:" → server-generated name

required
shape Sequence[int]

Array shape

required
dtype str

Data type string (numpy format)

required
chunk_shape Sequence[int]

Chunk size per dimension

required
dim_labels Optional[Sequence[str]]

Optional dimension labels

None
ome_metadata Optional[dict]

Optional OME metadata dict

None

Returns:

Type Description
str

source_id assigned by server

Source code in src/main/python/biopb/tensor/client.py
def create_source(
    self,
    source_name: str,
    shape: Sequence[int],
    dtype: str,
    chunk_shape: Sequence[int],
    dim_labels: Optional[Sequence[str]] = None,
    ome_metadata: Optional[dict] = None,
) -> str:
    """Create source on server (internal).

    Args:
        source_name: "cache:name" → cache-backed; "ome_zarr:name" → zarr-backed
                     "cache:" or "ome_zarr:" → server-generated name
        shape: Array shape
        dtype: Data type string (numpy format)
        chunk_shape: Chunk size per dimension
        dim_labels: Optional dimension labels
        ome_metadata: Optional OME metadata dict

    Returns:
        source_id assigned by server
    """
    req_desc = TensorDescriptor(
        array_id=source_name,
        shape=list(shape),
        dtype=dtype,
        chunk_shape=list(chunk_shape),
        dim_labels=list(dim_labels or []),
        metadata_json=json.dumps(ome_metadata) if ome_metadata else "",
    )

    action = flight.Action("create_source", req_desc.SerializeToString())
    results = self._client.do_action(action, options=self._call_options)
    try:
        result = next(results)
    except StopIteration as exc:
        raise RuntimeError("create_source: server returned no result") from exc

    response_desc = TensorDescriptor.FromString(result.body.to_pybytes())
    logger.info(f"create_source: created {response_desc.array_id}")
    return response_desc.array_id

upload_chunk

upload_chunk(
    source_id: str, bounds: ChunkBounds, data: ndarray
) -> None

Upload single chunk (internal).

Parameters:

Name Type Description Default
source_id str

Source identifier

required
bounds ChunkBounds

Chunk start/stop coordinates

required
data ndarray

Numpy array with chunk data

required
Source code in src/main/python/biopb/tensor/client.py
def upload_chunk(
    self,
    source_id: str,
    bounds: ChunkBounds,
    data: np.ndarray,
) -> None:
    """Upload single chunk (internal).

    Args:
        source_id: Source identifier
        bounds: Chunk start/stop coordinates
        data: Numpy array with chunk data
    """
    upload = ChunkUpload(
        source_id=source_id,
        bounds=bounds,
    )

    desc = flight.FlightDescriptor.for_command(upload.SerializeToString())
    schema = pa.schema([pa.field("data", pa.from_numpy_dtype(data.dtype))])

    writer, reader = self._client.do_put(desc, schema, options=self._call_options)
    batch = pa.RecordBatch.from_arrays([pa.array(data.ravel())], ["data"])
    writer.write_batch(batch)
    writer.done_writing()
    writer.close()
    reader.read()
    logger.debug(f"upload_chunk: uploaded {data.nbytes} bytes to {source_id}")

close

close()

Close the Flight client.

Source code in src/main/python/biopb/tensor/client.py
def close(self):
    """Close the Flight client."""
    logger.info("Closing Flight client")
    self._client.close()

health_check

health_check() -> Dict[str, Any]

Check server health status via Flight action.

Returns:

Type Description
Dict[str, Any]

Dictionary with health status information:

Dict[str, Any]
  • status: "SERVING" or other status string. Note: with progressive discovery, SERVING means "up and serving the possibly-still- populating catalog," not "catalog complete" -- use the freshness fields below to tell whether indexing is still in progress.
Dict[str, Any]
  • source_count: Number of registered sources
Dict[str, Any]
  • metadata_db_enabled: Whether metadata database is enabled
Dict[str, Any]
  • writable: Whether server accepts uploads
Dict[str, Any]
  • uptime_seconds: Server uptime in seconds
Dict[str, Any]
  • full_scan_in_progress: Whether a full catalog rescan is running now (absent on older servers)
Dict[str, Any]
  • last_full_scan_finished_at: Epoch seconds when a full scan last succeeded, or None until the first one does (absent on older servers)

Raises:

Type Description
FlightError

If server is unreachable or action fails

Source code in src/main/python/biopb/tensor/client.py
def health_check(self) -> Dict[str, Any]:
    """Check server health status via Flight action.

    Returns:
        Dictionary with health status information:
        - status: "SERVING" or other status string. Note: with progressive
          discovery, SERVING means "up and serving the possibly-still-
          populating catalog," not "catalog complete" -- use the freshness
          fields below to tell whether indexing is still in progress.
        - source_count: Number of registered sources
        - metadata_db_enabled: Whether metadata database is enabled
        - writable: Whether server accepts uploads
        - uptime_seconds: Server uptime in seconds
        - full_scan_in_progress: Whether a full catalog rescan is running now
          (absent on older servers)
        - last_full_scan_finished_at: Epoch seconds when a full scan last
          succeeded, or None until the first one does (absent on older
          servers)

    Raises:
        FlightError: If server is unreachable or action fails
    """
    action = flight.Action("health", b"")
    results = self._client.do_action(action, options=self._call_options)
    for result in results:
        return json.loads(result.body.to_pybytes())
    return {"status": "UNKNOWN"}

cache_stats

cache_stats() -> Dict[str, Any]

Fetch server-side cache statistics via Flight action.

Returns:

Type Description
Dict[str, Any]

Dictionary of CacheStats fields: total_entries, total_bytes,

Dict[str, Any]

max_entries, max_bytes, hits, misses, evictions, pending_waits,

Dict[str, Any]

ref_held_evictions_skipped, oversized_skips, and (file backend)

Dict[str, Any]

per-pool stats under "pool_stats".

Raises:

Type Description
FlightError

If server is unreachable or action fails

Source code in src/main/python/biopb/tensor/client.py
def cache_stats(self) -> Dict[str, Any]:
    """Fetch server-side cache statistics via Flight action.

    Returns:
        Dictionary of CacheStats fields: total_entries, total_bytes,
        max_entries, max_bytes, hits, misses, evictions, pending_waits,
        ref_held_evictions_skipped, oversized_skips, and (file backend)
        per-pool stats under "pool_stats".

    Raises:
        FlightError: If server is unreachable or action fails
    """
    action = flight.Action("cache_stats", b"")
    results = self._client.do_action(action, options=self._call_options)
    for result in results:
        return json.loads(result.body.to_pybytes())
    return {}

get_upload_status

get_upload_status(source_id: str) -> Dict[str, Any]

Get upload status for a writable source.

Parameters:

Name Type Description Default
source_id str

Source identifier returned by create_source()

required

Returns:

Type Description
Dict[str, Any]

Dictionary with source_id, state, expected_chunks, and uploaded_chunks.

Source code in src/main/python/biopb/tensor/client.py
def get_upload_status(self, source_id: str) -> Dict[str, Any]:
    """Get upload status for a writable source.

    Args:
        source_id: Source identifier returned by create_source()

    Returns:
        Dictionary with source_id, state, expected_chunks, and uploaded_chunks.
    """
    action = flight.Action("upload_status", source_id.encode("utf-8"))
    results = self._client.do_action(action, options=self._call_options)
    for result in results:
        return json.loads(result.body.to_pybytes())
    return {
        "source_id": source_id,
        "state": "UNKNOWN",
        "expected_chunks": 0,
        "uploaded_chunks": 0,
    }

get_upload_status_pb

get_upload_status_pb(
    pb: SerializedTensor,
) -> Dict[str, Any]

Get upload status for a registration-first SerializedTensor handle.

This helper is intended for cache-backed handles returned before upload completion, where tensor_descriptor.array_id is the source identifier.

Parameters:

Name Type Description Default
pb SerializedTensor

SerializedTensor handle returned by a registration-first flow.

required

Returns:

Type Description
Dict[str, Any]

Dictionary with source_id, state, expected_chunks, and uploaded_chunks.

Source code in src/main/python/biopb/tensor/client.py
def get_upload_status_pb(self, pb: SerializedTensor) -> Dict[str, Any]:
    """Get upload status for a registration-first SerializedTensor handle.

    This helper is intended for cache-backed handles returned before upload
    completion, where tensor_descriptor.array_id is the source identifier.

    Args:
        pb: SerializedTensor handle returned by a registration-first flow.

    Returns:
        Dictionary with source_id, state, expected_chunks, and uploaded_chunks.
    """
    return self.get_upload_status(_upload_source_id_from_pb(pb))

wait_for_upload_ready

wait_for_upload_ready(
    source_id: str,
    timeout_seconds: float = 60.0,
    poll_interval_seconds: float = 0.5,
) -> Dict[str, Any]

Poll upload status until the source reports READY.

Parameters:

Name Type Description Default
source_id str

Source identifier returned by create_source().

required
timeout_seconds float

Maximum time to wait before timing out.

60.0
poll_interval_seconds float

Delay between status checks.

0.5

Returns:

Type Description
Dict[str, Any]

Final upload status dictionary when READY.

Raises:

Type Description
TimeoutError

If the upload does not reach READY within the timeout.

RuntimeError

If the upload reports FAILED.

Source code in src/main/python/biopb/tensor/client.py
def wait_for_upload_ready(
    self,
    source_id: str,
    timeout_seconds: float = 60.0,
    poll_interval_seconds: float = 0.5,
) -> Dict[str, Any]:
    """Poll upload status until the source reports READY.

    Args:
        source_id: Source identifier returned by create_source().
        timeout_seconds: Maximum time to wait before timing out.
        poll_interval_seconds: Delay between status checks.

    Returns:
        Final upload status dictionary when READY.

    Raises:
        TimeoutError: If the upload does not reach READY within the timeout.
        RuntimeError: If the upload reports FAILED.
    """
    deadline = time.monotonic() + timeout_seconds
    while True:
        status = self.get_upload_status(source_id)
        state = status.get("state")
        if state == "READY":
            return status
        if state == "FAILED":
            raise RuntimeError(f"Upload failed for source '{source_id}'")
        if time.monotonic() >= deadline:
            raise TimeoutError(
                f"Timed out waiting for upload readiness for source '{source_id}'"
            )
        time.sleep(poll_interval_seconds)

wait_for_upload_ready_pb

wait_for_upload_ready_pb(
    pb: SerializedTensor,
    timeout_seconds: float = 60.0,
    poll_interval_seconds: float = 0.5,
) -> Dict[str, Any]

Poll upload status until a registration-first SerializedTensor is READY.

Source code in src/main/python/biopb/tensor/client.py
def wait_for_upload_ready_pb(
    self,
    pb: SerializedTensor,
    timeout_seconds: float = 60.0,
    poll_interval_seconds: float = 0.5,
) -> Dict[str, Any]:
    """Poll upload status until a registration-first SerializedTensor is READY."""
    return self.wait_for_upload_ready(
        _upload_source_id_from_pb(pb),
        timeout_seconds=timeout_seconds,
        poll_interval_seconds=poll_interval_seconds,
    )

cache_info

cache_info() -> Dict

Return cache statistics from the pooled cache for this connection.

Returns:

Type Description
Dict

Dictionary with cache size and item count

Source code in src/main/python/biopb/tensor/client.py
def cache_info(self) -> Dict:
    """Return cache statistics from the pooled cache for this connection.

    Returns:
        Dictionary with cache size and item count
    """
    key = (self._location, self._token)
    pool_entry = _CACHE_POOL.get(key)
    if pool_entry is None:
        # No cache allocated -- either nothing fetched yet, or caching is
        # disabled for this connection (e.g. localhost). Report the resolved
        # size so a disabled cache truthfully shows max_bytes == 0.
        resolved = _resolve_cache_bytes(self._location, self._cache_bytes)
        return {"size_bytes": 0, "max_bytes": resolved, "item_count": 0}
    cache = pool_entry[1]  # Extract cache from (pid, cache) tuple
    if cache is None:
        # Pinned off by configure_cache(): report max_bytes == 0 truthfully.
        return {"size_bytes": 0, "max_bytes": 0, "item_count": 0}
    return {
        "size_bytes": cache.total_bytes,
        "max_bytes": cache.available_bytes,
        "item_count": len(cache.data),
    }

cache_clear

cache_clear()

Clear the pooled cache for this connection namespace.

Source code in src/main/python/biopb/tensor/client.py
def cache_clear(self):
    """Clear the pooled cache for this connection namespace."""
    key = (self._location, self._token)
    pool_entry = _CACHE_POOL.get(key)
    if pool_entry is not None and pool_entry[1] is not None:
        pool_entry[1].clear()

configure_cache

configure_cache(
    location: str, token: Optional[str], cache_bytes: int
) -> int

Pin this process's chunk-cache budget for a connection, authoritatively.

Sets the per-process chunk cache to cache_bytes and keeps it there: every later fetch honors this budget regardless of the cache_bytes it requests. Idempotent. Call it once per worker process (e.g. from a dask worker-init plugin) to fix the budget deterministically across a dynamically-sized cluster.

A localhost server (the default) or cache_bytes <= 0 pins the cache OFF; later fetches then skip caching rather than recreating one of their own.

Parameters:

Name Type Description Default
location str

Flight server location string

required
token Optional[str]

Bearer token (or None for no auth)

required
cache_bytes int

Requested per-process cache size in bytes

required

Returns:

Type Description
int

The effective (resolved) cache size that was pinned, in bytes.

Source code in src/main/python/biopb/tensor/client.py
def configure_cache(location: str, token: Optional[str], cache_bytes: int) -> int:
    """Pin this process's chunk-cache budget for a connection, authoritatively.

    Sets the per-process chunk cache to ``cache_bytes`` and keeps it there: every
    later fetch honors this budget regardless of the ``cache_bytes`` it requests.
    Idempotent. Call it once per worker process (e.g. from a dask worker-init
    plugin) to fix the budget deterministically across a dynamically-sized
    cluster.

    A localhost server (the default) or ``cache_bytes <= 0`` pins the cache OFF;
    later fetches then skip caching rather than recreating one of their own.

    Args:
        location: Flight server location string
        token: Bearer token (or None for no auth)
        cache_bytes: Requested per-process cache size in bytes

    Returns:
        The effective (resolved) cache size that was pinned, in bytes.
    """
    # Unlike the lazy first-touch creation in _get_shared_cache, this (re)sizes
    # the pooled cache now and records a None sentinel for the disabled case, so
    # a later fetch can't undo the decision from its own cache_bytes.
    effective = _resolve_cache_bytes(location, cache_bytes)
    key = (location, token)
    current_pid = os.getpid()

    with _POOL_LOCK:
        existing = _CACHE_POOL.get(key)
        if effective <= 0:
            # Pin OFF authoritatively: store a None-cache sentinel instead of
            # deleting, so a later fetch's _get_shared_cache honors the decision
            # rather than recreating a cache from its own cache_bytes.
            _CACHE_POOL[key] = (current_pid, None)
            return 0
        # (Re)create when absent, inherited from a fork, or a different size.
        if (
            existing is None
            or existing[0] != current_pid
            or existing[1].available_bytes != effective
        ):
            _CACHE_POOL[key] = (current_pid, Cache(available_bytes=effective))

    return effective

make_debug_serialized_tensor

make_debug_serialized_tensor(
    arr: Array, array_id: str = "debug"
) -> SerializedTensor

Create a SerializedTensor with debug_pickled_array for testing.

Eagerly computes the array and pickles it, bypassing Flight server. Preserves original chunk structure for testing chunk-related behavior. Populates inferable tensor_descriptor fields.

Parameters:

Name Type Description Default
arr Array

Dask array to serialize

required
array_id str

Optional array identifier

'debug'

Returns:

Type Description
SerializedTensor

SerializedTensor with debug_pickled_array populated

Source code in src/main/python/biopb/tensor/client.py
def make_debug_serialized_tensor(
    arr: da.Array, array_id: str = "debug"
) -> SerializedTensor:
    """Create a SerializedTensor with debug_pickled_array for testing.

    Eagerly computes the array and pickles it, bypassing Flight server.
    Preserves original chunk structure for testing chunk-related behavior.
    Populates inferable tensor_descriptor fields.

    Args:
        arr: Dask array to serialize
        array_id: Optional array identifier

    Returns:
        SerializedTensor with debug_pickled_array populated
    """
    import pickle

    # Eager compute
    np_arr = arr.compute()

    # Rechunk to original chunk structure (preserves chunk boundaries for testing)
    computed_da = da.from_array(np_arr, chunks=arr.chunksize)

    descriptor = TensorDescriptor(
        array_id=array_id,
        shape=list(arr.shape),
        dtype=np.dtype(arr.dtype).str,
        chunk_shape=list(arr.chunksize),
    )

    return SerializedTensor(
        tensor_descriptor=descriptor,
        debug_pickled_array=pickle.dumps(computed_da),
    )