hirundo 0.1.21__py3-none-any.whl → 0.2.3.post1__py3-none-any.whl

hirundo/dataset_qa.py

@@ -1,11 +1,9 @@
 import datetime
-import json
 import typing
 from collections.abc import AsyncGenerator, Generator
 from enum import Enum
 from typing import overload
 
-import httpx
 from pydantic import BaseModel, Field, model_validator
 from tqdm import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
@@ -14,7 +12,16 @@ from hirundo._constraints import validate_labeling_info, validate_url
 from hirundo._env import API_HOST
 from hirundo._headers import get_headers
 from hirundo._http import raise_for_status_with_reason, requests
-from hirundo._iter_sse_retrying import aiter_sse_retrying, iter_sse_retrying
+from hirundo._run_checking import (
+    STATUS_TO_PROGRESS_MAP,
+    RunStatus,
+    aiter_run_events,
+    build_status_text_map,
+    get_state,
+    handle_run_failure,
+    iter_run_events,
+    update_progress_from_result,
+)
 from hirundo._timeouts import MODIFY_TIMEOUT, READ_TIMEOUT
 from hirundo._urls import HirundoUrl
 from hirundo.dataset_enum import DatasetMetadataType, LabelingType
@@ -35,87 +42,57 @@ class HirundoError(Exception):
     pass
 
 
-MAX_RETRIES = 200  # Max 200 retries for HTTP SSE connection
-
-
-class RunStatus(Enum):
-    PENDING = "PENDING"
-    STARTED = "STARTED"
-    SUCCESS = "SUCCESS"
-    FAILURE = "FAILURE"
-    AWAITING_MANUAL_APPROVAL = "AWAITING MANUAL APPROVAL"
-    REVOKED = "REVOKED"
-    REJECTED = "REJECTED"
-    RETRY = "RETRY"
-
-
-STATUS_TO_TEXT_MAP = {
-    RunStatus.STARTED.value: "Dataset QA run in progress. Downloading dataset",
-    RunStatus.PENDING.value: "Dataset QA run queued and not yet started",
-    RunStatus.SUCCESS.value: "Dataset QA run completed successfully",
-    RunStatus.FAILURE.value: "Dataset QA run failed",
-    RunStatus.AWAITING_MANUAL_APPROVAL.value: "Awaiting manual approval",
-    RunStatus.RETRY.value: "Dataset QA run failed. Retrying",
-    RunStatus.REVOKED.value: "Dataset QA run was cancelled",
-    RunStatus.REJECTED.value: "Dataset QA run was rejected",
-}
-STATUS_TO_PROGRESS_MAP = {
-    RunStatus.STARTED.value: 0.0,
-    RunStatus.PENDING.value: 0.0,
-    RunStatus.SUCCESS.value: 100.0,
-    RunStatus.FAILURE.value: 100.0,
-    RunStatus.AWAITING_MANUAL_APPROVAL.value: 100.0,
-    RunStatus.RETRY.value: 0.0,
-    RunStatus.REVOKED.value: 100.0,
-    RunStatus.REJECTED.value: 0.0,
-}
+STATUS_TO_TEXT_MAP = build_status_text_map(
+    "Dataset QA",
+    started_detail="Dataset QA run in progress. Downloading dataset",
+)
 
 
 class ClassificationRunArgs(BaseModel):
-    image_size: typing.Optional[tuple[int, int]] = (224, 224)
+    image_size: tuple[int, int] | None = (224, 224)
     """
     Size (width, height) to which to resize classification images.
     It is recommended to keep this value at (224, 224) unless your classes are differentiated by very small differences.
     """
-    upsample: typing.Optional[bool] = False
+    upsample: bool | None = False
     """
     Whether to upsample the dataset to attempt to balance the classes.
     """
 
 
 class ObjectDetectionRunArgs(ClassificationRunArgs):
-    min_abs_bbox_size: typing.Optional[int] = None
+    min_abs_bbox_size: int | None = None
     """
     Minimum valid size (in pixels) of a bounding box to keep it in the dataset for QA.
     """
-    min_abs_bbox_area: typing.Optional[int] = None
+    min_abs_bbox_area: int | None = None
     """
     Minimum valid absolute area (in pixels²) of a bounding box to keep it in the dataset for QA.
     """
