segrails.pipeline.weights
=========================

.. py:module:: segrails.pipeline.weights

.. autoapi-nested-parse::

   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
-------

.. autoapisummary::

   segrails.pipeline.weights.WeightStep


Module Contents
---------------

.. py:class:: 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).

   :param blob_manager: BlobManager instance. Pass ``None`` when all sources are local.
   :param cache_manager: CacheManager instance — manages the local TTL cache.
   :param max_retries: Maximum number of download attempts on network errors.
   :param retry_delay: Initial delay between retries in seconds (doubled each attempt).

   .. rubric:: Example

   >>> step = WeightStep(blob_manager, cache_manager)
   >>> path = step.resolve("azure://weights-dev/segrails/U-Net-segrails")


   .. py:attribute:: _blob_manager


   .. py:attribute:: _cache_manager


   .. py:attribute:: _max_retries
      :value: 3



   .. py:attribute:: _retry_delay
      :value: 1.0



   .. py:method:: 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.

      :param 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
      :param 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.
      :raises ValueError: If the URI scheme is not recognised, or a local directory
          does not contain an ``MLmodel`` file.
      :raises RuntimeError: If the source is remote but no ``BlobManager`` was provided.
      :raises AzureDownloadError: If all download attempts fail and no fallback exists.
      :raises BlobNotFoundError: If the blob or prefix does not exist on the remote.



   .. py:method:: _resolve_local(source: str) -> pathlib.Path


   .. py:method:: _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.



   .. py:method:: _download_file_with_retry(*, local_path: pathlib.Path, container: str, blob_path: str, source: str) -> None


   .. py:method:: _download_dir_with_retry(*, local_dir: pathlib.Path, container: str, prefix: str, blobs: list[dict], source: str) -> None