Skip to content

commands

zeus.optimizer.batch_size.server.job.commands

Commands to use JobStateRepository.

UpdateExpDefaultBs

Bases: BaseModel

Parameters to update the exploration default batch size.

Attributes:

Name Type Description
job_id str

Job Id.
exp_default_batch_size int

new default batch size to use.
Source code in zeus/optimizer/batch_size/server/job/commands.py 
15
16
17
18
19
20
21
22
23
24
class UpdateExpDefaultBs(BaseModel):
    """Parameters to update the exploration default batch size.

    Attributes:
        job_id: Job Id.
        exp_default_batch_size: new default batch size to use.
    """

    job_id: str
    exp_default_batch_size: int = Field(gt=0)

UpdateJobStage

Bases: BaseModel

Parameters to update the job stage.

Attributes:

Name Type Description
job_id str

Job Id.
stage Stage

Set it to MAB since we only go from Pruning to MAB.
Source code in zeus/optimizer/batch_size/server/job/commands.py 
27
28
29
30
31
32
33
34
35
36
class UpdateJobStage(BaseModel):
    """Parameters to update the job stage.

    Attributes:
        job_id: Job Id.
        stage: Set it to MAB since we only go from Pruning to MAB.
    """

    job_id: str
    stage: Stage = Field(Stage.MAB, const=True)

UpdateGeneratorState

Bases: BaseModel

Parameters to update the generator state.

Attributes:

Name Type Description
job_id str

Job Id.
state str

Generator state.
Source code in zeus/optimizer/batch_size/server/job/commands.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class UpdateGeneratorState(BaseModel):
    """Parameters to update the generator state.

    Attributes:
        job_id: Job Id.
        state: Generator state.
    """

    job_id: str
    state: str

    @validator("state")
    def _validate_state(cls, state: str) -> str:
        """Validate the sanity of state."""
        try:
            np.random.default_rng(1).__setstate__(json.loads(state))
            return state
        except (TypeError, ValueError) as err:
            raise ValueError(f"Invalid generator state ({state})") from err

_validate_state 

_validate_state(state)

Validate the sanity of state.

Source code in zeus/optimizer/batch_size/server/job/commands.py 
50
51
52
53
54
55
56
57
@validator("state")
def _validate_state(cls, state: str) -> str:
    """Validate the sanity of state."""
    try:
        np.random.default_rng(1).__setstate__(json.loads(state))
        return state
    except (TypeError, ValueError) as err:
        raise ValueError(f"Invalid generator state ({state})") from err

UpdateJobMinCost

Bases: BaseModel

Parameters to update the min training cost and corresponding batch size.

Attributes:

Name Type Description
job_id str

Job Id.
min_cost float

Min training cost.
min_cost_batch_size int

Corresponding batch size.
Source code in zeus/optimizer/batch_size/server/job/commands.py 
60
61
62
63
64
65
66
67
68
69
70
71
class UpdateJobMinCost(BaseModel):
    """Parameters to update the min training cost and corresponding batch size.

    Attributes:
        job_id: Job Id.
        min_cost: Min training cost.
        min_cost_batch_size: Corresponding batch size.
    """

    job_id: str
    min_cost: float = Field(ge=0)
    min_cost_batch_size: int = Field(gt=0)

CreateJob

Bases: GpuConfig, JobParams

Parameters to create a new job.

Attributes:

Name Type Description
exp_default_batch_size int

Exploration default batch size that is used during Pruning stage.
min_cost None

Min training cost observed. Initially, None.
min_cost_batch_size int

Batch size that has minimum training cost observed.
stage Stage

Stage of the job.
mab_random_generator_state Optional[str]

Generator state if mab_seed is not None. Otherwise, None.

For the rest of attributes, refer to JobParams[zeus.optimizer.batch_size.common.JobParams] and GpuConfig[zeus.optimizer.batch_size.common.GpuConfig]

