diff --git a/src/access/config/esm1p6_layout_input.py b/src/access/config/esm1p6_layout_input.py new file mode 100644 index 0000000..f9784ad --- /dev/null +++ b/src/access/config/esm1p6_layout_input.py @@ -0,0 +1,603 @@ +import logging + +from access.config.layout_config import ( + convert_num_nodes_to_ncores, + find_layouts_with_maxncore, + return_layout_tuple, +) + +logger = logging.getLogger(__name__) + + +# The noqa comment is to suppress the complexity warning from ruff/flake8 +# The complexity of this function is high due to the nested loops and multiple conditionals. Some day +# I or someone else will refactor it to reduce the complexity. - MS 7th Oct, 2025 +def _generate_esm1p6_layout_from_core_counts( # noqa: C901 + min_atm_ncores: int, + max_atm_ncores: int, + ncores_for_atm_and_ocn: int, + ice_ncores: int, + min_ncores_needed: int, + *, + atm_ncore_delta: int = 2, + abs_maxdiff_nx_ny: int = 4, + prefer_atm_nx_greater_than_ny: bool = True, + prefer_mom_nx_greater_than_ny: bool = True, + prefer_atm_ncores_greater_than_mom_ncores: bool = True, + mom_ncores_over_atm_ncores_range: (float, float) = (0.75, 1.25), +) -> list: + """ + Returns a list of possible core layouts for the Atmosphere and Ocean for the ESM 1.6 PI config + + Parameters + ---------- + + min_atm_ncores : int, required + Minimum number of ATM cores to consider when generating layouts. + Must be at least 2 and less than or equal to max_atm_ncores. + + max_atm_ncores : int, required + Maximum number of ATM cores to consider when generating layouts. + Must be at least 2 and greater than or equal to min_atm_ncores. + + ncores_for_atm_and_ocn : int, required + Total number of cores available for ATM and MOM. + Must be at least 3 (2 for atm and 1 for mom). + + ice_ncores : int, required + Number of cores allocated to ICE. Must be at least 1. + + min_ncores_needed : int, required + Minimum number of cores that must be used by ATM, MOM and ICE combined. + Must be at least 3 + ice_ncores (2 for ATM, 1 for MOM and ``ice_ncores`` for ice). + Layouts using fewer cores will be discarded. + + atm_ncore_delta : int, optional, default=2 + Step size to use when iterating between min_atm_ncores and max_atm_ncores. + Must be a non-zero and positive integer. + + abs_maxdiff_nx_ny : int, optional, default=4 + Absolute max. of the difference between nx and ny (in the solved layout) to consider + when generating layouts. Must be a non-negative integer. + + Setting to 0 will return only square layouts. Applies to both ATM and MOM layouts. + + prefer_atm_nx_greater_than_ny : bool, optional, default=True + If True, only consider ATM layouts with nx >= ny. + + prefer_mom_nx_greater_than_ny : bool, optional, default=True + If True, only consider MOM layouts with nx >= ny. + + prefer_atm_ncores_greater_than_mom_ncores : bool, optional, default=True + If True, only consider layouts with ATM ncores >= MOM ncores. + + mom_ncores_over_atm_ncores_range : tuple of float, optional, default=(0.75, 1.25) + A tuple of two floats representing the minimum and maximum fractions of MOM ncores + over ATM ncores to consider when generating layouts. + """ + + min_atm_and_mom_ncores = 3 # atm requires min 2 ncores (2x1 layout), mom requires min 1 ncore (1x1 layout) + + if min_atm_ncores < 2 or max_atm_ncores < 2 or min_atm_ncores > max_atm_ncores: + raise ValueError(f"Invalid ATM ncores range. Got ({min_atm_ncores}, {max_atm_ncores}) instead") + + if atm_ncore_delta <= 0: + raise ValueError( + "Stepsize in core counts to cover min. and max. ATM ncores must be a positive integer. " + f"Got {atm_ncore_delta} instead" + ) + + if ncores_for_atm_and_ocn < min_atm_and_mom_ncores: + raise ValueError( + "Number of cores available for ATM and OCN must be at least {min_atm_and_mom_ncores} " + f"(2 for atm and 1 for mom). Got {ncores_for_atm_and_ocn} instead" + ) + + if ice_ncores < 1: + raise ValueError(f"ice_ncores must be at least 1. Got {ice_ncores} instead") + + if min_ncores_needed > (ncores_for_atm_and_ocn + ice_ncores): + raise ValueError( + f"Min. number of cores needed ({min_ncores_needed}) cannot be greater than the total " + f"number of available cores ({ncores_for_atm_and_ocn + ice_ncores})" + ) + + if min_ncores_needed < ncores_for_atm_and_ocn: + logger.warning( + f"Min. total cores required for a valid config ({min_ncores_needed}) should be greater " + f"than the number of ATM + OCN cores ({ncores_for_atm_and_ocn}). " + f"Currently, any config that satisfies the ATM + OCN core requirements will also satisfy " + "the requirement for the min. total cores" + ) + + if ( + mom_ncores_over_atm_ncores_range[0] <= 0.0 + or mom_ncores_over_atm_ncores_range[1] <= 0.0 + or mom_ncores_over_atm_ncores_range[0] > mom_ncores_over_atm_ncores_range[1] + ): + raise ValueError( + f"Invalid MOM ncores over ATM ncores fractions. Got {mom_ncores_over_atm_ncores_range} instead" + ) + + layout_tuple = return_layout_tuple() + all_layouts = [] + logger.debug( + f"Generating layouts with {min_atm_ncores=}, {max_atm_ncores=}, {atm_ncore_delta=}, " + f"{ncores_for_atm_and_ocn=}, {ice_ncores=}, {min_ncores_needed=}, " + f"{mom_ncores_over_atm_ncores_range=}, {abs_maxdiff_nx_ny=}, " + f"{prefer_atm_nx_greater_than_ny=}, {prefer_mom_nx_greater_than_ny=}, " + f"{prefer_atm_ncores_greater_than_mom_ncores=}" + ) + for atm_ncores in range(min_atm_ncores, max_atm_ncores + 1, atm_ncore_delta): + logger.debug(f"Trying atm_ncores = {atm_ncores}") + atm_layout = find_layouts_with_maxncore( + atm_ncores, + abs_maxdiff_nx_ny=abs_maxdiff_nx_ny, + even_nx=True, + prefer_nx_greater_than_ny=prefer_atm_nx_greater_than_ny, + ) + if not atm_layout: + continue + + logger.debug(f" Found {len(atm_layout)} atm layouts for atm_ncores = {atm_ncores}: {atm_layout}") + + min_mom_ncores = int(atm_ncores * mom_ncores_over_atm_ncores_range[0]) + max_mom_ncores = int(atm_ncores * mom_ncores_over_atm_ncores_range[1]) + for atm in atm_layout: + atm_nx, atm_ny = atm + + mom_ncores = ncores_for_atm_and_ocn - atm_nx * atm_ny + logger.debug(f" Trying atm layout {atm_nx}x{atm_ny} with {atm_nx * atm_ny} ncores") + mom_layout = find_layouts_with_maxncore( + mom_ncores, + abs_maxdiff_nx_ny=abs_maxdiff_nx_ny, + prefer_nx_greater_than_ny=prefer_mom_nx_greater_than_ny, + ) + if not mom_layout: + continue + + # filter mom_layout to only include layouts with ncores in the range [min_mom_ncores, max_mom_ncores] + layout = [] + for mom_nx, mom_ny in mom_layout: + mom_ncores = mom_nx * mom_ny + if mom_ncores < min_mom_ncores or mom_ncores > max_mom_ncores: + logger.debug( + f"Skipping mom layout {mom_nx}x{mom_ny} with {mom_ncores} ncores " + f"not in the range [{min_mom_ncores}, {max_mom_ncores}]" + ) + continue + + if prefer_atm_ncores_greater_than_mom_ncores and (atm_nx * atm_ny < mom_ncores): + logger.debug( + f"Skipping mom layout since mom ncores = {mom_nx}x{mom_ny} is not less " + f"than atm ncores = {atm_nx * atm_ny}" + ) + continue + + ncores_used = mom_nx * mom_ny + atm_nx * atm_ny + ice_ncores + if ncores_used < min_ncores_needed: + logger.debug( + f"Skipping layout atm {atm_nx}x{atm_ny} mom {mom_nx}x{mom_ny} ice {ice_ncores}, " + f"with {ncores_used=} is less than {min_ncores_needed=}" + ) + continue + + logger.debug( + f"Adding layout atm {atm_nx}x{atm_ny} mom {mom_nx}x{mom_ny} ice {ice_ncores} with {ncores_used=}" + ) + layout.append(layout_tuple(ncores_used, atm_nx, atm_ny, mom_nx, mom_ny, ice_ncores)) + + # create a set of layouts to avoid duplicates + all_layouts.extend(set(layout)) + + if all_layouts: + # sort the layouts by ncores_used (descending, fewer wasted cores first), and then + # the sum of the absolute differences between nx and ny for atm and mom (ascending, + # i.e., more square layouts first) + all_layouts = sorted( + all_layouts, key=lambda x: (-x.ncores_used, abs(x.atm_nx - x.atm_ny) + abs(x.mom_nx - x.mom_ny)) + ) + + return all_layouts + + +# The noqa comment is to suppress the complexity warning from ruff/flake8 +# The complexity of this function is high due to the nested loops and multiple conditionals. Some day +# I or someone else will refactor it to reduce the complexity. - MS 7th Oct, 2025 +def generate_esm1p6_core_layouts_from_node_count( # noqa: C901 + num_nodes_list: float, + *, + queue: str = "normalsr", + tol_around_ctrl_ratio: float = None, + atm_ncore_delta: int = 2, + prefer_atm_nx_greater_than_ny: bool = True, + prefer_mom_nx_greater_than_ny: bool = True, + prefer_atm_ncores_greater_than_mom_ncores: bool = True, + abs_maxdiff_nx_ny: int = 4, + mom_ncores_over_atm_ncores_range: (float, float) = (0.75, 1.25), + max_wasted_ncores_frac: float = 0.01, + allocate_unused_cores_to_ice: bool = False, +) -> list: + """ + Given a list of target number of nodes to use, this function generates + possible core layouts for the Atmosphere and Ocean for the ESM 1.6 PI config. + + Parameters + ---------- + num_nodes_list : scalar or a list of integer/floats, required + A positive number or a list of positive numbers representing the number of nodes to use. + + queue : str, optional, default="normalsr" + Queue name on ``gadi``. Allowed values are "normalsr" and "normal". + + tol_around_ctrl_ratio : float, optional, default=None + If set, the min and max fractions of MOM ncores over ATM ncores will be set to (at most) within + (1 ± tol_around_ctrl_ratio) of the released PI config. Must be in the range [0.0, 1.0]. + If not set, the min and max fractions of MOM ncores over ATM ncores are used from the + ``mom_ncores_over_atm_ncores_range`` parameter. + + mom_ncores_over_atm_ncores_range : tuple of float, optional, default=(0.75, 1.25) + A tuple of two floats representing the minimum and maximum fractions of MOM ncores over ATM + ncores to consider when generating layouts. Must be greater than 0.0, and the second + value (i.e, the max.) must be at least equal to the first value (i.e., the min.) + Layouts with MOM ncores over ATM ncores outside this range will be discarded. + + *Note*: This parameter is ignored if ``tol_around_ctrl_ratio`` is set. + + atm_ncore_delta : int, optional, default=100 + Number of steps to take between the min. and max. ATM ncores when generating layouts. + Must be a positive integer. + + prefer_atm_nx_greater_than_ny : bool, optional, default=True + If True, only consider ATM layouts with nx >= ny. + + prefer_mom_nx_greater_than_ny : bool, optional, default=True + If True, only consider MOM layouts with nx >= ny. + + prefer_atm_ncores_greater_than_mom_ncores : bool, optional, default=True + If True, only consider layouts with ATM ncores >= MOM ncores. + + abs_maxdiff_nx_ny : int, optional, default=4 + Absolute max. of the difference between nx and ny (in the solved layout) to + consider when generating layouts. Must be a non-negative integer. + + allocate_unused_cores_to_ice : bool, optional, default=False + If True, any unused cores (i.e., total cores - atm cores - mom cores) will be allocated to ice. + + max_wasted_ncores_frac : float, optional, default=0.01 + Maximum fraction of wasted cores (i.e. not used by atm, mom or ice) to allow when generating layouts. + Must be in the range [0.0, 1.0]. + + Returns + ------- + list + A list of lists of layout_tuples. Each inner list corresponds to the layouts for the respective + number of nodes in ``num_nodes_list``. Each layout_tuple has the following fields: + - ncores_used : int + - atm_nx : int + - atm_ny : int + - mom_nx : int + - mom_ny : int + - ice_ncores : int + + An empty list is returned for a given number of nodes if no valid layouts could be generated. + + Raises + ------ + ValueError + If any of the input parameters are invalid. + + Notes + ----- + - atm requires nx to be even -> atm requires min 2 ncores (2x1 layout), + mom requires min 1 ncore (1x1 layout), ice requires min 1 ncore + - The released configuration used is: + - atm: 16x13 (208 cores) + - ocn: 14x14 (196 cores) + - ice: 12 (12 cores) + - queue: normalsr + - num_nodes: 4 (416 cores) + """ + + if tol_around_ctrl_ratio is None and mom_ncores_over_atm_ncores_range is None: + raise ValueError("Either tol_around_ctrl_ratio or mom_ncores_over_atm_ncores_range must be provided") + + if tol_around_ctrl_ratio is not None: + if tol_around_ctrl_ratio < 0.0: + raise ValueError( + f"The tolerance fraction for setting the MOM to ATM core ratio to the control ratio must " + f"be in >= 0.0. Got {tol_around_ctrl_ratio} instead" + ) + else: # tol_around_ctrl_ratio is None -> use mom_ncores_over_atm_ncores_range + if not isinstance(mom_ncores_over_atm_ncores_range, tuple) or len(mom_ncores_over_atm_ncores_range) != 2: + raise ValueError( + f"The min. and max. fraction of MOM ncores over ATM ncores must be a tuple of two floats. " + f"Got {mom_ncores_over_atm_ncores_range} instead" + ) + + min_frac_mom_ncores_over_atm_ncores, max_frac_mom_ncores_over_atm_ncores = mom_ncores_over_atm_ncores_range + if not min_frac_mom_ncores_over_atm_ncores or not max_frac_mom_ncores_over_atm_ncores: + raise ValueError( + f"The min. and max. fraction of MOM ncores over ATM ncores must be a valid float. Got " + f"min={min_frac_mom_ncores_over_atm_ncores} and max={max_frac_mom_ncores_over_atm_ncores} instead" + ) + + if min_frac_mom_ncores_over_atm_ncores <= 0: + raise ValueError( + "The minimum fraction of MOM ncores over ATM ncores must be greater than 0. " + f"Got {min_frac_mom_ncores_over_atm_ncores} instead" + ) + + if max_frac_mom_ncores_over_atm_ncores <= 0: + raise ValueError( + f"The maximum fraction of MOM ncores over ATM ncores must be greater than 0. " + f"Got {max_frac_mom_ncores_over_atm_ncores} instead" + ) + + if min_frac_mom_ncores_over_atm_ncores > max_frac_mom_ncores_over_atm_ncores: + raise ValueError( + f"Invalid MOM ncores over ATM ncores fractions - min. must be <= max." + f" Got min={min_frac_mom_ncores_over_atm_ncores} and " + f"max={max_frac_mom_ncores_over_atm_ncores} instead" + ) + + if atm_ncore_delta <= 0: + raise ValueError( + f"The stepsize in core counts to take between the min. and max. ATM ncores must " + f"be a positive integer. Got {atm_ncore_delta} instead" + ) + + if abs_maxdiff_nx_ny < 0: + raise ValueError( + "The absolute max. diff. between nx and ny (in the solved layout) must be a non-zero integer. " + f"Got {abs_maxdiff_nx_ny} instead" + ) + + if max_wasted_ncores_frac < 0 or max_wasted_ncores_frac > 1.0: + raise ValueError( + "The max. fraction of wasted cores must be in the range [0.0, 1.0]. Got {max_wasted_ncores_frac} instead" + ) + + if not isinstance(num_nodes_list, list): + num_nodes_list = [num_nodes_list] + + if any(not isinstance(n, (int, float)) for n in num_nodes_list): + raise ValueError(f"Number of nodes must be a float or an integer. Got {num_nodes_list} instead") + + if any(n <= 0 for n in num_nodes_list): + raise ValueError(f"Number of nodes must be > 0. Got {num_nodes_list} instead") + + # atm requires nx to be even -> atm requires min 2 ncores + # (2x1 layout), mom requires min 1 ncore (1x1 layout), ice requires min 1 ncore + min_cores_required = 2 + 1 + 1 + + ctrl_num_nodes, ctrl_config, ctrl_queue = 4, {"atm": (16, 13), "ocn": (14, 14), "ice": (12)}, "normalsr" + ctrl_ratio_mom_over_atm = ( + ctrl_config["ocn"][0] * ctrl_config["ocn"][1] / (ctrl_config["atm"][0] * ctrl_config["atm"][1]) + ) + ctrl_totncores = convert_num_nodes_to_ncores(ctrl_num_nodes, queue=ctrl_queue) + ctrl_ice_ncores = ctrl_config["ice"] + + if tol_around_ctrl_ratio is not None: + logger.debug( + "The min and max fractions of MOM ncores over ATM ncores will be set to " + f"(at most) within (1 \u00b1 {tol_around_ctrl_ratio})" + f" of the control ratio={ctrl_ratio_mom_over_atm:0.3g}" + ) + + # create the named tuple for holding the layouts + layout_tuple = return_layout_tuple() + final_layouts = [] + for num_nodes in num_nodes_list: + # cast num_nodes to int if it is an integer value + if isinstance(num_nodes, float) and num_nodes.is_integer(): + num_nodes = int(num_nodes) + + totncores = convert_num_nodes_to_ncores(num_nodes, queue=queue) + if totncores < min_cores_required: + logger.warning( + f"Total ncores = {totncores} is less than the min. ncores required = {min_cores_required}. Skipping" + ) + final_layouts.append([]) + continue + + logger.debug(f"Generating layouts for {num_nodes = } nodes") + ice_ncores = max(1, int(ctrl_ice_ncores / ctrl_totncores * totncores)) + + ncores_left = totncores - ice_ncores + max_wasted_ncores = int(totncores * max_wasted_ncores_frac) + min_ncores_needed = totncores - max_wasted_ncores + + # atm + mom = totncores - cice_ncores + # => 1 + mom/atm = (totncores - cice_ncores)/atm + # => 1 + min_frac_mom_ncores_over_atm_ncores = (totncores - cice_ncores)/max_atm_ncores + # => max_atm_ncores = (totncores - cice_ncores)/(1 + min_frac_mom_ncores_over_atm_ncores) + if tol_around_ctrl_ratio is not None: + target_atm_ncores = ncores_left / ( + 1.0 + ctrl_ratio_mom_over_atm + ) # intentionally not converted to int here -> done in the next lines + max_atm_ncores = max( + 2, int(target_atm_ncores * (1.0 + tol_around_ctrl_ratio)) + ) # allow some variation around the control ratio + min_atm_ncores = max( + 2, int(target_atm_ncores * (1.0 - tol_around_ctrl_ratio)) + ) # allow some variation around the control ratio + mom_ncores_over_atm_ncores_range = ( + ctrl_ratio_mom_over_atm * (1.0 - tol_around_ctrl_ratio), + ctrl_ratio_mom_over_atm * (1.0 + tol_around_ctrl_ratio), + ) + logger.debug( + f"The min. and max. frac. of MOM ncores over ATM ncores are set to {mom_ncores_over_atm_ncores_range} " + f"based on the control ratio={ctrl_ratio_mom_over_atm:0.3g} and {tol_around_ctrl_ratio=}" + ) + else: + max_atm_ncores = max(2, int(ncores_left / (1.0 + min_frac_mom_ncores_over_atm_ncores))) + min_atm_ncores = max(2, int(ncores_left / (1.0 + max_frac_mom_ncores_over_atm_ncores))) + mom_ncores_over_atm_ncores_range = ( + min_frac_mom_ncores_over_atm_ncores, + max_frac_mom_ncores_over_atm_ncores, + ) + logger.debug( + f"The min. and max. frac. of MOM ncores over ATM ncores are set to {mom_ncores_over_atm_ncores_range} " + f"based on the provided mom_ncores_over_atm_ncores_range={mom_ncores_over_atm_ncores_range}" + ) + + # If we want ATM ncores to be >= MOM ncores, then the max. fraction of MOM ncores over ATM ncores must be <= 1.0 + if prefer_atm_ncores_greater_than_mom_ncores: + mom_ncores_over_atm_ncores_range = (mom_ncores_over_atm_ncores_range[0], 1.0) + + logger.debug(f"ATM ncores range, steps = ({min_atm_ncores}, {max_atm_ncores}, {atm_ncore_delta})") + logger.debug(f"MOM ncores range = ({ncores_left - max_atm_ncores}, {ncores_left - min_atm_ncores})") + layout = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=min_atm_ncores, + max_atm_ncores=max_atm_ncores, + atm_ncore_delta=atm_ncore_delta, + prefer_atm_nx_greater_than_ny=prefer_atm_nx_greater_than_ny, + prefer_mom_nx_greater_than_ny=prefer_mom_nx_greater_than_ny, + prefer_atm_ncores_greater_than_mom_ncores=prefer_atm_ncores_greater_than_mom_ncores, + mom_ncores_over_atm_ncores_range=mom_ncores_over_atm_ncores_range, + abs_maxdiff_nx_ny=abs_maxdiff_nx_ny, + ncores_for_atm_and_ocn=ncores_left, + ice_ncores=ice_ncores, + min_ncores_needed=min_ncores_needed, + ) + + if allocate_unused_cores_to_ice and layout: + # update the ice_ncores in each layout to include any unused cores + updated_layouts = [ + layout_tuple( + totncores, + x.atm_nx, + x.atm_ny, + x.mom_nx, + x.mom_ny, + x.ice_ncores + (totncores - x.ncores_used), + ) + for x in layout + ] + layout = list(set(updated_layouts)) + # sort the layouts by ncores_used (descending, fewer wasted cores first), and then + # the sum of the absolute differences between nx and ny for atm and mom (ascending, i.e., + # more square layouts first) + layout = sorted(layout, key=lambda x: (-x.ncores_used, abs(x.atm_nx - x.atm_ny) + abs(x.mom_nx - x.mom_ny))) + + final_layouts.append(layout) + + logger.info(f"Generated a total of {len(final_layouts)} layouts for {num_nodes_list} nodes") + return final_layouts + + +def generate_esm1p6_perturb_block( + num_nodes: (float | int), + layouts: list, + branch_name_prefix: str, + *, + queue: str = "normalsr", + start_blocknum: int = 1, +) -> str: + """ + + Generates a block for "perturbation" experiments in the ESM 1.6 PI config. + + Parameters + ---------- + num_nodes : float or int, required + A positive number representing the number of nodes to use. + + layouts : list, required + A list of layout_tuples as returned by ``generate_esm1p6_core_layouts_from_node_count``. + Each layout_tuple has the following fields: + - ncores_used : int + - atm_nx : int + - atm_ny : int + - mom_nx : int + - mom_ny : int + - ice_ncores : int + + The layouts will be used in the order they appear in the list. + + branch_name_prefix : str, required + Prefix to use for the branch names in the generated block. + + queue : str, optional, default="normalsr" + Queue name on ``gadi``. Allowed values are "normalsr" and "normal". + + start_blocknum : int, optional, default=1 + The starting block number to use in the generated block. Must be a positive integer greater than 0. + + Returns + ------- + str + A string representing the generated block. + + Raises + ------ + ValueError + If any of the input parameters are invalid. + + """ + + if num_nodes is None: + raise ValueError("num_nodes must be provided.}") + + if not isinstance(num_nodes, (int, float)) or num_nodes <= 0: + raise ValueError( + f"Number of nodes must be a positive number or a list of positive numbers. Got {num_nodes} instead" + ) + + if branch_name_prefix is None: + raise ValueError("branch_name_prefix must be provided") + + if not layouts: + raise ValueError("No layouts provided") + + if not isinstance(layouts, list): + layouts = [layouts] + + if any(len(x) != 6 for x in layouts): + raise ValueError(f"Invalid layouts provided. Layouts = {layouts}, {len(layouts[0])=} instead of 6") + + if not start_blocknum or start_blocknum < 1: + raise ValueError("start_blocknum must be a positive integer greater than 0") + + totncores = convert_num_nodes_to_ncores(num_nodes, queue=queue) + blocknum = start_blocknum + block = "" + for layout in layouts: + atm_nx, atm_ny = layout.atm_nx, layout.atm_ny + mom_nx, mom_ny = layout.mom_nx, layout.mom_ny + ice_ncores = layout.ice_ncores + atm_ncores = atm_nx * atm_ny + mom_ncores = mom_nx * mom_ny + branch_name = f"{branch_name_prefix}_atm_{atm_nx}x{atm_ny}_mom_{mom_nx}x{mom_ny}_ice_{ice_ncores}x1" + ncores_used = atm_ncores + mom_ncores + ice_ncores + block += f""" + Scaling_numnodes_{num_nodes}_totncores_{totncores}_ncores_used_{ncores_used}_seqnum_{blocknum}: + branches: + - {branch_name} + config.yaml: + submodels: + - - ncpus: # atmosphere + - {atm_ncores} # ncores for atmosphere + - ncpus: # ocean + - {mom_ncores} # ncores for ocean + - ncpus: # ice + - {ice_ncores} # ncores for ice + + atmosphere/um_env.yaml: + UM_ATM_NPROCX: {atm_nx} + UM_ATM_NPROCY: {atm_ny} + UM_NPES: {atm_ncores} + + ocean/input.nml: + ocean_model_nml: + layout: + - {mom_nx},{mom_ny} + + ice/cice_in.nml: + domain_nml: + - {ice_ncores} + """ + blocknum += 1 + + return block, blocknum diff --git a/src/access/config/layout_config.py b/src/access/config/layout_config.py new file mode 100644 index 0000000..7d00af0 --- /dev/null +++ b/src/access/config/layout_config.py @@ -0,0 +1,137 @@ +from typing import NamedTuple + + +def return_layout_tuple() -> NamedTuple: + """ + Define a named tuple to hold layout information. + Returns + ------- + NamedTuple + A named tuple with fields: + - ncores_used (int): Total number of cores used. + - atm_nx (int): Number of cores in the x-direction for the atmosphere model. + - atm_ny (int): Number of cores in the y-direction for the atmosphere model. + - mom_nx (int): Number of cores in the x-direction for the ocean model. + - mom_ny (int): Number of cores in the y-direction for the ocean model. + - ice_ncores (int): Number of cores used for the ice model. + """ + # The noqa comment is to suppress the "convert to class" warning from ruff/flake8 + layout_tuple = NamedTuple( # noqa: UP014 + "layout_tuple", + [ + ("ncores_used", int), # This can be derived from the other fields -> perhaps remove it? MS: 3rd Oct, 2025 + ("atm_nx", int), + ("atm_ny", int), + ("mom_nx", int), + ("mom_ny", int), + ("ice_ncores", int), + ], + ) + return layout_tuple + + +def convert_num_nodes_to_ncores(num_nodes: (int | float), queue: str = "normalsr") -> int: + """ + Convert number of nodes to number of cores based on queue properties. + + Parameters + ---------- + num_nodes : int or float, required + Number of nodes to convert. Must be a positive number. + queue : str, optional + Queue name. Allowed values are "normalsr" and "normal". + Default is "normalsr". + Returns + ------- + int + Total number of cores corresponding to the given number of nodes. + + Raises + ------ + ValueError + If the queue name is not recognized or if num_nodes is not a positive number. + + """ + queue_properties = { + "normalsr": {"ncores_per_node": 104}, + "normal": {"ncores_per_node": 48}, + } + if queue not in list(queue_properties.keys()): + raise ValueError(f"Queue = {queue} not allowed. Allowed values are {list(queue_properties.keys())}") + + if not isinstance(num_nodes, (int, float)) or num_nodes <= 0: + raise ValueError("Number of nodes must be a positive number (integer or float).") + + return int(num_nodes * queue_properties[queue]["ncores_per_node"]) + + +def find_layouts_with_maxncore( + maxncore: int, + *, # keyword-only arguments follow + abs_maxdiff_nx_ny: int = 4, + even_nx: bool = False, + prefer_nx_greater_than_ny: bool = False, +) -> list: + """ + Find possible (nx, ny) layouts for a given maximum number of cores (maxncore). + + The function returns a list of tuples (nx, ny) such that ``nx * ny <= maxncore``. + The function tries to find layouts with nx and ny as close as possible to + sqrt(maxncore). + + Parameters + ---------- + maxncore : int, required + Maximum number of cores to use. + + abs_maxdiff_nx_ny : int, optional + Maximum absolute difference between nx and ny in the layout. Default is 4. + + even_nx : bool, optional + If True, only layouts with even nx are returned. Default is False. + + prefer_nx_greater_than_ny : bool, optional + If True, only layouts with nx >= ny are returned. Default is False. + + Returns + ------- + list of tuples + List of (nx, ny) tuples representing the unique layouts found. + If no layouts are found, an empty list is returned. + + Raises + ------ + ValueError + If maxncore is not a positive integer or if abs_maxdiff_nx_ny is negative. + """ + import math + + if maxncore < 1: + raise ValueError(f"Max. number of cores to use must be a positive integer. Got {maxncore} instead") + if abs_maxdiff_nx_ny < 0: + raise ValueError( + "The max. absolute difference between nx and ny in the layout " + f" must be a non-negative integer. Got {abs_maxdiff_nx_ny} instead" + ) + + if maxncore < 2 and even_nx: + return [] + + best_ncore = int(math.sqrt(maxncore)) + layouts = [] + start = max(1, best_ncore - abs_maxdiff_nx_ny) + if prefer_nx_greater_than_ny: + start = best_ncore + + for nx in range(start, best_ncore + abs_maxdiff_nx_ny + 1): + if even_nx and nx % 2 != 0: + continue + ny = maxncore // nx + if abs(nx - ny) > abs_maxdiff_nx_ny: + continue + if prefer_nx_greater_than_ny and nx < ny: + continue + + layouts.append((nx, ny)) + + return layouts diff --git a/tests/test_esm1p6_layout_input.py b/tests/test_esm1p6_layout_input.py new file mode 100644 index 0000000..d509497 --- /dev/null +++ b/tests/test_esm1p6_layout_input.py @@ -0,0 +1,335 @@ +# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from access.config.esm1p6_layout_input import ( + _generate_esm1p6_layout_from_core_counts, + generate_esm1p6_core_layouts_from_node_count, + generate_esm1p6_perturb_block, + return_layout_tuple, +) + + +@pytest.fixture(scope="module") +def layout_tuple(): + return return_layout_tuple() + + +@pytest.fixture(scope="module") +def esm1p6_ctrl_layout(): + layout_tuple = return_layout_tuple() + return layout_tuple(ncores_used=416, atm_nx=16, atm_ny=13, mom_nx=14, mom_ny=14, ice_ncores=12) # Example layout + + +def test_generate_esm1p6_layout_from_core_counts(layout_tuple): + # Test the the validation works + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=120, + max_atm_ncores=96, + ice_ncores=6, + ncores_for_atm_and_ocn=208 - 6, + min_ncores_needed=0, + ) + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=1, + max_atm_ncores=96, + ice_ncores=6, + ncores_for_atm_and_ocn=208 - 6, + min_ncores_needed=0, + ) + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=96, + max_atm_ncores=120, + atm_ncore_delta=0, + ice_ncores=6, + ncores_for_atm_and_ocn=208 - 6, + min_ncores_needed=0, + ) + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=96, + max_atm_ncores=120, + ice_ncores=0, + ncores_for_atm_and_ocn=208 - 6, + min_ncores_needed=0, + ) + + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=96, + max_atm_ncores=120, + ice_ncores=6, + ncores_for_atm_and_ocn=208 - 6, + min_ncores_needed=0, + mom_ncores_over_atm_ncores_range=(-1.0, 1.0), + ) + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=96, + max_atm_ncores=120, + ice_ncores=6, + ncores_for_atm_and_ocn=208 - 6, + min_ncores_needed=0, + mom_ncores_over_atm_ncores_range=(1.0, -1.0), + ) + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=96, + max_atm_ncores=120, + ice_ncores=6, + ncores_for_atm_and_ocn=208 - 6, + min_ncores_needed=0, + mom_ncores_over_atm_ncores_range=(1.1, 0.9), + ) + + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=96, + max_atm_ncores=120, + ice_ncores=6, + ncores_for_atm_and_ocn=208 - 2 * 6, + min_ncores_needed=208, + ) + + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + min_atm_ncores=2, + max_atm_ncores=2, + ice_ncores=6, + ncores_for_atm_and_ocn=2, + min_ncores_needed=1, + ) + + # Test with a valid core count + core_count = 208 + max_atm_ncores = 120 + min_atm_ncores = 96 + ice_ncores = 6 + ncores_for_atm_and_ocn = core_count - ice_ncores + min_ncores_needed = core_count - 1 # Allow for some unused cores + + layouts = _generate_esm1p6_layout_from_core_counts( + max_atm_ncores=max_atm_ncores, + min_atm_ncores=min_atm_ncores, + ice_ncores=ice_ncores, + ncores_for_atm_and_ocn=ncores_for_atm_and_ocn, + min_ncores_needed=min_ncores_needed, + ) + assert all(layout.ncores_used >= min_ncores_needed for layout in layouts), ( + f"Some layouts have ncores_used < min_ncores_needed. Min ncores needed: {min_ncores_needed}, " + f"Min ncores used: {min([x.ncores_used for x in layouts])}" + ) + assert all(layout.ncores_used <= (ncores_for_atm_and_ocn + layout.ice_ncores) for layout in layouts), ( + f"Some layouts have ncores_used > ncores_for_atm_and_ocn + ice_ncores. Max ncores for " + f"atm and ocn: {ncores_for_atm_and_ocn}, Max ncores used: {max([x.ncores_used for x in layouts])}" + ) + + # Test that setting min_ncores_needed less than ncores_for_atm_and_ocn produces larger number of layouts + layouts_without_min_ncores = _generate_esm1p6_layout_from_core_counts( + max_atm_ncores=max_atm_ncores, + min_atm_ncores=min_atm_ncores, + ice_ncores=ice_ncores, + ncores_for_atm_and_ocn=ncores_for_atm_and_ocn, + min_ncores_needed=ncores_for_atm_and_ocn - 1, + ) + + assert len(layouts_without_min_ncores) >= len(layouts), ( + f"Expected more layouts when min_ncores_needed is less than " + f"ncores_for_atm_and_ocn. Got {len(layouts_without_min_ncores)} vs {len(layouts)}" + ) + assert all(x in layouts_without_min_ncores for x in layouts), ( + "All layouts from the first call should be in the second call" + ) + + # Test that the continue statement in the loop works by setting abs_maxdiff_nx_ny to 0 + min_atm_ncores = 98 + max_atm_ncores = 102 + ice_ncores = 6 + ncores_for_atm_and_ocn = 10 * 10 + (10 * 10 - 1) + min_ncores_needed = 1 + abs_maxdiff_nx_ny = 0 + layouts = _generate_esm1p6_layout_from_core_counts( + max_atm_ncores=max_atm_ncores, + min_atm_ncores=min_atm_ncores, + ice_ncores=ice_ncores, + ncores_for_atm_and_ocn=ncores_for_atm_and_ocn, + min_ncores_needed=min_ncores_needed, + abs_maxdiff_nx_ny=abs_maxdiff_nx_ny, + ) + assert layouts == [], f"Expected *no* layouts to be returned. Got layouts = {layouts}" + + # Test with zero cores + core_count = 0 + with pytest.raises(ValueError): + layouts = _generate_esm1p6_layout_from_core_counts( + max_atm_ncores=max_atm_ncores, + min_atm_ncores=min_atm_ncores, + ice_ncores=ice_ncores, + ncores_for_atm_and_ocn=0, + min_ncores_needed=ice_ncores, + ) + + # Test that the layouts are returned with ncores_used <= ncores_for_atm_and_ocn + assert all(x.ncores_used <= ncores_for_atm_and_ocn for x in layouts), ( + f"Some layouts have ncores_used > ncores_for_atm_and_ocn. " + f"Max. ncores used : {max([x.ncores_used for x in layouts])}" + ) + + # Test that the cores_used are sorted in descending order + assert all(layouts[i].ncores_used >= layouts[i + 1].ncores_used for i in range(len(layouts) - 1)), ( + "Layouts are not sorted in descending order of ncores_used" + ) + + +def test_generate_esm1p6_core_layouts_from_node_count(esm1p6_ctrl_layout): + # Test the the validation works + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, tol_around_ctrl_ratio=-0.1) + + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count( + 4, tol_around_ctrl_ratio=None, mom_ncores_over_atm_ncores_range=None + ) + + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, tol_around_ctrl_ratio=-0.1) + + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=(None,)) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=None) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=(None, None)) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=(1.0, None)) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=(None, 1.0)) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=(1.2, 0.8)) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=(0.8, -1.0)) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, mom_ncores_over_atm_ncores_range=(-0.8, 1.0)) + + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, atm_ncore_delta=0) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, atm_ncore_delta=-1) + + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, abs_maxdiff_nx_ny=-1) + + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, max_wasted_ncores_frac=-0.1) + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(4, max_wasted_ncores_frac=1.01) + + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count([4, "abcd"]) + + # Test with negative nodes + node_count = -3 + with pytest.raises(ValueError): + generate_esm1p6_core_layouts_from_node_count(node_count) + + # Test that with a very low node count, no layouts are returned (i.e. empty list of an empty list) + layouts = generate_esm1p6_core_layouts_from_node_count([0.2], max_wasted_ncores_frac=0.2) + assert layouts != [[]], f"Expected layouts to be returned even with small node fraction. Got layouts = {layouts}" + + layouts = generate_esm1p6_core_layouts_from_node_count([0.001], max_wasted_ncores_frac=0.5) + assert layouts == [[]], f"Expected no layouts to be returned for nearly zero nodes. Got layouts = {layouts}" + + # Test with a valid node count that should return the control layout + node_count = 4 + layouts = generate_esm1p6_core_layouts_from_node_count(node_count, tol_around_ctrl_ratio=0.0)[0] + assert len(layouts) == 1, f"Expected *exactly* one layout to be returned. Got layouts = {layouts}" + layouts = layouts[0] + assert esm1p6_ctrl_layout == layouts, f"Control config layout={esm1p6_ctrl_layout} not found in solved {layouts}" + + # Test with a valid node count as a float that should return the control layout + node_count = 4.0 + layouts = generate_esm1p6_core_layouts_from_node_count(node_count, tol_around_ctrl_ratio=0.0)[0] + assert len(layouts) == 1, f"Expected *exactly* one layout to be returned. Got layouts = {layouts}" + layouts = layouts[0] + assert esm1p6_ctrl_layout == layouts, f"Control config layout={esm1p6_ctrl_layout} not found in solved {layouts}" + + # Test with zero nodes + node_count = 0 + with pytest.raises(ValueError): + layouts = generate_esm1p6_core_layouts_from_node_count(node_count) + + # Test with non-integer nodes + node_count = 2.5 + layouts = generate_esm1p6_core_layouts_from_node_count(node_count) + assert layouts != [[]], f"Expected layouts to be returned for non-integer nodes. Got layouts = {layouts}" + + # Test that specifying mom_ncores_over_atm_ncores_range works + node_count = 4 + mom_ncores_over_atm_ncores_range = (0.8, 1.2) + layouts = generate_esm1p6_core_layouts_from_node_count( + node_count, mom_ncores_over_atm_ncores_range=mom_ncores_over_atm_ncores_range + ) + assert layouts != [[]], f"Expected layouts to be returned for non-integer nodes. Got layouts = {layouts}" + + # Test that allocating remaining cores to ICE works + from access.config.layout_config import convert_num_nodes_to_ncores + + node_count, queue = 4, "normalsr" + totncores = convert_num_nodes_to_ncores(node_count, queue=queue) + mom_ncores_over_atm_ncores_range = (0.8, 1.2) + layouts = generate_esm1p6_core_layouts_from_node_count( + node_count, + mom_ncores_over_atm_ncores_range=mom_ncores_over_atm_ncores_range, + allocate_unused_cores_to_ice=True, + queue=queue, + ) + assert layouts != [[]], f"Expected layouts to be returned for non-integer nodes. Got layouts = {layouts}" + assert all(layout.ice_ncores >= esm1p6_ctrl_layout.ice_ncores for layout in layouts[0]), ( + f"Expected ice_ncores to be >= {esm1p6_ctrl_layout.ice_ncores}. Got layout = {layouts[0]}" + ) + assert all(layout.ncores_used == totncores for layout in layouts[0]), ( + f"Expected ncores used to be *exactly* equal to {totncores}. Got layout = {layouts[0]}" + ) + + +def test_generate_esm1p6_perturb_block(esm1p6_ctrl_layout): + # Test that the validation works + with pytest.raises(ValueError): + generate_esm1p6_perturb_block(num_nodes=None, layouts=esm1p6_ctrl_layout, branch_name_prefix="test_block") + with pytest.raises(ValueError): + generate_esm1p6_perturb_block(num_nodes=-1, layouts=esm1p6_ctrl_layout, branch_name_prefix="test_block") + + with pytest.raises(ValueError): + generate_esm1p6_perturb_block(num_nodes=4, layouts=esm1p6_ctrl_layout, branch_name_prefix=None) + + # Test with invalid layout + with pytest.raises(ValueError): + generate_esm1p6_perturb_block(num_nodes=4, layouts=None, branch_name_prefix="test_block") + + # Test with empty layout + with pytest.raises(ValueError): + generate_esm1p6_perturb_block(num_nodes=4, layouts=[[]], branch_name_prefix="test_block") + + # Test that the validation works for layouts with missing fields + with pytest.raises(ValueError): + missing_ice_ncores_layout = [[416, 16, 13, 14, 14]] + generate_esm1p6_perturb_block(num_nodes=4, layouts=missing_ice_ncores_layout, branch_name_prefix="test_block") + + with pytest.raises(ValueError): + generate_esm1p6_perturb_block( + num_nodes=4, layouts=esm1p6_ctrl_layout, branch_name_prefix="test_block", start_blocknum=-1 + ) + + # Test with valid parameters + branch_name_prefix = "test_block" + perturb_block, _ = generate_esm1p6_perturb_block( + num_nodes=4, layouts=esm1p6_ctrl_layout, branch_name_prefix=branch_name_prefix + ) + assert isinstance(perturb_block, str), f"Expected perturb block to be a string, but got: {type(perturb_block)}" + assert branch_name_prefix in perturb_block, ( + f"Expected branch name prefix '{branch_name_prefix}' to be in perturb block, but got: {perturb_block}" + ) diff --git a/tests/test_layout_config.py b/tests/test_layout_config.py new file mode 100644 index 0000000..52c6c9f --- /dev/null +++ b/tests/test_layout_config.py @@ -0,0 +1,125 @@ +# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from access.config.layout_config import convert_num_nodes_to_ncores, find_layouts_with_maxncore, return_layout_tuple + + +@pytest.fixture(scope="module") +def layout_tuple(): + return return_layout_tuple() + + +def test_layout_tuple(layout_tuple): + # Test that the tuple has the correct field names and types + atm_nx, atm_ny = 6, 4 + mom_nx, mom_ny = 5, 4 + ice_ncores = 2 + ncores_used = atm_nx * atm_ny * mom_nx * mom_ny * ice_ncores + layout = layout_tuple(ncores_used, atm_nx, atm_ny, mom_nx, mom_ny, ice_ncores) + assert isinstance(layout.ncores_used, int), f"Expected int, got {type(layout.ncores_used)}" + assert isinstance(layout.atm_nx, int), f"Expected int, got {type(layout.atm_nx)}" + assert isinstance(layout.atm_ny, int), f"Expected int, got {type(layout.atm_ny)}" + assert isinstance(layout.mom_nx, int), f"Expected int, got {type(layout.mom_nx)}" + assert isinstance(layout.mom_ny, int), f"Expected int, got {type(layout.mom_ny)}" + assert isinstance(layout.ice_ncores, int), f"Expected int, got {type(layout.ice_ncores)}" + + assert layout.ncores_used == ncores_used, f"Expected {ncores_used}, got {layout.ncores_used}" + assert layout.atm_nx == atm_nx, f"Expected {atm_nx}, got {layout.atm_nx}" + assert layout.atm_ny == atm_ny, f"Expected {atm_ny}, got {layout.atm_ny}" + assert layout.mom_nx == mom_nx, f"Expected {mom_nx}, got {layout.mom_nx}" + assert layout.mom_ny == mom_ny, f"Expected {mom_ny}, got {layout.mom_ny}" + assert layout.ice_ncores == ice_ncores, f"Expected {ice_ncores}, got {layout.ice_ncores}" + + +def test_find_layouts_with_maxncore(): + maxncores = 20 + layouts = find_layouts_with_maxncore(maxncores) + assert isinstance(layouts, list), f"Expected list, got {type(layouts)}" + assert all(isinstance(layout, tuple) for layout in layouts), "All items in the list should be tuples" + assert all(len(layout) == 2 for layout in layouts), "All tuples should have length 2" + assert all(isinstance(n, int) for layout in layouts for n in layout), ( + "All elements in the tuples should be integers" + ) + assert all(layout[0] * layout[1] <= maxncores for layout in layouts), ( + f"All layouts should have nx * ny <= {maxncores}" + ) + assert (5, 4) in layouts, f"(5, 4) should be in the layouts for maxncores={maxncores}" + assert (4, 5) in layouts, f"(4, 5) should be in the layouts for maxncores={maxncores}" + assert (4, 4) not in layouts, f"(4, 4) should *not* be in the layouts for maxncores={maxncores}" + assert (4, 1) not in layouts, f"(4, 1) should *not* be in the layouts for maxncores={maxncores}" + + layouts_even_nx = find_layouts_with_maxncore(maxncores, even_nx=True) + assert (5, 4) not in layouts_even_nx, ( + "(5, 4) should *not* be in the layouts for maxncores={maxncores} with even_nx=True" + ) + assert all(layout[0] % 2 == 0 for layout in layouts_even_nx), "All nx should be even when even_nx=True" + + layouts_nx_ge_ny = find_layouts_with_maxncore(maxncores, prefer_nx_greater_than_ny=True) + assert all(layout[0] >= layout[1] for layout in layouts_nx_ge_ny), ( + "All layouts should have nx >= ny when prefer_nx_greater_than_ny=True" + ) + + layouts_both = find_layouts_with_maxncore(maxncores, even_nx=True, prefer_nx_greater_than_ny=True) + assert all(layout[0] % 2 == 0 and layout[0] >= layout[1] for layout in layouts_both), ( + "All layouts should have even nx and nx >= ny when both options are set" + ) + + layout_none = find_layouts_with_maxncore(1, even_nx=True) + assert layout_none == [], "No layouts should be found for maxncores=1 with even_nx=True" + + layout_abs_maxdiff = find_layouts_with_maxncore(maxncores, abs_maxdiff_nx_ny=1) + assert all(abs(layout[0] - layout[1]) <= 1 for layout in layout_abs_maxdiff), ( + "All layouts should have abs(nx - ny) <= 1 when abs_maxdiff_nx_ny=1" + ) + + layout_abs_maxdiff = find_layouts_with_maxncore(16, abs_maxdiff_nx_ny=0) + assert len(layout_abs_maxdiff) == 1 and layout_abs_maxdiff[0] == (4, 4), ( + "Only (4, 4) should be found for maxncores=16 with abs_maxdiff_nx_ny=0" + ) + + with pytest.raises(ValueError): + find_layouts_with_maxncore(0) + with pytest.raises(ValueError): + find_layouts_with_maxncore(-4) + with pytest.raises(ValueError): + find_layouts_with_maxncore(16, abs_maxdiff_nx_ny=-1) + with pytest.raises(ValueError): + find_layouts_with_maxncore(-16, even_nx=True) + with pytest.raises(ValueError): + find_layouts_with_maxncore(-16, prefer_nx_greater_than_ny=True) + with pytest.raises(ValueError): + find_layouts_with_maxncore(-16, even_nx=True, prefer_nx_greater_than_ny=True) + + +def test_convert_num_nodes_to_ncores(): + with pytest.raises(ValueError): + convert_num_nodes_to_ncores(2.5, queue="broadwell") + with pytest.raises(ValueError): + convert_num_nodes_to_ncores(2.5, queue="unknown_queue") + + assert convert_num_nodes_to_ncores(2, queue="normalsr") == 208, ( + f"Expected 208, got {convert_num_nodes_to_ncores(2, queue='normalsr')}" + ) + assert convert_num_nodes_to_ncores(2.0, queue="normalsr") == 208, ( + f"Expected 208, got {convert_num_nodes_to_ncores(2.0, queue='normalsr')}" + ) + assert convert_num_nodes_to_ncores(1, queue="normalsr") == 104, ( + f"Expected 104, got {convert_num_nodes_to_ncores(1, queue='normalsr')}" + ) + assert convert_num_nodes_to_ncores(1.0, queue="normalsr") == 104, ( + f"Expected 104, got {convert_num_nodes_to_ncores(1.0, queue='normalsr')}" + ) + assert convert_num_nodes_to_ncores(0.5) == 52, f"Expected 52, got {convert_num_nodes_to_ncores(0.5)}" + assert convert_num_nodes_to_ncores(0.5, queue="normal") == 24, ( + f"Expected 24, got {convert_num_nodes_to_ncores(0.5, queue='normal')}" + ) + assert convert_num_nodes_to_ncores(3, queue="normal") == 144, ( + f"Expected 144, got {convert_num_nodes_to_ncores(3, queue='normal')}" + ) + assert convert_num_nodes_to_ncores(3.0, queue="normal") == 144, ( + f"Expected 144, got {convert_num_nodes_to_ncores(3.0, queue='normal')}" + ) + with pytest.raises(ValueError): + convert_num_nodes_to_ncores([2])