def _make_cluster_parameters()

in launcher/nemo/stages.py [0:0]

• 12 lines of code
• 4 McCabe index (conditional complexity)

    def _make_cluster_parameters(self, cluster: str) -> Dict:
        """
        Make a cluster-specific parameters for jobs on different clusters.

        :param str cluster: i.e. `bcm`, `bcp`, `interactive`, etc.
        :return: a dictionary of cluster parameters, e.g. `ntasks_per_node`
        :rtype: Dict
        """
        with omegaconf.open_dict(self.cfg):
            # Patch self.cfg.cluster to align with
            # https://github.com/NVIDIA/NeMo-Framework-Launcher/blob/23.11/launcher_scripts/nemo_launcher/core/stages.py#L312
            origin_cluster = self.cfg.cluster
            self.cfg.cluster = self.cfg.cluster.cluster_config
            cluster_parameters = super()._make_cluster_parameters(cluster)
            cluster_type = origin_cluster.get("cluster_type")
            if cluster_type == "k8s":
                env_vars = cluster_parameters.get("env_vars")
                if env_vars and "SLURM_NTASKS_PER_NODE" in env_vars:
                    env_vars.pop("SLURM_NTASKS_PER_NODE")
            self.cfg.cluster = origin_cluster
        return cluster_parameters