Skip to content

Commit

Permalink
addressing review comments; added value checks for batch and launch a…
Browse files Browse the repository at this point in the history
…rguments
  • Loading branch information
juliaputko committed Oct 15, 2024
1 parent 6aa9991 commit 67ba3a5
Show file tree
Hide file tree
Showing 7 changed files with 168 additions and 34 deletions.
25 changes: 22 additions & 3 deletions smartsim/settings/arguments/batch/lsf.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,10 @@ def set_smts(self, smts: int) -> None:
takes precedence.
:param smts: SMT (e.g on Summit: 1, 2, or 4)
:raises TypeError: if not an int
"""
if not isinstance(smts, int):
raise TypeError("smts argument was not of type int")
self.set("alloc_flags", str(smts))

def set_project(self, project: str) -> None:
Expand All @@ -81,7 +84,10 @@ def set_project(self, project: str) -> None:
This sets ``-P``.
:param time: project name
:raises TypeError: if not a str
"""
if not isinstance(project, str):
raise TypeError("project argument was not of type str")
self.set("P", project)

def set_account(self, account: str) -> None:
Expand All @@ -90,7 +96,10 @@ def set_account(self, account: str) -> None:
this function is an alias for `set_project`.
:param account: project name
:raises TypeError: if not a str
"""
if not isinstance(account, str):
raise TypeError("account argument was not of type str")
return self.set_project(account)

def set_nodes(self, num_nodes: int) -> None:
Expand All @@ -99,7 +108,10 @@ def set_nodes(self, num_nodes: int) -> None:
This sets ``-nnodes``.
:param nodes: number of nodes
:raises TypeError: if not an int
"""
if not isinstance(num_nodes, int):
raise TypeError("num_nodes argument was not of type int")
self.set("nnodes", str(num_nodes))

def set_hostlist(self, host_list: t.Union[str, t.List[str]]) -> None:
Expand All @@ -110,10 +122,11 @@ def set_hostlist(self, host_list: t.Union[str, t.List[str]]) -> None:
"""
if isinstance(host_list, str):
host_list = [host_list.strip()]
if not isinstance(host_list, list):
if not (
isinstance(host_list, list)
and all(isinstance(item, str) for item in host_list)
):
raise TypeError("host_list argument must be a list of strings")
if not all(isinstance(host, str) for host in host_list):
raise TypeError("host_list argument must be list of strings")
self.set("m", '"' + " ".join(host_list) + '"')

def set_tasks(self, tasks: int) -> None:
Expand All @@ -122,7 +135,10 @@ def set_tasks(self, tasks: int) -> None:
This sets ``-n``
:param tasks: number of tasks
:raises TypeError: if not an int
"""
if not isinstance(tasks, int):
raise TypeError("tasks argument was not of type int")
self.set("n", str(tasks))

def set_queue(self, queue: str) -> None:
Expand All @@ -131,7 +147,10 @@ def set_queue(self, queue: str) -> None:
This sets ``-q``
:param queue: The queue to submit the job on
:raises TypeError: if not a str
"""
if not isinstance(queue, str):
raise TypeError("queue argument was not of type str")
self.set("q", queue)

def format_batch_args(self) -> t.List[str]:
Expand Down
27 changes: 23 additions & 4 deletions smartsim/settings/arguments/batch/pbs.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@

from __future__ import annotations

import re
import typing as t
from copy import deepcopy

Expand Down Expand Up @@ -61,7 +62,10 @@ def set_nodes(self, num_nodes: int) -> None:
nodes here is sets the 'nodes' resource.
:param num_nodes: number of nodes
:raises TypeError: if not an int
"""
if not isinstance(num_nodes, int):
raise TypeError("num_nodes argument was not of type int")

self.set("nodes", str(num_nodes))

Expand All @@ -73,9 +77,10 @@ def set_hostlist(self, host_list: t.Union[str, t.List[str]]) -> None:
"""
if isinstance(host_list, str):
host_list = [host_list.strip()]
if not isinstance(host_list, list):
raise TypeError("host_list argument must be a list of strings")
if not all(isinstance(host, str) for host in host_list):
if not (
isinstance(host_list, list)
and all(isinstance(item, str) for item in host_list)
):
raise TypeError("host_list argument must be a list of strings")
self.set("hostname", ",".join(host_list))

Expand All @@ -89,14 +94,22 @@ def set_walltime(self, walltime: str) -> None:
this value will be overridden
:param walltime: wall time
:raises ValueError: if walltime format is invalid
"""
self.set("walltime", walltime)
pattern = r"^\d{2}:\d{2}:\d{2}$"
if walltime and re.match(pattern, walltime):
self.set("walltime", walltime)
else:
raise ValueError("Invalid walltime format. Please use 'HH:MM:SS' format.")

