Skip to content

commands

zeus.optimizer.batch_size.server.batch_size_state.commands

Commands to use BatchSizeStateRepository.

ReadTrial

Bases: BatchSizeBase

Command to read a trial.

Equivalent to primary key of Trial.

Attributes:

Name Type Description
job_id

ID of job

batch_size

batch size of a given trial

trial_number int

number of trial

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
17
18
19
20
21
22
23
24
25
26
27
28
class ReadTrial(BatchSizeBase):
    """Command to read a trial.

    Equivalent to primary key of Trial.

    Attributes:
        job_id: ID of job
        batch_size: batch size of a given trial
        trial_number: number of trial
    """

    trial_number: int = Field(gt=0)

CreateTrialBase

Bases: BatchSizeBase

Base command to create trial.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class CreateTrialBase(BatchSizeBase):
    """Base command to create trial."""

    type: TrialType
    start_timestamp: datetime = Field(default_factory=datetime.now)
    status: TrialStatus = Field(default=TrialStatus.Dispatched, const=True)

    class Config:  # type: ignore
        """Model configuration.

        Make it immutable after it's created.
        """

        frozen = True

Config

Model configuration.

Make it immutable after it's created.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
38
39
40
41
42
43
44
class Config:  # type: ignore
    """Model configuration.

    Make it immutable after it's created.
    """

    frozen = True

CreateTrial

Bases: CreateTrialBase

Internal command to create trial.

trial_number is populate within ZeusService.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class CreateTrial(CreateTrialBase):
    """Internal command to create trial.

    trial_number is populate within ZeusService.
    """

    trial_number: int = Field(gt=0)

    class Config:  # type: ignore
        """Model configuration.

        Make it immutable after it's created.
        """

        frozen = True

    def to_orm(self) -> TrialTable:
        """Create an ORM object from pydantic model.

        Returns:
            `TrialTable`: ORM object representing the trial.
        """
        d = self.dict()
        t = TrialTable()
        for k, v in d.items():
            setattr(t, k, v)
        return t

Config

Model configuration.

Make it immutable after it's created.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
55
56
57
58
59
60
61
class Config:  # type: ignore
    """Model configuration.

    Make it immutable after it's created.
    """

    frozen = True

to_orm

to_orm()

Create an ORM object from pydantic model.

Returns:

Type Description
TrialTable

TrialTable: ORM object representing the trial.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
63
64
65
66
67
68
69
70
71
72
73
def to_orm(self) -> TrialTable:
    """Create an ORM object from pydantic model.

    Returns:
        `TrialTable`: ORM object representing the trial.
    """
    d = self.dict()
    t = TrialTable()
    for k, v in d.items():
        setattr(t, k, v)
    return t

CreateExplorationTrial

Bases: CreateTrialBase

Create a exploration.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
76
77
78
79
80
81
82
83
84
85
86
87
class CreateExplorationTrial(CreateTrialBase):
    """Create a exploration."""

    type: TrialType = Field(default=TrialType.Exploration, const=True)

    class Config:  # type: ignore
        """Model configuration.

        Make it immutable after it's created.
        """

        frozen = True

Config

Model configuration.

Make it immutable after it's created.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
81
82
83
84
85
86
87
class Config:  # type: ignore
    """Model configuration.

    Make it immutable after it's created.
    """

    frozen = True

CreateMabTrial

Bases: CreateTrialBase

Create a MAB trial.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
class CreateMabTrial(CreateTrialBase):
    """Create a MAB trial."""

    type: TrialType = Field(default=TrialType.MAB, const=True)

    class Config:  # type: ignore
        """Model configuration.

        Make it immutable after it's created.
        """

        frozen = True

Config

Model configuration.

Make it immutable after it's created.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
 95
 96
 97
 98
 99
100
101
class Config:  # type: ignore
    """Model configuration.

    Make it immutable after it's created.
    """

    frozen = True

CreateConcurrentTrial

Bases: CreateTrialBase

Create a exploration.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
104
105
106
107
108
109
110
111
112
113
114
115
class CreateConcurrentTrial(CreateTrialBase):
    """Create a exploration."""

    type: TrialType = Field(default=TrialType.Concurrent, const=True)

    class Config:  # type: ignore
        """Model configuration.

        Make it immutable after it's created.
        """

        frozen = True

Config

Model configuration.

Make it immutable after it's created.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
109
110
111
112
113
114
115
class Config:  # type: ignore
    """Model configuration.

    Make it immutable after it's created.
    """

    frozen = True

UpdateTrial

Bases: BatchSizeBase

Report the result of trial.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
class UpdateTrial(BatchSizeBase):
    """Report the result of trial."""

    trial_number: int = Field(gt=0)
    end_timestamp: datetime = Field(default_factory=datetime.now)
    status: TrialStatus
    time: Optional[float] = Field(default=None, ge=0)
    energy: Optional[float] = Field(default=None, ge=0)
    converged: Optional[bool] = None

    class Config:  # type: ignore
        """Model configuration.

        Make it immutable after it's created.
        """

        frozen = True

    @validator("status")
    def _check_status(cls, s: TrialStatus) -> TrialStatus:
        """Check if status is equal to Dispatched."""
        if s != TrialStatus.Dispatched:
            return s
        else:
            raise ValueError(
                f"{s} shouldn't be Dispatched since this is reporting the result."
            )

    @root_validator(skip_on_failure=True)
    def _validate_sanity(cls, values: dict[str, Any]) -> dict[str, Any]:
        """Validate result.

        We are checking
            - if status == Failed, time/energy/converged == None.
                else, time/energy/converged != None.
        """
        status: TrialStatus = values["status"]

        time: float | None = values["time"]
        energy: float | None = values["energy"]
        converged: bool | None = values["converged"]

        if status != TrialStatus.Failed and (
            time is None or energy is None or converged is None
        ):
            raise ValueError(
                f"Result is incomplete: time({time}), energy({energy}), converged({converged})"
            )

        return values

Config

Model configuration.

Make it immutable after it's created.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
128
129
130
131
132
133
134
class Config:  # type: ignore
    """Model configuration.

    Make it immutable after it's created.
    """

    frozen = True

_check_status

_check_status(s)

Check if status is equal to Dispatched.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
136
137
138
139
140
141
142
143
144
@validator("status")
def _check_status(cls, s: TrialStatus) -> TrialStatus:
    """Check if status is equal to Dispatched."""
    if s != TrialStatus.Dispatched:
        return s
    else:
        raise ValueError(
            f"{s} shouldn't be Dispatched since this is reporting the result."
        )

_validate_sanity

_validate_sanity(values)

Validate result.

We are checking - if status == Failed, time/energy/converged == None. else, time/energy/converged != None.

Source code in zeus/optimizer/batch_size/server/batch_size_state/commands.py
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
@root_validator(skip_on_failure=True)
def _validate_sanity(cls, values: dict[str, Any]) -> dict[str, Any]:
    """Validate result.

    We are checking
        - if status == Failed, time/energy/converged == None.
            else, time/energy/converged != None.
    """
    status: TrialStatus = values["status"]

    time: float | None = values["time"]
    energy: float | None = values["energy"]
    converged: bool | None = values["converged"]

    if status != TrialStatus.Failed and (
        time is None or energy is None or converged is None
    ):
        raise ValueError(
            f"Result is incomplete: time({time}), energy({energy}), converged({converged})"
        )

    return values