Skip to content

StepResources

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
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.",
    )