-    min_rel_bbox_size: typing.Optional[float] = None
+    min_rel_bbox_size: float | None = None
     """
     Minimum valid size (as a fraction of both image height and width) for a bounding box
     to keep it in the dataset for QA, relative to the corresponding dimension size,
     i.e. if the bounding box is 10% of the image width and 5% of the image height, it will be kept if this value is 0.05, but not if the
     value is 0.06 (since both width and height are checked).
     """
-    min_rel_bbox_area: typing.Optional[float] = None
+    min_rel_bbox_area: float | None = None
     """
     Minimum valid relative area (as a fraction of the image area) of a bounding box to keep it in the dataset for QA.
     """
-    crop_ratio: typing.Optional[float] = None
+    crop_ratio: float | None = None
     """
     Ratio of the bounding box to crop.
     Change this value at your own risk. It is recommended to keep it at 1.0 unless you know what you are doing.
     """
-    add_mask_channel: typing.Optional[bool] = None
+    add_mask_channel: bool | None = None
     """
     Whether to add a mask channel to the image.
     Change at your own risk. It is recommended to keep it at False unless you know what you are doing.
     """
 
 
-RunArgs = typing.Union[ClassificationRunArgs, ObjectDetectionRunArgs]
+RunArgs = ClassificationRunArgs | ObjectDetectionRunArgs
 
 
 class AugmentationName(str, Enum):
@@ -128,32 +105,32 @@ class AugmentationName(str, Enum):
     GAUSSIAN_BLUR = "GaussianBlur"
 
 
-class Domain(str, Enum):
+class ModalityType(str, Enum):
     RADAR = "RADAR"
     VISION = "VISION"
     SPEECH = "SPEECH"
     TABULAR = "TABULAR"
 
 
