from enum import StrEnum
from typing import Any
from pydantic import BaseModel, ConfigDict
[docs]
class ProgressState(StrEnum):
"""The status of a bulk inference job."""
IN_PROGRESS = "in_progress"
FAILED = "failed"
SUCCEEDED = "succeeded"
INITIALIZED = "initialized"
RUNNING = "running"
RUNNING_FAILED = "running_failed"
ABORTED = "aborted"
[docs]
class NewBulkInferenceJobRequest(BaseModel):
"""A request to create a new bulk inference job.
:param model_project_id: The project ID of the project containing the model.
:type model_project_id: str
:param model_id: The model ID of the model to use for inference.
:type model_id: str
:param dataset_project_id: The project ID of the project containing the snapshot to infer on.
:type dataset_project_id: str
:param dataset_id: The dataset ID dataset containing the snapshot to infer on.
:type dataset_id: str
:param dataset_snapshot_id: The ID of the snapshot to infer on.
:type dataset_snapshot_id: str
:param dataset_snapshot_split: The split of the snapshot to infer on.
:type dataset_snapshot_split: str
:param inference_method: The inference method (AKA action) to use for inference. Only
needed for task types that support multiple inference methods. Defaults to "",
meaning the default action will be chosen for task types supporting only a single
action.
:type inference_method: str
:param evaluate_metrics: Whether to use Valor to perform an evalaution after inference
is complete. When ``evaluate_metrics`` is true, only snapshot datums with annotations
will be infered on. To infer over unannotated data, set this to false. Defaults to false.
:type evaluate_metrics: bool
:param batch_size: The batch size to use for inference. Datums will be sent to inference
servers in groups of this size.
:type batch_size: int
:param inference_kwargs: Extra args to pass to the bulk inference job.
:type inference_kwargs: dict[str, Any] | None
"""
model_project_id: str
model_id: str
dataset_project_id: str
dataset_id: str
dataset_snapshot_id: str
dataset_snapshot_split: str
inference_method: str = ""
evaluate_metrics: bool = False
batch_size: int = 1
inference_kwargs: dict[str, Any] | None = None
model_config = ConfigDict(protected_namespaces=())
[docs]
class BulkInferenceJob(BaseModel):
"""Information about a previously created bulk inference job.
:param execution_id: The ID of this bulk inference job.
:type execution_id: str
:param execution_status: The status of this job.
:type execution_status: str
:param num_computed: The number of inferences completed so far in this job.
:type num_computed: int | None
:param expected_total: The total number of datums to infer on in this job.
:type expected_total: int | None
"""
execution_id: str
execution_status: str
num_computed: int | None = None
expected_total: int | None = None
@property
def job_id(self):
"""The ID of this bulk inference job."""
return self.execution_id
@property
def job_status(self) -> ProgressState:
"""The status of this bulk inference job."""
return ProgressState(self.execution_status)