Module audiocraft.utils.cluster

Utility functions for SLURM configuration and cluster settings.

Functions

def get_cluster_type(cluster_type: ClusterType | None = None) ‑> ClusterType | None

Expand source code

def get_cluster_type(
    cluster_type: tp.Optional[ClusterType] = None,
) -> tp.Optional[ClusterType]:
    if cluster_type is None:
        return _guess_cluster_type()

    return cluster_type

def get_slurm_parameters(cfg: omegaconf.dictconfig.DictConfig,
cluster_type: ClusterType | None = None) ‑> omegaconf.dictconfig.DictConfig

Expand source code

def get_slurm_parameters(
    cfg: omegaconf.DictConfig, cluster_type: tp.Optional[ClusterType] = None
) -> omegaconf.DictConfig:
    """Update SLURM parameters in configuration based on cluster type.
    If the cluster type is not specify, it infers it automatically.
    """
    from ..environment import AudioCraftEnvironment
    cluster_type = get_cluster_type(cluster_type)
    # apply cluster-specific adjustments
    if cluster_type == ClusterType.AWS:
        cfg["mem_per_gpu"] = None
        cfg["constraint"] = None
        cfg["setup"] = []
    elif cluster_type == ClusterType.RSC:
        cfg["mem_per_gpu"] = None
        cfg["setup"] = []
        cfg["constraint"] = None
        cfg["partition"] = "learn"
    slurm_exclude = AudioCraftEnvironment.get_slurm_exclude()
    if slurm_exclude is not None:
        cfg["exclude"] = slurm_exclude
    return cfg

Update SLURM parameters in configuration based on cluster type. If the cluster type is not specify, it infers it automatically.

Classes

class ClusterType (value, names=None, *, module=None, qualname=None, type=None, start=1)

Expand source code

class ClusterType(Enum):
    AWS = "aws"
    FAIR = "fair"
    RSC = "rsc"
    LOCAL_DARWIN = "darwin"
    DEFAULT = "default"  # used for any other cluster.

An enumeration.

Ancestors

• enum.Enum

Class variables

var AWS

var DEFAULT

var FAIR

var LOCAL_DARWIN

var RSC