def update_stage_specific_k8s_values()

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

• 18 lines of code
• 5 McCabe index (conditional complexity)

    def update_stage_specific_k8s_values(self, values_template):
        """
        Update the k8s configs that is related to the current stage
        """
        values_template.trainingConfig.ntasksPerNode = self.stage_cfg.trainer.devices
        values_template.trainingConfig.nodes = self.stage_cfg.trainer.num_nodes
        choice_model_type, _ = self.get_stage_config_choice()
        if self.cfg.git.get("entry_script", None) is not None:
            # Override with entry script provided by the customer
            values_template.trainingConfig.scriptPath = self.cfg.git.entry_script
        else:
            values_template.trainingConfig.scriptPath = str(self._entry_script_path)

        if OmegaConf.select(self.cfg, "recipes.model.multi_modal", default=False):
            transformers_upgrade_cmd = "pip install transformers==4.45.2"
            values_template.trainingConfig.pre_script.append(transformers_upgrade_cmd)
        if OmegaConf.select(self.cfg, "recipes.model.model_type", default=False) == "deepseek_r1":
            transformers_upgrade_cmd = "pip install transformers==4.48.2"
            values_template.trainingConfig.pre_script.append(transformers_upgrade_cmd)
        if OmegaConf.select(self.cfg, "recipes.model.model_type", default=None) == "llama_v4":
            transformers_upgrade_cmd = "pip install transformers==4.51.3"
            values_template.trainingConfig.pre_script.append(transformers_upgrade_cmd)

        return values_template