Source code for tidy3d.web.core.task_info

"""Defines information about a task"""

from __future__ import annotations

from abc import ABC
from datetime import datetime
from enum import Enum
from typing import Optional

import pydantic.v1 as pydantic


class TaskBase(pydantic.BaseModel, ABC):
    """Base configuration for all task objects."""

    class Config:
        """Configuration for TaskBase"""

        arbitrary_types_allowed = True
        """Allow arbitrary types to be used within the model."""


class ChargeType(str, Enum):
    """The payment method of the task."""

    FREE = "free"
    """No payment required."""

    PAID = "paid"
    """Payment required."""


class TaskBlockInfo(TaskBase):
    """Information about the task's block status.

    This includes details about how the task can be blocked by various features
    such as user limits and insufficient balance.
    """

    chargeType: ChargeType = None
    """The type of charge applicable to the task (free or paid)."""

    maxFreeCount: int = None
    """The maximum number of free tasks allowed."""

    maxGridPoints: int = None
    """The maximum number of grid points permitted."""

    maxTimeSteps: int = None
    """The maximum number of time steps allowed."""




class TaskInfo(TaskBase):
    """General information about a task."""

    taskId: str
    """Unique identifier for the task."""

    taskName: str = None
    """Name of the task."""

    nodeSize: int = None
    """Size of the node allocated for the task."""

    completedAt: Optional[datetime] = None
    """Timestamp when the task was completed."""

    status: str = None
    """Current status of the task."""

    realCost: float = None
    """Actual cost incurred by the task."""

    timeSteps: int = None
    """Number of time steps involved in the task."""

    solverVersion: str = None
    """Version of the solver used for the task."""

    createAt: Optional[datetime] = None
    """Timestamp when the task was created."""

    estCostMin: float = None
    """Estimated minimum cost for the task."""

    estCostMax: float = None
    """Estimated maximum cost for the task."""

    realFlexUnit: float = None
    """Actual flexible units used by the task."""

    oriRealFlexUnit: float = None
    """Original real flexible units."""

    estFlexUnit: float = None
    """Estimated flexible units for the task."""

    estFlexCreditTimeStepping: float = None
    """Estimated flexible credits for time stepping."""

    estFlexCreditPostProcess: float = None
    """Estimated flexible credits for post-processing."""

    estFlexCreditMode: float = None
    """Estimated flexible credits based on the mode."""

    s3Storage: float = None
    """Amount of S3 storage used by the task."""

    startSolverTime: Optional[datetime] = None
    """Timestamp when the solver started."""

    finishSolverTime: Optional[datetime] = None
    """Timestamp when the solver finished."""

    totalSolverTime: int = None
    """Total time taken by the solver."""

    callbackUrl: str = None
    """Callback URL for task notifications."""

    taskType: str = None
    """Type of the task."""

    metadataStatus: str = None
    """Status of the metadata for the task."""

    taskBlockInfo: TaskBlockInfo = None
    """Blocking information for the task."""

    version: str = None
    """Version of the task."""



class RunInfo(TaskBase):
    """Information about the run of a task."""

    perc_done: pydantic.confloat(ge=0.0, le=100.0)
    """Percentage of the task that is completed (0 to 100)."""

    field_decay: pydantic.confloat(ge=0.0, le=1.0)
    """Field decay from the maximum value (0 to 1)."""

    def display(self) -> None:
        """Print some info about the task's progress."""
        print(f" - {self.perc_done:.2f} (%) done")
        print(f" - {self.field_decay:.2e} field decay from max")


# ---------------------- Batch (Modeler) detail schema ---------------------- #


class BatchTaskBlockInfo(TaskBlockInfo):
    """
    Extends `TaskBlockInfo` with specific details for batch task blocking.

    Attributes:
        accountLimit: A usage or cost limit imposed by the user's account.
        taskBlockMsg: A human-readable message describing the reason for the block.
        taskBlockType: The specific type of block (e.g., 'balance', 'limit').
        blockStatus: The current blocking status for the batch.
        taskStatus: The status of the task when it was blocked.
    """

    accountLimit: float = None
    taskBlockMsg: str = None
    taskBlockType: str = None
    blockStatus: str = None
    taskStatus: str = None


