Skip to content

StepResources

Bases: RuntimeParametersMixin, BaseModel

A class to define the resources assigned to a _Step.

Attributes:

Name Type Description
replicas RuntimeParameter[PositiveInt]

The number of replicas for the step.
cpus Optional[RuntimeParameter[PositiveInt]]

The number of CPUs assigned to each step replica.
gpus Optional[RuntimeParameter[PositiveInt]]

The number of GPUs assigned to each step replica.
memory Optional[RuntimeParameter[PositiveInt]]

The memory in bytes required for each step replica.
resources Optional[RuntimeParameter[Dict[str, int]]]

A dictionary containing the number of custom resources required for each step replica.
Source code in src/distilabel/steps/base.py 
 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
class StepResources(RuntimeParametersMixin, BaseModel):
    """A class to define the resources assigned to a `_Step`.

    Attributes:
        replicas: The number of replicas for the step.
        cpus: The number of CPUs assigned to each step replica.
        gpus: The number of GPUs assigned to each step replica.
        memory: The memory in bytes required for each step replica.
        resources: A dictionary containing the number of custom resources required for
            each step replica.
    """

    replicas: RuntimeParameter[PositiveInt] = Field(
        default=1, description="The number of replicas for the step."
    )
    cpus: Optional[RuntimeParameter[PositiveInt]] = Field(
        default=None, description="The number of CPUs assigned to each step replica."
    )
    gpus: Optional[RuntimeParameter[PositiveInt]] = Field(
        default=None, description="The number of GPUs assigned to each step replica."
    )
    memory: Optional[RuntimeParameter[PositiveInt]] = Field(
        default=None, description="The memory in bytes required for each step replica."
    )
    resources: Optional[RuntimeParameter[Dict[str, int]]] = Field(
        default=None,
        description="A dictionary containing names of custom resources and the"
        " number of those resources required for each step replica.",
    )