segrails.pipeline.weights

Weight resolution pipeline step.

Resolves a source URI (azure://, https://, local://, or a plain path) to a local Path, managing the TTL cache and retrying on network errors.

Classes

WeightStep

Resolves model weights from a source URI to a local path.

Module Contents

class segrails.pipeline.weights.WeightStep(blob_manager: altametris.azure_sdk.storage_manager.blob_manager.BlobManager | None, cache_manager: altametris.azure_sdk.storage_manager.cache_manager.CacheManager, max_retries: int = 3, retry_delay: float = 1.0)

Resolves model weights from a source URI to a local path.

Supports four source types:

  • azure://container/path — Azure Blob Storage SDK

  • https://apim.host/container/path — APIM HTTP

  • local:///absolute/path or plain path — local file or MLflow directory

File vs. directory is detected automatically for remote sources: the path is listed as a blob prefix; if results are found the whole directory is downloaded, otherwise the path is treated as a single blob.

The local cache is checked before any download. On network failure, successive attempts are made with exponential back-off. If all retries fail and an expired cached entry exists on disk, it is used as a last resort (with a warning).

Parameters:

  • blob_manager – BlobManager instance. Pass None when all sources are local.

  • cache_manager – CacheManager instance — manages the local TTL cache.

  • max_retries – Maximum number of download attempts on network errors.

  • retry_delay – Initial delay between retries in seconds (doubled each attempt).

Example

>>> step = WeightStep(blob_manager, cache_manager)
>>> path = step.resolve("azure://weights-dev/segrails/U-Net-segrails")
_blob_manager
_cache_manager
_max_retries = 3
_retry_delay = 1.0
resolve(source: str, force_download: bool = False) → pathlib.Path

Resolve a source URI to a local path (file or directory).

File vs. directory is auto-detected: - Local paths use Path.is_dir() / Path.is_file(). - Remote sources list blobs under the path as a prefix; results → directory,

no results → single file.

Parameters:

  • source – Source URI. Accepted formats: - azure://container/blob_or_prefix — Blob Storage SDK - https://apim.host/container/blob_or_prefix — APIM HTTP - local:///absolute/path — local file or MLflow directory - /absolute/path — local file or MLflow directory

  • force_download – If True, skip the cache and always download.

Returns:

Path to the local file or directory. Directories are validated to contain an MLmodel file (MLflow model layout).

Raises:

  • FileNotFoundError – If the local path does not exist.

  • ValueError – If the URI scheme is not recognised, or a local directory does not contain an MLmodel file.

  • RuntimeError – If the source is remote but no BlobManager was provided.

  • AzureDownloadError – If all download attempts fail and no fallback exists.

  • BlobNotFoundError – If the blob or prefix does not exist on the remote.

_resolve_local(source: str) → pathlib.Path
_resolve_remote(source: str, container: str, blob_path: str, force_download: bool) → pathlib.Path

Auto-detect file vs dir via blob listing, check cache, then download.

_download_file_with_retry(*, local_path: pathlib.Path, container: str, blob_path: str, source: str) → None
_download_dir_with_retry(*, local_dir: pathlib.Path, container: str, prefix: str, blobs: list[dict], source: str) → None