class BatchMember(TaskBase):
    """
    Represents a single task within a larger batch operation.

    Attributes:
        refId: A reference identifier for the member task.
        folderId: The identifier of the folder containing the task.
        sweepId: The identifier for the parameter sweep, if applicable.
        taskId: The unique identifier of the task.
        linkedTaskId: The identifier of a task linked to this one.
        groupId: The identifier of the group this task belongs to.
        taskName: The name of the individual task.
        status: The current status of this specific task.
        sweepData: Data associated with a parameter sweep.
        validateInfo: Information related to the task's validation.
        replaceData: Data used for replacements or modifications.
        protocolVersion: The version of the protocol used.
        variable: The variable parameter for this task in a sweep.
        createdAt: The timestamp when the member task was created.
        updatedAt: The timestamp when the member task was last updated.
        denormalizeStatus: The status of the data denormalization process.
        summary: A dictionary containing summary information for the task.
    """

    refId: str = None
    folderId: str = None
    sweepId: str = None
    taskId: str = None
    linkedTaskId: str = None
    groupId: str = None
    taskName: str = None
    status: str = None
    sweepData: str = None
    validateInfo: str = None
    replaceData: str = None
    protocolVersion: str = None
    variable: str = None
    createdAt: Optional[datetime] = None
    updatedAt: Optional[datetime] = None
    denormalizeStatus: str = None
    summary: dict = None


class BatchDetail(TaskBase):
    """
    Provides a detailed, top-level view of a batch of tasks.

    This model serves as the main payload for retrieving comprehensive
    information about a batch operation.

    Attributes:
        refId: A reference identifier for the entire batch.
        optimizationId: Identifier for the optimization process, if any.
        groupId: Identifier for the group the batch belongs to.
        name: The user-defined name of the batch.
        status: The current status of the batch.
        totalTask: The total number of tasks in the batch.
        preprocessSuccess: The count of tasks that completed preprocessing.
        postprocessStatus: The status of the batch's postprocessing stage.
        validateSuccess: The count of tasks that passed validation.
        runSuccess: The count of tasks that ran successfully.
        postprocessSuccess: The count of tasks that completed postprocessing.
        taskBlockInfo: Information on what might be blocking the batch.
        estFlexUnit: The estimated total flexible compute units for the batch.
        totalSeconds: The total time in seconds the batch has taken.
        totalCheckMillis: Total time in milliseconds spent on checks.
        message: A general message providing information about the batch status.
        tasks: A list of `BatchMember` objects, one for each task in the batch.
        taskType: The type of tasks contained in the batch.
    """

    refId: str = None
    optimizationId: str = None
    groupId: str = None
    name: str = None
    status: str = None
    totalTask: int = 0
    preprocessSuccess: int = 0
    postprocessStatus: str = None
    validateSuccess: int = 0
    runSuccess: int = 0
    postprocessSuccess: int = 0
    taskBlockInfo: BatchTaskBlockInfo = None
    estFlexUnit: float = None
    realFlexUnit: float = None
    totalSeconds: int = None
    totalCheckMillis: int = None
    message: str = None
    tasks: list[BatchMember] = []
    validateErrors: dict = None
    taskType: str = "RF"
    version: str = None


class AsyncJobDetail(TaskBase):
    """
    Provides a detailed view of an asynchronous job and its sub-tasks.

    This model represents a long-running operation. The 'result' attribute holds
    the output of a completed job, which for orchestration jobs, is often a
    JSON string mapping sub-task names to their unique IDs.

    Attributes:
        asyncId: The unique identifier for the asynchronous job.
        status: The current overall status of the job (e.g., 'RUNNING', 'COMPLETED').
        progress: The completion percentage of the job (from 0.0 to 100.0).
        createdAt: The timestamp when the job was created.
        completedAt: The timestamp when the job finished (successfully or not).
        tasks: A dictionary mapping logical task keys to their unique task IDs.
               This is often populated by parsing the 'result' of an orchestration task.
        result: The raw string output of the completed job. If the job spawns other
                tasks, this is expected to be a JSON string detailing those tasks.
        taskBlockInfo: Information on any dependencies blocking the job from running.
        message: A human-readable message about the job's status.
    """

    asyncId: str
    status: str
    progress: Optional[float] = None
    createdAt: Optional[datetime] = None
    completedAt: Optional[datetime] = None
    tasks: Optional[dict[str, str]] = None
    result: Optional[str] = None
    taskBlockInfo: Optional[TaskBlockInfo] = None
    message: Optional[str] = None


AsyncJobDetail.update_forward_refs()