def _compute_mpi_vars()

in lib/ramble/ramble/experiment_set.py [0:0]

• 58 lines of code
• 24 McCabe index (conditional complexity)

    def _compute_mpi_vars(self, expander, variables):
        """Compute required MPI variables

        Perform computation of required MPI variables, including:
        - n_ranks
        - n_nodes
        - processes_per_node
        - n_threads
        """
        n_ranks = (
            variables[self.keywords.n_ranks] if self.keywords.n_ranks in variables.keys() else None
        )
        ppn = (
            variables[self.keywords.processes_per_node]
            if self.keywords.processes_per_node in variables.keys()
            else None
        )
        n_nodes = (
            variables[self.keywords.n_nodes] if self.keywords.n_nodes in variables.keys() else None
        )
        n_threads = (
            variables[self.keywords.n_threads]
            if self.keywords.n_threads in variables.keys()
            else None
        )

        def _expand(var):
            try:
                return int(var)
            except ValueError:
                return int(expander.expand_var(var))

        if n_ranks:
            n_ranks = _expand(n_ranks)
            if n_ranks <= 0:
                logger.error("n_ranks must be positive")

        if ppn:
            ppn = _expand(ppn)
            if ppn <= 0:
                logger.error("processes_per_node must be positive")

        if n_nodes:
            n_nodes = _expand(n_nodes)
            if n_nodes <= 0:
                logger.error("n_nodes must be positive")

        if n_threads:
            n_threads = _expand(n_threads)

        if n_ranks and ppn:
            test_n_nodes = math.ceil(n_ranks / ppn)

            if n_nodes and n_nodes < test_n_nodes:
                logger.error(
                    f"n_nodes in {self.experiment_namespace} is {n_nodes} "
                    f"and should be {test_n_nodes}"
                )
            elif not n_nodes:
                logger.debug(f"Defining n_nodes in {self.experiment_namespace}")
                variables[self.keywords.n_nodes] = test_n_nodes
        elif n_ranks and n_nodes:
            ppn = math.ceil(int(n_ranks) / int(n_nodes))
            logger.debug(f"Defining processes_per_node in {self.experiment_namespace}")
            variables[self.keywords.processes_per_node] = ppn
        elif ppn and n_nodes:
            n_ranks = ppn * n_nodes
            logger.debug(f"Defining n_ranks in {self.experiment_namespace}")
            variables[self.keywords.n_ranks] = n_ranks
        elif not n_nodes:
            variables[self.keywords.n_nodes] = 1

        if not n_threads:
            variables[self.keywords.n_threads] = 1