Method create_for

axlearn/cloud/gcp/node_pool_provisioner.py:95–211

Creates named node pools for the job.

(self, job: GKEJob)

class _JobProvisioningConfig(NamedTuple):
    """Extracted provisioning parameters for a TPU job.

    Attributes:
        acc_cfg: Accelerator config (provides num_replicas, instance_type, etc.).
        job_sys_property: System characteristics for the TPU type.
        topology: GKE topology string, or None if not applicable.
        use_spot_vm: Whether to use spot VMs.
        reservation: Reservation name (fully-qualified if cross-project), or None.
        location_hint: Location hint for node pool placement, or None.
        enable_tpu_ici_resiliency: Whether TPU ICI resiliency is enabled.
        enable_tpu_smart_repair: Whether TPU smart repair is enabled.
        job_priority: Job priority from the bastion jobspec, or None.
    """

    acc_cfg: AcceleratorConfig
    job_sys_property: _SystemCharacteristics
    topology: Optional[str]
    use_spot_vm: bool
    reservation: Optional[str]
    location_hint: Optional[str]
    enable_tpu_ici_resiliency: bool
    enable_tpu_smart_repair: bool
    job_priority: Optional[int]


class _NodePoolBatch(NamedTuple):
    """A batch of node pools to create or delete.

    Attributes:
        names: Node pool names (e.g. from construct_node_pool_name).
        additional_labels_list: Per-pool label dicts, parallel to names.
            Empty list for delete batches (labels are not used on delete).
    """

    names: list[str]
    additional_labels_list: list[dict[str, str]]


class TPUNodePoolProvisioner(NodePoolProvisioner):
    """TPU node pool provisioner."""

    def _get_job_provisioning_config(self, job: GKEJob) -> _JobProvisioningConfig:
        """Extracts and resolves all provisioning parameters from the job config.

        Args:
            job: The GKEJob whose builder config carries TPU accelerator settings.

        Returns:
            A _JobProvisioningConfig with all parameters needed by create_node_pools
            or delete_node_pools.

        Raises:
            TypeError: If the job's builder config is not a supported type.
        """
        job_cfg: GKEJob.Config = job.config
        builder_cfg: TPUReplicatedJob.Config | PathwaysLeaderWorkerTemplate.Config = job_cfg.builder