def set_queue(self, queue: str) -> None:
"""Set the queue for the batch job
:param queue: queue name
:raises TypeError: if not a str
"""
if not isinstance(queue, str):
raise TypeError("queue argument was not of type str")
self.set("q", str(queue))

def set_ncpus(self, num_cpus: int) -> None:
Expand All @@ -107,14 +120,20 @@ def set_ncpus(self, num_cpus: int) -> None:
this value will be overridden
:param num_cpus: number of cpus per node in select
:raises TypeError: if not an int
"""
if not isinstance(num_cpus, int):
raise TypeError("num_cpus argument was not of type int")
self.set("ppn", str(num_cpus))

def set_account(self, account: str) -> None:
"""Set the account for this batch job
:param acct: account id
:raises TypeError: if not a str
"""
if not isinstance(account, str):
raise TypeError("account argument was not of type str")
self.set("A", str(account))

def format_batch_args(self) -> t.List[str]:
Expand Down
19 changes: 16 additions & 3 deletions smartsim/settings/arguments/batch/slurm.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def set_walltime(self, walltime: str) -> None:
format = "HH:MM:SS"
:param walltime: wall time
:raises ValueError: if walltime format is invalid
"""
pattern = r"^\d{2}:\d{2}:\d{2}$"
if walltime and re.match(pattern, walltime):
Expand All @@ -69,7 +70,10 @@ def set_nodes(self, num_nodes: int) -> None:
This sets ``--nodes``.
:param num_nodes: number of nodes
:raises TypeError: if not an int
"""
if not isinstance(num_nodes, int):
raise TypeError("num_nodes argument was not of type int")
self.set("nodes", str(num_nodes))

def set_account(self, account: str) -> None:
Expand All @@ -78,7 +82,10 @@ def set_account(self, account: str) -> None:
This sets ``--account``.
:param account: account id
:raises TypeError: if not a str
"""
if not isinstance(account, str):
raise TypeError("account argument was not of type str")
self.set("account", account)

def set_partition(self, partition: str) -> None:
Expand All @@ -96,7 +103,10 @@ def set_queue(self, queue: str) -> None:
Sets the partition for the slurm batch job
:param queue: the partition to run the batch job on
:raises TypeError: if not a str
"""
if not isinstance(queue, str):
raise TypeError("queue argument was not of type str")
return self.set_partition(queue)

def set_cpus_per_task(self, cpus_per_task: int) -> None:
Expand All @@ -118,10 +128,13 @@ def set_hostlist(self, host_list: t.Union[str, t.List[str]]) -> None:
"""
if isinstance(host_list, str):
host_list = [host_list.strip()]
if not isinstance(host_list, list):

if not (
isinstance(host_list, list)
and all(isinstance(item, str) for item in host_list)
):
raise TypeError("host_list argument must be a list of strings")
if not all(isinstance(host, str) for host in host_list):
raise TypeError("host_list argument must be list of strings")

self.set("nodelist", ",".join(host_list))