-DOMAIN_TO_SUPPORTED_LABELING_TYPES = {
-    Domain.RADAR: [
+MODALITY_TO_SUPPORTED_LABELING_TYPES = {
+    ModalityType.RADAR: [
         LabelingType.SINGLE_LABEL_CLASSIFICATION,
         LabelingType.OBJECT_DETECTION,
     ],
-    Domain.VISION: [
+    ModalityType.VISION: [
         LabelingType.SINGLE_LABEL_CLASSIFICATION,
         LabelingType.OBJECT_DETECTION,
         LabelingType.OBJECT_SEGMENTATION,
         LabelingType.SEMANTIC_SEGMENTATION,
         LabelingType.PANOPTIC_SEGMENTATION,
     ],
-    Domain.SPEECH: [LabelingType.SPEECH_TO_TEXT],
-    Domain.TABULAR: [LabelingType.SINGLE_LABEL_CLASSIFICATION],
+    ModalityType.SPEECH: [LabelingType.SPEECH_TO_TEXT],
+    ModalityType.TABULAR: [LabelingType.SINGLE_LABEL_CLASSIFICATION],
 }
 
 
 class QADataset(BaseModel):
-    id: typing.Optional[int] = Field(default=None)
+    id: int | None = Field(default=None)
     """
     The ID of the dataset created on the server.
     """
@@ -169,17 +146,15 @@ class QADataset(BaseModel):
     - `LabelingType.OBJECT_DETECTION`: Indicates that the dataset is for object detection tasks
     - `LabelingType.SPEECH_TO_TEXT`: Indicates that the dataset is for speech-to-text tasks
     """
-    language: typing.Optional[str] = None
+    language: str | None = None
     """
     Language of the Speech-to-Text audio dataset. This is required for Speech-to-Text datasets.
     """
-    storage_config_id: typing.Optional[int] = None
+    storage_config_id: int | None = None
     """
     The ID of the storage config used to store the dataset and metadata.
     """
-    storage_config: typing.Optional[
-        typing.Union[StorageConfig, ResponseStorageConfig]
-    ] = None
+    storage_config: StorageConfig | ResponseStorageConfig | None = None
     """
     The `StorageConfig` instance to link to.
     """
@@ -193,41 +168,44 @@ class QADataset(BaseModel):
     Note: All CSV `image_path` entries in the metadata file should be relative to this folder.
     """
 
-    classes: typing.Optional[list[str]] = None
+    classes: list[str] | None = None
     """
     A full list of possible classes used in classification / object detection.
     It is currently required for clarity and performance.
     """
-    labeling_info: typing.Union[LabelingInfo, list[LabelingInfo]]
+    labeling_info: LabelingInfo | list[LabelingInfo]
 
-    augmentations: typing.Optional[list[AugmentationName]] = None
+    augmentations: list[AugmentationName] | None = None
     """
     Used to define which augmentations are apply to a vision dataset.
     For audio datasets, this field is ignored.
     If no value is provided, all augmentations are applied to vision datasets.
     """
-    domain: Domain = Domain.VISION
+    modality: ModalityType = ModalityType.VISION
     """
-    Used to define the domain of the dataset.
+    Used to define the modality of the dataset.
     Defaults to Image.
     """
 
-    run_id: typing.Optional[str] = Field(default=None, init=False)
+    run_id: str | None = Field(default=None, init=False)
     """
     The ID of the Dataset QA run created on the server.
     """
 
-    status: typing.Optional[RunStatus] = None
+    status: RunStatus | None = None
 
     @model_validator(mode="after")
     def validate_dataset(self):
-        if self.domain not in DOMAIN_TO_SUPPORTED_LABELING_TYPES:
+        if self.modality not in MODALITY_TO_SUPPORTED_LABELING_TYPES:
             raise ValueError(
-                f"Domain {self.domain} is not supported. Supported domains are: {list(DOMAIN_TO_SUPPORTED_LABELING_TYPES.keys())}"
+                f"Modality {self.modality} is not supported. Supported modalities are: {list(MODALITY_TO_SUPPORTED_LABELING_TYPES.keys())}"
             )
-        if self.labeling_type not in DOMAIN_TO_SUPPORTED_LABELING_TYPES[self.domain]:
+        if (
+            self.labeling_type
+            not in MODALITY_TO_SUPPORTED_LABELING_TYPES[self.modality]
+        ):
             raise ValueError(
-                f"Labeling type {self.labeling_type} is not supported for domain {self.domain}. Supported labeling types are: {DOMAIN_TO_SUPPORTED_LABELING_TYPES[self.domain]}"
+                f"Labeling type {self.labeling_type} is not supported for modality {self.modality}. Supported labeling types are: {MODALITY_TO_SUPPORTED_LABELING_TYPES[self.modality]}"
             )
         if self.storage_config is None and self.storage_config_id is None:
             raise ValueError(
@@ -307,7 +285,7 @@ class QADataset(BaseModel):
 
     @staticmethod
     def list_datasets(
-        organization_id: typing.Optional[int] = None,
+        organization_id: int | None = None,
     ) -> list["QADatasetOut"]:
         """
         Lists all the datasets created by user's default organization
@@ -333,7 +311,8 @@ class QADataset(BaseModel):
 
     @staticmethod
     def list_runs(
-        organization_id: typing.Optional[int] = None,
+        organization_id: int | None = None,
+        archived: bool | None = False,
     ) -> list["DataQARunOut"]:
         """
         Lists all the `QADataset` instances created by user's default organization
@@ -342,10 +321,11 @@ class QADataset(BaseModel):
 
         Args:
             organization_id: The ID of the organization to list the datasets for.
+            archived: Whether to list archived runs.
         """
         response = requests.get(
             f"{API_HOST}/dataset-qa/run/list",
-            params={"dataset_organization_id": organization_id},
+            params={"dataset_organization_id": organization_id, "archived": archived},
             headers=get_headers(),
             timeout=READ_TIMEOUT,
         )
@@ -396,7 +376,7 @@ class QADataset(BaseModel):
 
     def create(
         self,
-        organization_id: typing.Optional[int] = None,
+        organization_id: int | None = None,
         replace_if_exists: bool = False,
     ) -> int:
         """
@@ -453,8 +433,8 @@ class QADataset(BaseModel):
     @staticmethod
     def launch_qa_run(
         dataset_id: int,
-        organization_id: typing.Optional[int] = None,
-        run_args: typing.Optional[RunArgs] = None,
+        organization_id: int | None = None,
+        run_args: RunArgs | None = None,
     ) -> str:
         """
         Run the dataset QA process on the server using the dataset with the given ID
@@ -503,9 +483,9 @@ class QADataset(BaseModel):
 
     def run_qa(
         self,
-        organization_id: typing.Optional[int] = None,
+        organization_id: int | None = None,
         replace_dataset_if_exists: bool = False,
-        run_args: typing.Optional[RunArgs] = None,
+        run_args: RunArgs | None = None,
     ) -> str:
         """
         If the dataset was not created on the server yet, it is created.
@@ -556,53 +536,20 @@ class QADataset(BaseModel):
 
     @staticmethod
     def _check_run_by_id(run_id: str, retry=0) -> Generator[dict, None, None]:
-        if retry > MAX_RETRIES:
-            raise HirundoError("Max retries reached")
-        last_event = None
-        with httpx.Client(timeout=httpx.Timeout(None, connect=5.0)) as client:
-            for sse in iter_sse_retrying(
-                client,
-                "GET",
-                f"{API_HOST}/dataset-qa/run/{run_id}",
-                headers=get_headers(),
-            ):
-                if sse.event == "ping":
-                    continue
-                logger.debug(
-                    "[SYNC] received event: %s with data: %s and ID: %s and retry: %s",
-                    sse.event,
-                    sse.data,
-                    sse.id,
-                    sse.retry,
-                )
-                last_event = json.loads(sse.data)
-                if not last_event:
-                    continue
-                if "data" in last_event:
-                    data = last_event["data"]
-                else:
-                    if "detail" in last_event:
-                        raise HirundoError(last_event["detail"])
-                    elif "reason" in last_event:
-                        raise HirundoError(last_event["reason"])
-                    else:
-                        raise HirundoError("Unknown error")
-                yield data
-        if not last_event or last_event["data"]["state"] == RunStatus.PENDING.value:
-            QADataset._check_run_by_id(run_id, retry + 1)
-
-    @staticmethod
-    def _handle_failure(iteration: dict):
-        if iteration["result"]:
-            raise HirundoError(f"QA run failed with error: {iteration['result']}")
-        else:
-            raise HirundoError("QA run failed with an unknown error in _handle_failure")
+        yield from iter_run_events(
+            f"{API_HOST}/dataset-qa/run/{run_id}",
+            headers=get_headers(),
+            retry=retry,
+            status_keys=("state",),
+            error_cls=HirundoError,
+            log=logger,
+        )
 
     @staticmethod
     @overload
     def check_run_by_id(
         run_id: str, stop_on_manual_approval: typing.Literal[True]
-    ) -> typing.Optional[DatasetQAResults]: ...
+    ) -> DatasetQAResults | None: ...
 
     @staticmethod
     @overload
@@ -614,12 +561,12 @@ class QADataset(BaseModel):
     @overload
     def check_run_by_id(
         run_id: str, stop_on_manual_approval: bool
-    ) -> typing.Optional[DatasetQAResults]: ...
+    ) -> DatasetQAResults | None: ...
 
     @staticmethod
     def check_run_by_id(
         run_id: str, stop_on_manual_approval: bool = False
-    ) -> typing.Optional[DatasetQAResults]:
+    ) -> DatasetQAResults | None:
         """
         Check the status of a run given its ID
 
@@ -637,22 +584,25 @@ class QADataset(BaseModel):
         with logging_redirect_tqdm():
             t = tqdm(total=100.0)
             for iteration in QADataset._check_run_by_id(run_id):
-                if iteration["state"] in STATUS_TO_PROGRESS_MAP:
-                    t.set_description(STATUS_TO_TEXT_MAP[iteration["state"]])
-                    t.n = STATUS_TO_PROGRESS_MAP[iteration["state"]]
+                state = get_state(iteration, ("state",))
+                if state in STATUS_TO_PROGRESS_MAP:
+                    t.set_description(STATUS_TO_TEXT_MAP[state])
+                    t.n = STATUS_TO_PROGRESS_MAP[state]
                     logger.debug("Setting progress to %s", t.n)
                     t.refresh()
-                    if iteration["state"] in [
+                    if state in [
                         RunStatus.FAILURE.value,
                         RunStatus.REJECTED.value,
                         RunStatus.REVOKED.value,
                     ]:
                         logger.error(
                             "State is failure, rejected, or revoked: %s",
-                            iteration["state"],
+                            state,
                         )
-                        QADataset._handle_failure(iteration)
-                    elif iteration["state"] == RunStatus.SUCCESS.value:
+                        handle_run_failure(
+                            iteration, error_cls=HirundoError, run_label="QA"
+                        )
+                    elif state == RunStatus.SUCCESS.value:
                         t.close()
                         zip_temporary_url = iteration["result"]
                         logger.debug("QA run completed. Downloading results")
@@ -662,45 +612,24 @@ class QADataset(BaseModel):
                             zip_temporary_url,
                         )
                     elif (
-                        iteration["state"] == RunStatus.AWAITING_MANUAL_APPROVAL.value
+                        state == RunStatus.AWAITING_MANUAL_APPROVAL.value
                         and stop_on_manual_approval
                     ):
                         t.close()
                         return None
-                elif iteration["state"] is None:
-                    if (
-                        iteration["result"]
-                        and isinstance(iteration["result"], dict)
-                        and iteration["result"]["result"]
-                        and isinstance(iteration["result"]["result"], str)
-                    ):
-                        result_info = iteration["result"]["result"].split(":")
-                        if len(result_info) > 1:
-                            stage = result_info[0]
-                            current_progress_percentage = float(
-                                result_info[1].removeprefix(" ").removesuffix("% done")
-                            )
-                        elif len(result_info) == 1:
-                            stage = result_info[0]
-                            current_progress_percentage = t.n  # Keep the same progress
-                        else:
-                            stage = "Unknown progress state"
-                            current_progress_percentage = t.n  # Keep the same progress
-                        desc = (
-                            "QA run completed. Uploading results"
-                            if current_progress_percentage == 100.0
-                            else stage
-                        )
-                        t.set_description(desc)
-                        t.n = current_progress_percentage
-                        logger.debug("Setting progress to %s", t.n)
-                        t.refresh()
+                elif state is None:
+                    update_progress_from_result(
+                        iteration,
+                        t,
+                        uploading_text="QA run completed. Uploading results",
+                        log=logger,
+                    )
         raise HirundoError("QA run failed with an unknown error in check_run_by_id")
 
     @overload
     def check_run(
         self, stop_on_manual_approval: typing.Literal[True]
-    ) -> typing.Optional[DatasetQAResults]: ...
+    ) -> DatasetQAResults | None: ...
 
     @overload
     def check_run(
@@ -709,7 +638,7 @@ class QADataset(BaseModel):
 
     def check_run(
         self, stop_on_manual_approval: bool = False
-    ) -> typing.Optional[DatasetQAResults]:
+    ) -> DatasetQAResults | None:
         """
         Check the status of the current active instance's run.
 
@@ -741,32 +670,15 @@ class QADataset(BaseModel):
 
         """
         logger.debug("Checking run with ID: %s", run_id)
-        if retry > MAX_RETRIES:
-            raise HirundoError("Max retries reached")
-        last_event = None
-        async with httpx.AsyncClient(
-            timeout=httpx.Timeout(None, connect=5.0)
-        ) as client:
-            async_iterator = await aiter_sse_retrying(
-                client,
-                "GET",
-                f"{API_HOST}/dataset-qa/run/{run_id}",
-                headers=get_headers(),
-            )
-            async for sse in async_iterator:
-                if sse.event == "ping":
-                    continue
-                logger.debug(
-                    "[ASYNC] Received event: %s with data: %s and ID: %s and retry: %s",
-                    sse.event,
-                    sse.data,
-                    sse.id,
-                    sse.retry,
-                )
-                last_event = json.loads(sse.data)
-                yield last_event["data"]
-        if not last_event or last_event["data"]["state"] == RunStatus.PENDING.value:
-            QADataset.acheck_run_by_id(run_id, retry + 1)
+        async for iteration in aiter_run_events(
+            f"{API_HOST}/dataset-qa/run/{run_id}",
+            headers=get_headers(),
+            retry=retry,
+            status_keys=("state",),
+            error_cls=HirundoError,
+            log=logger,
+        ):
+            yield iteration
 
     async def acheck_run(self) -> AsyncGenerator[dict, None]:
         """
@@ -776,6 +688,8 @@ class QADataset(BaseModel):
 
         This generator will produce values to show progress of the run.
 
+        Note: This function does not handle errors nor show progress. It is expected that you do that.
+
         Yields:
             Each event will be a dict, where:
             - `"state"` is PENDING, STARTED, RETRY, FAILURE or SUCCESS
@@ -846,11 +760,11 @@ class QADatasetOut(BaseModel):
 
     data_root_url: HirundoUrl
 
-    classes: typing.Optional[list[str]] = None
-    labeling_info: typing.Union[LabelingInfo, list[LabelingInfo]]
+    classes: list[str] | None = None
+    labeling_info: LabelingInfo | list[LabelingInfo]
 
-    organization_id: typing.Optional[int]
-    creator_id: typing.Optional[int]
+    organization_id: int | None
+    creator_id: int | None
     created_at: datetime.datetime
     updated_at: datetime.datetime
 
@@ -863,4 +777,6 @@ class DataQARunOut(BaseModel):
     status: RunStatus
     approved: bool
     created_at: datetime.datetime
-    run_args: typing.Optional[RunArgs]
+    run_args: RunArgs | None
+
+    deleted_at: datetime.datetime | None = None

hirundo/dataset_qa_results.py

@@ -11,11 +11,11 @@ DataFrameType = TypeAliasType("DataFrameType", None)
 if has_pandas:
     from hirundo._dataframe import pd
 
-    DataFrameType = TypeAliasType("DataFrameType", typing.Union[pd.DataFrame, None])
+    DataFrameType = TypeAliasType("DataFrameType", pd.DataFrame | None)
 if has_polars:
     from hirundo._dataframe import pl
 
-    DataFrameType = TypeAliasType("DataFrameType", typing.Union[pl.DataFrame, None])
+    DataFrameType = TypeAliasType("DataFrameType", pl.DataFrame | None)
 
 
 T = typing.TypeVar("T")
@@ -32,7 +32,7 @@ class DatasetQAResults(BaseModel, typing.Generic[T]):
     """
     A polars/pandas DataFrame containing the results of the data QA run
     """
-    object_suspects: typing.Optional[T]
+    object_suspects: T | None
     """
     A polars/pandas DataFrame containing the object-level results of the data QA run
     """

hirundo/git.py

@@ -1,6 +1,5 @@
 import datetime
 import re
-import typing
 
 import pydantic
 from pydantic import BaseModel, field_validator
@@ -32,14 +31,14 @@ class GitSSHAuth(BaseModel):
     """
     The SSH key for the Git repository
     """
-    ssh_password: typing.Optional[str]
+    ssh_password: str | None
     """
     The password for the SSH key for the Git repository.
     """
 
 
 class GitRepo(BaseModel):
-    id: typing.Optional[int] = None
+    id: int | None = None
     """
     The ID of the Git repository.
     """
@@ -48,25 +47,25 @@ class GitRepo(BaseModel):
     """
     A name to identify the Git repository in the Hirundo system.
     """
-    repository_url: typing.Union[str, RepoUrl]
+    repository_url: str | RepoUrl
     """
     The URL of the Git repository, it should start with `ssh://` or `https://` or be in the form `user@host:path`.
     If it is in the form `user@host:path`, it will be rewritten to `ssh://user@host/path`.
     """
-    organization_id: typing.Optional[int] = None
+    organization_id: int | None = None
     """
     The ID of the organization that the Git repository belongs to.
     If not provided, it will be assigned to your default organization.
     """
 
-    plain_auth: typing.Optional[GitPlainAuth] = pydantic.Field(
+    plain_auth: GitPlainAuth | None = pydantic.Field(
         default=None, examples=[None, {"username": "ben", "password": "password"}]
     )
     """
     The plain authentication details for the Git repository.
     Use this if using a special user with a username and password for authentication.
     """
-    ssh_auth: typing.Optional[GitSSHAuth] = pydantic.Field(
+    ssh_auth: GitSSHAuth | None = pydantic.Field(
         default=None,
         examples=[
             {
@@ -84,7 +83,7 @@ class GitRepo(BaseModel):
 
     @field_validator("repository_url", mode="before", check_fields=True)
     @classmethod
-    def check_valid_repository_url(cls, repository_url: typing.Union[str, RepoUrl]):
+    def check_valid_repository_url(cls, repository_url: str | RepoUrl):
         # Check if the URL has the `@` and `:` pattern with a non-numeric section before the next slash
         match = re.match("([^@]+@[^:]+):([^0-9/][^/]*)/(.+)", str(repository_url))
         if match: