MolSSI / QCFractal
1 8
import datetime
2 8
from enum import Enum
3 8
from typing import Any, Dict, List, Optional
4

5 8
from pydantic import Field, validator
6 8
from qcelemental.models import ComputeError
7

8 8
from .common_models import ObjectId, ProtoModel
9

10

11 8
class TaskStatusEnum(str, Enum):
12
    """
13
    The state of a Task object. The states which are available are a finite set.
14
    """
15

16 8
    running = "RUNNING"
17 8
    waiting = "WAITING"
18 8
    error = "ERROR"
19 8
    complete = "COMPLETE"
20

21

22 8
class ManagerStatusEnum(str, Enum):
23
    """
24
    The state of a Queue Manager. The states which are available are a finite set.
25
    """
26

27 8
    active = "ACTIVE"
28 8
    inactive = "INACTIVE"
29

30

31 8
class PriorityEnum(int, Enum):
32
    """
33
    The priority of a Task. Higher priority will be pulled first. The priorities which are available are a finite set.
34
    """
35

36 8
    HIGH = 2
37 8
    NORMAL = 1
38 8
    LOW = 0
39

40

41 8
class BaseResultEnum(str, Enum):
42 8
    result = "result"
43 8
    procedure = "procedure"
44

45

46 8
class PythonComputeSpec(ProtoModel):
47 8
    function: str = Field(
48
        ..., description="The module and function name of a Python-callable to call. Of the form 'module.function'."
49
    )
50 8
    args: List[Any] = Field(
51
        ..., description="A List of positional arguments to pass into ``function`` in order they appear."
52
    )
53 8
    kwargs: Dict[str, Any] = Field(..., description="Dictionary of keyword arguments to pass into ``function``.")
54

55

56 8
class TaskRecord(ProtoModel):
57

58 8
    id: ObjectId = Field(None, description="The Database assigned Id of the Task, if it has been assigned yet.")
59

60 8
    spec: PythonComputeSpec = Field(..., description="The Python function specification for this Task.")
61 8
    parser: str = Field(..., description="The type of operation this is Task is. Can be 'single' or 'optimization'.")
62 8
    status: TaskStatusEnum = Field(TaskStatusEnum.waiting, description="What stage of processing this task is at.")
63

64
    # Compute blockers and prevention
65 8
    program: str = Field(
66
        ..., description="Name of the quantum chemistry program which must be present to execute this task."
67
    )
68 8
    procedure: Optional[str] = Field(
69
        None, description="Name of the procedure the compute platform must be able to perform to execute this task."
70
    )
71 8
    manager: Optional[str] = Field(None, description="The Queue Manager that evaluated this task.")
72

73
    # Sortables
74 8
    priority: PriorityEnum = Field(PriorityEnum.NORMAL, description=str(PriorityEnum.__doc__))
75 8
    tag: Optional[str] = Field(
76
        None,
77
        description="The optional tag assigned to this Task. Tagged tasks can only be pulled by Queue Managers which "
78
        "explicitly reference this tag. If no Tag is specified, any Queue Manager can pull this Task.",
79
    )
80
    # Link back to the base Result
81 8
    base_result: ObjectId = Field(
82
        ..., description="Reference to the output Result from this Task as it exists within the database."
83
    )
84

85
    # Modified data
86 8
    modified_on: datetime.datetime = Field(None, description="The last time this task was updated in the Database.")
87 8
    created_on: datetime.datetime = Field(None, description="The time when this task was created in the Database.")
88

89 8
    def __init__(self, **data):
90

91
        # Set datetime defaults if not present
92 8
        dt = datetime.datetime.utcnow()
93 8
        data.setdefault("modified_on", dt)
94 8
        data.setdefault("created_on", dt)
95

96 8
        super().__init__(**data)
97

98 8
    @validator("priority", pre=True)
99 3
    def munge_priority(cls, v):
100 8
        if isinstance(v, str):
101 0
            v = PriorityEnum[v.upper()]
102 8
        elif v is None:
103 8
            v = PriorityEnum.NORMAL
104 8
        return v
105

106 8
    @validator("program")
107 3
    def check_program(cls, v):
108 8
        return v.lower()
109

110 8
    @validator("procedure")
111 3
    def check_procedure(cls, v):
112 8
        if v:
113 8
            v = v.lower()
114 8
        return v

Read our documentation on viewing source code .

Loading