Skip to content

models

zeus.optimizer.batch_size.server.job.models

Pydantic models for Job.

Stage

Bases: Enum

Job Stage.

Source code in zeus/optimizer/batch_size/server/job/models.py 
15
16
17
18
19
class Stage(Enum):
    """Job Stage."""

    Pruning = "Pruning"
    MAB = "MAB"

JobGetter

Bases: GetterDict

Getter for batch size to convert ORM batch size object to integer.

Source code in zeus/optimizer/batch_size/server/job/models.py 
22
23
24
25
26
27
28
29
30
31
class JobGetter(GetterDict):
    """Getter for batch size to convert ORM batch size object to integer."""

    def get(self, key: str, default: Any = None) -> Any:
        """Get value from dict."""
        if key == "batch_sizes":
            # If the key is batch_sizes, parse the integer from object.
            return [bs.batch_size for bs in self._obj.batch_sizes]

        return super().get(key, default)

get 

get(key, default=None)

Get value from dict.

Source code in zeus/optimizer/batch_size/server/job/models.py 
25
26
27
28
29
30
31
def get(self, key: str, default: Any = None) -> Any:
    """Get value from dict."""
    if key == "batch_sizes":
        # If the key is batch_sizes, parse the integer from object.
        return [bs.batch_size for bs in self._obj.batch_sizes]

    return super().get(key, default)

JobState

Bases: JobParams, GpuConfig

Pydantic model for Job which includes job-level states.

Attributes:

Name Type Description
exp_default_batch_size int

Exploration default batch size that is used during Pruning stage.
min_cost Optional[float]

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 and GpuConfig

Source code in zeus/optimizer/batch_size/server/job/models.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
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
74
75
76
77
78
79
80
class JobState(JobParams, GpuConfig):
    """Pydantic model for Job which includes job-level states.

    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: Optional[float] = None
    min_cost_batch_size: int
    stage: Stage = Stage.Pruning

    mab_random_generator_state: Optional[str] = None

    class Config:
        """Model configuration.

        Allow instantiating the model from an ORM object.
        """

        orm_mode = True
        getter_dict = JobGetter

    @root_validator(skip_on_failure=True)
    def _validate_mab(cls, values: dict[str, Any]) -> dict[str, Any]:
        """Validate generator state."""
        state: str | None = values["mab_random_generator_state"]
        mab_seed: int | None = values["mab_seed"]

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

        return values

Config

Model configuration.

Allow instantiating the model from an ORM object.

Source code in zeus/optimizer/batch_size/server/job/models.py 
55
56
57
58
59
60
61
62
class Config:
    """Model configuration.

    Allow instantiating the model from an ORM object.
    """

    orm_mode = True
    getter_dict = JobGetter

_validate_mab 

_validate_mab(values)

Validate generator state.

Source code in zeus/optimizer/batch_size/server/job/models.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@root_validator(skip_on_failure=True)
def _validate_mab(cls, values: dict[str, Any]) -> dict[str, Any]:
    """Validate generator state."""
    state: str | None = values["mab_random_generator_state"]
    mab_seed: int | None = values["mab_seed"]

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

    return values