Source code in zeus/optimizer/batch_size/server/job/commands.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
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
class CreateJob(GpuConfig, JobParams):
    """Parameters to create a new job.

    Attributes:
        exp_default_batch_size: Exploration default batch size that is used during Pruning stage.
        min_cost: Min training cost observed. Initially, None.
        min_cost_batch_size: Batch size that has minimum training cost observed.
        stage: Stage of the job.
        mab_random_generator_state: Generator state if mab_seed is not None. Otherwise, None.

    For the rest of attributes, refer to `JobParams`[zeus.optimizer.batch_size.common.JobParams] and `GpuConfig`[zeus.optimizer.batch_size.common.GpuConfig]
    """

    exp_default_batch_size: int
    min_cost: None = Field(None, const=True)
    min_cost_batch_size: int
    stage: Stage = Field(Stage.Pruning, const=True)
    mab_random_generator_state: Optional[str] = None

    class Config:
        """Model configuration.

        Make it immutable after creation.
        """

        frozen = True

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

        We are checking,
            - If mab seed and generator state is matching.
            - If default, exp_default, min batch sizes are correctly intialized.
            - If default batch size is in the list of batch sizes.
        """
        state: str | None = values["mab_random_generator_state"]
        mab_seed: int | None = values["mab_seed"]
        bss: list[int] = values["batch_sizes"]
        dbs: int = values["default_batch_size"]
        ebs: int = values["exp_default_batch_size"]
        mbs: int = values["min_cost_batch_size"]

        if mab_seed is not None:
            if state is None:
                raise ValueError("mab_seed is not none, but generator state is none")
            else:
                try:
                    np.random.default_rng(1).__setstate__(json.loads(state))
                except (TypeError, ValueError) as err:
                    raise ValueError(f"Invalid generator state ({state})") from err

        if not (dbs == ebs == mbs):
            raise ValueError(
                f"During initialization, default_batch_size({dbs}), exp_default_batch_size({ebs}), min_batch_size({mbs}) should be all the same"
            )
        if dbs not in bss:
            raise ValueError(
                f"default_batch_size({dbs}) is not in the batch size list({bss})"
            )

        return values

    @classmethod
    def from_job_config(cls, js: JobSpecFromClient) -> "CreateJob":
        """From JobConfig, instantiate `CreateJob`.

        Initialize generator state, exp_default_batch_size, and min_cost_batch_size.
        """
        d = js.dict()
        d["exp_default_batch_size"] = js.default_batch_size
        if js.mab_seed is not None:
            rng = np.random.default_rng(js.mab_seed)
            d["mab_random_generator_state"] = json.dumps(rng.__getstate__())
        d["min_cost_batch_size"] = js.default_batch_size
        return cls.parse_obj(d)

    def to_orm(self) -> JobTable:
        """Convert pydantic model `CreateJob` to ORM object Job."""
        d = self.dict()
        job = JobTable()
        for k, v in d.items():
            if k != "batch_sizes":
                setattr(job, k, v)
        job.batch_sizes = [
            BatchSizeTable(job_id=self.job_id, batch_size=bs) for bs in self.batch_sizes
        ]
        return job

Config

Model configuration.

Make it immutable after creation.

Source code in zeus/optimizer/batch_size/server/job/commands.py 
93
94
95
96
97
98
99
class Config:
    """Model configuration.

    Make it immutable after creation.
    """

    frozen = True

_validate_states 

_validate_states(values)

Validate Job states.

We are checking, - If mab seed and generator state is matching. - If default, exp_default, min batch sizes are correctly intialized. - If default batch size is in the list of batch sizes.

Source code in zeus/optimizer/batch_size/server/job/commands.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
@root_validator(skip_on_failure=True)
def _validate_states(cls, values: dict[str, Any]) -> dict[str, Any]:
    """Validate Job states.

    We are checking,
        - If mab seed and generator state is matching.
        - If default, exp_default, min batch sizes are correctly intialized.
        - If default batch size is in the list of batch sizes.
    """
    state: str | None = values["mab_random_generator_state"]
    mab_seed: int | None = values["mab_seed"]
    bss: list[int] = values["batch_sizes"]
    dbs: int = values["default_batch_size"]
    ebs: int = values["exp_default_batch_size"]
    mbs: int = values["min_cost_batch_size"]

    if mab_seed is not None:
        if state is None:
            raise ValueError("mab_seed is not none, but generator state is none")
        else:
            try:
                np.random.default_rng(1).__setstate__(json.loads(state))
            except (TypeError, ValueError) as err:
                raise ValueError(f"Invalid generator state ({state})") from err

    if not (dbs == ebs == mbs):
        raise ValueError(
            f"During initialization, default_batch_size({dbs}), exp_default_batch_size({ebs}), min_batch_size({mbs}) should be all the same"
        )
    if dbs not in bss:
        raise ValueError(
            f"default_batch_size({dbs}) is not in the batch size list({bss})"
        )

    return values

from_job_config classmethod 

from_job_config(js)

From JobConfig, instantiate CreateJob.

Initialize generator state, exp_default_batch_size, and min_cost_batch_size.

Source code in zeus/optimizer/batch_size/server/job/commands.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
@classmethod
def from_job_config(cls, js: JobSpecFromClient) -> "CreateJob":
    """From JobConfig, instantiate `CreateJob`.

    Initialize generator state, exp_default_batch_size, and min_cost_batch_size.
    """
    d = js.dict()
    d["exp_default_batch_size"] = js.default_batch_size
    if js.mab_seed is not None:
        rng = np.random.default_rng(js.mab_seed)
        d["mab_random_generator_state"] = json.dumps(rng.__getstate__())
    d["min_cost_batch_size"] = js.default_batch_size
    return cls.parse_obj(d)

to_orm 

to_orm()

Convert pydantic model CreateJob to ORM object Job.

Source code in zeus/optimizer/batch_size/server/job/commands.py 
151
152
153
154
155
156
157
158
159
160
161
def to_orm(self) -> JobTable:
    """Convert pydantic model `CreateJob` to ORM object Job."""
    d = self.dict()
    job = JobTable()
    for k, v in d.items():
        if k != "batch_sizes":
            setattr(job, k, v)
    job.batch_sizes = [
        BatchSizeTable(job_id=self.job_id, batch_size=bs) for bs in self.batch_sizes
    ]
    return job