Skip to content

Fix mypy type errors in mpi_resources.py #1759

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jlnav wants to merge 1 commit into
base: develop
Choose a base branch
from
+47 −37
Open

Fix mypy type errors in mpi_resources.py #1759

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 47 additions & 37 deletions libensemble/resources/mpi_resources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,13 @@ def rassert(test: int | bool | None, *args) -> None:
# logger.setLevel(logging.DEBUG)


def get_MPI_variant() -> str:
def get_MPI_variant() -> str | None:
"""Returns MPI base implementation

Returns
-------
mpi_variant: str
MPI variant 'aprun' or 'jsrun' or 'msmpi' or 'mpich' or 'openmpi' or 'srun'
mpi_variant: str | None
MPI variant 'aprun' or 'jsrun' or 'msmpi' or 'mpich' or 'openmpi' or 'srun', or None if not found

"""

Expand Down Expand Up @@ -85,7 +85,7 @@ def get_MPI_variant() -> str:
return None


def get_MPI_runner(mpi_runner=None) -> str:
def get_MPI_runner(mpi_runner: str | None = None) -> str | None:
"""Return whether ``mpirun`` is openmpi or mpich"""
var = mpi_runner or get_MPI_variant()
if var in ["mpich", "openmpi"]:
Expand All @@ -102,29 +102,31 @@ def task_partition(
"""

# Convert to int if string is provided
num_procs = int(num_procs) if num_procs else None
num_nodes = int(num_nodes) if num_nodes else None
procs_per_node = int(procs_per_node) if procs_per_node else None
num_procs_int: int | None = int(num_procs) if num_procs else None
num_nodes_int: int | None = int(num_nodes) if num_nodes else None
procs_per_node_int: int | None = int(procs_per_node) if procs_per_node else None

# If machinefile is provided - ignore everything else
if machinefile:
if num_procs or num_nodes or procs_per_node:
if num_procs_int or num_nodes_int or procs_per_node_int:
logger.warning("Machinefile provided - overriding " "procs/nodes/procs_per_node")
return None, None, None

if not num_procs:
rassert(num_nodes and procs_per_node, "Need num_procs, num_nodes/procs_per_node, or machinefile")
num_procs = num_nodes * procs_per_node
if not num_procs_int:
rassert(num_nodes_int and procs_per_node_int, "Need num_procs, num_nodes/procs_per_node, or machinefile")
assert num_nodes_int is not None and procs_per_node_int is not None # for mypy
num_procs_int = num_nodes_int * procs_per_node_int

elif not num_nodes:
procs_per_node = procs_per_node or num_procs
num_nodes = num_procs // procs_per_node
elif not num_nodes_int:
procs_per_node_int = procs_per_node_int or num_procs_int
num_nodes_int = num_procs_int // procs_per_node_int

elif not procs_per_node:
procs_per_node = num_procs // num_nodes
elif not procs_per_node_int:
procs_per_node_int = num_procs_int // num_nodes_int

rassert(num_procs == num_nodes * procs_per_node, "num_procs does not equal num_nodes*procs_per_node")
return num_procs, num_nodes, procs_per_node
assert num_nodes_int is not None and procs_per_node_int is not None # for mypy
rassert(num_procs_int == num_nodes_int * procs_per_node_int, "num_procs does not equal num_nodes*procs_per_node")
return num_procs_int, num_nodes_int, procs_per_node_int


def _max_rsets_per_node(worker_resources: WorkerResources) -> int:
Expand All @@ -137,9 +139,9 @@ def _max_rsets_per_node(worker_resources: WorkerResources) -> int:

def get_resources(
resources: Resources,
num_procs: int = None,
num_nodes: int = None,
procs_per_node: int = None,
num_procs: int | None = None,
num_nodes: int | None = None,
procs_per_node: int | None = None,
hyperthreads: bool = False,
) -> tuple[int, int, int]:
"""Reconciles user-supplied options with available worker
Expand All @@ -153,6 +155,7 @@ def get_resources(
raised if these are infeasible.
"""
wresources = resources.worker_resources
assert wresources is not None # for mypy
gresources = resources.glob_resources
node_list = wresources.local_nodelist
rassert(node_list, "Node list is empty - aborting")
Expand Down Expand Up @@ -190,57 +193,62 @@ def get_resources(
f"Nodes: {num_nodes} procs_per_node {procs_per_node}"
)
elif not num_nodes and not procs_per_node:
if num_procs <= cores_avail_per_node_per_worker:
if num_procs is not None and num_procs <= cores_avail_per_node_per_worker:
num_nodes = 1
else:
num_nodes = local_node_count
elif not num_procs and not num_nodes:
num_nodes = local_node_count

# Checks config is consistent and sufficient to express
num_procs, num_nodes, procs_per_node = task_partition(num_procs, num_nodes, procs_per_node)
result = task_partition(num_procs, num_nodes, procs_per_node)
if result == (None, None, None):
raise MPIResourcesException("task_partition returned all None unexpectedly")
num_procs_out, num_nodes_out, procs_per_node_out = result
# Type narrowing for mypy
assert num_procs_out is not None and num_nodes_out is not None and procs_per_node_out is not None

rassert(
num_nodes <= local_node_count,
"Not enough nodes to honor arguments. " f"Requested {num_nodes}. Only {local_node_count} available",
num_nodes_out <= local_node_count,
"Not enough nodes to honor arguments. " f"Requested {num_nodes_out}. Only {local_node_count} available",
)

if gresources.enforce_worker_core_bounds:
rassert(
procs_per_node <= cores_avail_per_node,
procs_per_node_out <= cores_avail_per_node,
"Not enough processors on a node to honor arguments. "
f"Requested {procs_per_node}. Only {cores_avail_per_node} available",
f"Requested {procs_per_node_out}. Only {cores_avail_per_node} available",
)

rassert(
procs_per_node <= cores_avail_per_node_per_worker,
procs_per_node_out <= cores_avail_per_node_per_worker,
"Not enough processors per worker to honor arguments. "
f"Requested {procs_per_node}. Only {cores_avail_per_node_per_worker} available",
f"Requested {procs_per_node_out}. Only {cores_avail_per_node_per_worker} available",
)

rassert(
num_procs <= (cores_avail_per_node * local_node_count),
num_procs_out <= (cores_avail_per_node * local_node_count),
"Not enough procs to honor arguments. "
f"Requested {num_procs}. Only {cores_avail_per_node * local_node_count} available",
f"Requested {num_procs_out}. Only {cores_avail_per_node * local_node_count} available",
)

if num_nodes < local_node_count:
if num_nodes_out < local_node_count:
logger.debug(
"User constraints mean fewer nodes being used "
f"than available. {num_nodes} nodes used. {local_node_count} nodes available"
f"than available. {num_nodes_out} nodes used. {local_node_count} nodes available"
)

return num_procs, num_nodes, procs_per_node
return num_procs_out, num_nodes_out, procs_per_node_out


def create_machinefile(
resources: Resources,
machinefile: str | None = None,
num_procs: int = None,
num_procs: int | None = None,
num_nodes: int | None = None,
procs_per_node: int | None = None,
hyperthreads: bool = False,
) -> tuple[bool, None, int, int]:
) -> tuple[bool, int | None, int | None, int | None]:
"""Creates a machinefile based on user-supplied config options,
completed by detected machine resources
"""
Expand All @@ -252,12 +260,13 @@ def create_machinefile(
except Exception as e:
logger.warning(f"Could not remove existing machinefile: {e}")

assert resources.worker_resources is not None # for mypy
node_list = resources.worker_resources.local_nodelist
logger.debug(f"Creating machinefile with {num_nodes} nodes and {procs_per_node} ranks per node")

with open(machinefile, "w") as f:
for node in node_list[:num_nodes]:
f.write((node + "\n") * procs_per_node)
f.write((node + "\n") * (procs_per_node or 1))

built_mfile = os.path.isfile(machinefile) and os.path.getsize(machinefile) > 0
return built_mfile, num_procs, num_nodes, procs_per_node
Expand All @@ -268,6 +277,7 @@ def get_hostlist(resources: Resources, num_nodes=None):

completed by detected machine resources
"""
assert resources.worker_resources is not None # for mypy
node_list = resources.worker_resources.local_nodelist
hostlist_str = ",".join([str(x) for x in node_list[:num_nodes]])
return hostlist_str
Loading