def format_batch_args(self) -> t.List[str]:
Expand Down
9 changes: 3 additions & 6 deletions smartsim/settings/batch_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,9 +118,9 @@ def __init__(
except ValueError:
raise ValueError(f"Invalid scheduler type: {batch_scheduler}") from None

if batch_args:
if batch_args is not None:
if not (
isinstance(batch_args, t.Mapping)
isinstance(batch_args, dict)
and all(isinstance(key, str) for key, val in batch_args.items())
):
raise TypeError(
Expand Down Expand Up @@ -152,10 +152,7 @@ def env_vars(self, value: t.Dict[str, str | None]) -> None:

if not (
isinstance(value, t.Mapping)
and all(
isinstance(key, str) and isinstance(val, str)
for key, val in value.items()
)
and all(isinstance(key, str) for key, val in value.items())
):
raise TypeError("env_vars argument was not of type dic of str and str")

Expand Down
17 changes: 3 additions & 14 deletions smartsim/settings/launch_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ def __init__(
except ValueError:
raise ValueError(f"Invalid launcher type: {launcher}")

if launch_args:
if launch_args is not None:
if not (
isinstance(launch_args, t.Mapping)
and all(isinstance(key, str) for key, val in launch_args.items())
Expand Down Expand Up @@ -175,11 +175,8 @@ def env_vars(self, value: dict[str, str | None]) -> None:
:param value: The new environment mapping
"""
if not (
isinstance(value, t.Mapping)
and all(
isinstance(key, str) and isinstance(val, str)
for key, val in value.items()
)
isinstance(value, dict)
and all(isinstance(key, str) for key, val in value.items())
):
raise TypeError("env_vars argument was not of type dic of str and str")

Expand Down Expand Up @@ -227,14 +224,6 @@ def update_env(self, env_vars: t.Dict[str, str | None]) -> None:
:param env_vars: environment variables to update or add
:raises TypeError: if env_vars values cannot be coerced to strings
"""
if not (
isinstance(env_vars, t.Mapping)
and all(
isinstance(key, str) and isinstance(val, str)
for key, val in env_vars.items()
)
):
raise TypeError("env_vars argument was not of type dic of str and str")

# Coerce env_vars values to str as a convenience to user
for env, val in env_vars.items():
Expand Down
93 changes: 91 additions & 2 deletions tests/temp_tests/test_settings/test_batchSettings.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,16 @@ def test_type_batch_scheduler():
)


def test_type_batch_args():
batch_args = "invalid"
@pytest.mark.parametrize(
"batch_args",
[
pytest.param("invalid", id="invalid"),
pytest.param("", id="empty string"),
pytest.param(0, id="0"),
pytest.param([], id="empty list"),
],
)
def test_type_batch_args(batch_args):
with pytest.raises(
TypeError, match="batch_args argument was not of type mapping of str and str"
):
Expand All @@ -108,3 +116,84 @@ def test_type_env_vars():
TypeError, match="env_vars argument was not of type dic of str and str"
):
BatchSettings(batch_scheduler="slurm", env_vars=env_vars)


@pytest.mark.parametrize(
"scheduler",
[
pytest.param("slurm", id="slurm scheduler"),
pytest.param("lsf", id="bsub scheduler"),
pytest.param("pbs", id="qsub scheduler"),
],
)
def test_batch_arguments_type_set_nodes(scheduler):
bs = BatchSettings(batch_scheduler=scheduler, env_vars={"ENV": "VAR"})
with pytest.raises(TypeError, match="num_nodes argument was not of type int"):
bs.batch_args.set_nodes("invalid")


@pytest.mark.parametrize(
"scheduler",
[
pytest.param("slurm", id="slurm scheduler"),
pytest.param("lsf", id="bsub scheduler"),
pytest.param("pbs", id="qsub scheduler"),
],
)
def test_batch_arguments_type_set_account(scheduler):
bs = BatchSettings(batch_scheduler=scheduler, env_vars={"ENV": "VAR"})

with pytest.raises(TypeError, match="account argument was not of type str"):
bs.batch_args.set_account(27)


@pytest.mark.parametrize(
"scheduler",
[
pytest.param("slurm", id="slurm scheduler"),
pytest.param("lsf", id="bsub scheduler"),
pytest.param("pbs", id="qsub scheduler"),
],
)
def test_batch_arguments_type_set_queue(scheduler):
bs = BatchSettings(batch_scheduler=scheduler, env_vars={"ENV": "VAR"})
with pytest.raises(TypeError, match="queue argument was not of type str"):
bs.batch_args.set_queue(27)


@pytest.mark.parametrize(
"scheduler",
[
pytest.param("slurm", id="slurm scheduler"),
pytest.param("lsf", id="bsub scheduler"),
pytest.param("pbs", id="qsub scheduler"),
],
)
def test_batch_arguments_type_set_hostlist(scheduler):
bs = BatchSettings(batch_scheduler=scheduler, env_vars={"ENV": "VAR"})
with pytest.raises(TypeError, match="host_list argument must be a list of strings"):
bs.batch_args.set_hostlist([25, 37])


def test_batch_arguments_type_set_ncpus():
bs = BatchSettings(batch_scheduler="pbs", env_vars={"ENV": "VAR"})
with pytest.raises(TypeError, match="num_cpus argument was not of type int"):
bs.batch_args.set_ncpus("invalid")


def test_batch_arguments_type_set_smts():
bs = BatchSettings(batch_scheduler="lsf", env_vars={"ENV": "VAR"})
with pytest.raises(TypeError, match="smts argument was not of type int"):
bs.batch_args.set_smts("invalid")


def test_batch_arguments_type_set_project():
bs = BatchSettings(batch_scheduler="lsf", env_vars={"ENV": "VAR"})
with pytest.raises(TypeError, match="project argument was not of type str"):
bs.batch_args.set_project(27)


def test_batch_arguments_type_set_tasks():
bs = BatchSettings(batch_scheduler="lsf", env_vars={"ENV": "VAR"})
with pytest.raises(TypeError, match="tasks argument was not of type int"):
bs.batch_args.set_tasks("invalid")
Loading

0 comments on commit 67ba3a5

Please sign in to comment.