def _internal_update_parameters()

in submitit/auto/auto.py [0:0]

• 50 lines of code
• 22 McCabe index (conditional complexity)

    def _internal_update_parameters(self, **kwargs: Any) -> None:
        """Updates submission parameters to srun/crun.

        Parameters
        ----------
        AutoExecutors provides shared parameters that are translated for each specific cluster.
        Those are: timeout_min (int), mem_gb (int), gpus_per_node (int), cpus_per_task (int),
        nodes (int), tasks_per_node (int) and name (str).
        Cluster specific parameters can be specified by prefixing them with the cluster name.

        Notes
        -----
        - Cluster specific parameters win over shared parameters.
            eg: if both `slurm_time` and `timeout_min` are provided, then:
                - `slurm_time` is used on the slurm cluster
                - `timeout_min` is used on other clusters
        """
        # We handle None as not set.
        kwargs = {k: v for k, v in kwargs.items() if v is not None}
        # check type of replaced variables
        generics = AutoExecutor._typed_parameters()
        for name, expected_type in generics.items():
            if expected_type == float:
                expected_type = (int, float)  # type: ignore
            if name in kwargs:
                assert isinstance(kwargs[name], expected_type), (
                    f'Parameter "{name}" expected type {expected_type} ' f'(but value: "{kwargs[name]}")'
                )

        _convert_deprecated_args(kwargs, self._deprecated_args)
        specific = [x.split("_", 1) for x in kwargs if x not in generics]

        invalid = []
        executors = plugins.get_executors()
        for ex_arg in specific:
            if len(ex_arg) != 2:
                invalid.append(f"Parameter '{ex_arg[0]}' need to be prefixed by an executor name.")
                continue
            ex, arg = ex_arg

            if ex not in executors:
                invalid.append(f"Unknown executor '{ex}' in parameter '{ex}_{arg}'.")
                continue

            valid = executors[ex]._valid_parameters()
            if arg not in valid and arg not in generics:
                invalid.append(
                    f"Unknown argument '{arg}' for executor '{ex}' in parameter '{ex}_{arg}'."
                    + " Valid arguments: "
                    + ", ".join(valid)
                )
                continue
        if invalid:
            invalid.extend(
                [
                    f"Known executors: {', '.join(executors.keys())}",
                    f"As a reminder, shared/generic (non-prefixed) parameters are: {generics}.",
                ]
            )
            raise NameError("\n".join(invalid))

        # add cluster specific generic overrides
        kwargs.update(
            **{
                arg: kwargs.pop(f"{ex}_{arg}")
                for ex, arg in specific
                if ex == self.cluster and arg in generics
            }
        )
        parameters = self._executor._convert_parameters({k: kwargs[k] for k in kwargs if k in generics})
        # update parameters in the core executor
        for (ex, arg) in specific:
            # update cluster specific non-generic arguments
            if arg not in generics and ex == self.cluster:
                parameters[arg] = kwargs[f"{ex}_{arg}"]

        self._executor._internal_update_parameters(**parameters)