diff --git a/smartsim/entity/application.py b/smartsim/entity/application.py index a8302fc1f..fb3ed2a7e 100644 --- a/smartsim/entity/application.py +++ b/smartsim/entity/application.py @@ -46,6 +46,14 @@ class Application(SmartSimEntity): + """The Application class enables users to execute computational tasks in an + Experiment workflow, such as launching compiled applications, running scripts, + or performing general computational operations. + + Applications are designed to be added to Jobs, where LaunchSettings are also + provided to inject launcher-specific behavior into the Job. + """ + def __init__( self, name: str, @@ -56,6 +64,16 @@ def __init__( ) -> None: """Initialize an ``Application`` + Applications require a name and an executable. Optionally, users may provide + executable arguments, files and file parameters. To create a simple Application + that echos `Hello World!`, consider the example below: + + .. highlight:: python + .. code-block:: python + + # Create an application that runs the 'echo' command + my_app = Application(name="my_app", exe="echo", exe_args="Hello World!") + :param name: name of the application :param exe: executable to run :param exe_args: executable arguments @@ -83,25 +101,25 @@ def __init__( @property def exe(self) -> str: - """Return executable to run. + """Return the executable. - :returns: application executable to run + :return: the executable """ return self._exe @exe.setter def exe(self, value: str) -> None: - """Set executable to run. + """Set the executable. - :param value: executable to run + :param value: the executable """ self._exe = copy.deepcopy(value) @property def exe_args(self) -> t.MutableSequence[str]: - """Return a list of attached executable arguments. + """Return the executable arguments. - :returns: application executable arguments + :return: the executable arguments """ return self._exe_args @@ -109,7 +127,7 @@ def exe_args(self) -> t.MutableSequence[str]: def exe_args(self, value: t.Union[str, t.Sequence[str], None]) -> None: """Set the executable arguments. - :param value: executable arguments + :param value: the executable arguments """ self._exe_args = self._build_exe_args(value) @@ -122,20 +140,20 @@ def add_exe_args(self, args: t.Union[str, t.List[str], None]) -> None: self._exe_args.extend(args) @property - def files(self) -> t.Optional[EntityFiles]: - """Return files to be copied, symlinked, and/or configured prior to - execution. + def files(self) -> t.Union[EntityFiles, None]: + """Return attached EntityFiles object. - :returns: files + :return: the EntityFiles object of files to be copied, symlinked, + and/or configured prior to execution """ return self._files @files.setter def files(self, value: t.Optional[EntityFiles]) -> None: - """Set files to be copied, symlinked, and/or configured prior to - execution. + """Set the EntityFiles object. - :param value: files + :param value: the EntityFiles object of files to be copied, symlinked, + and/or configured prior to execution """ self._files = copy.deepcopy(value) @@ -143,7 +161,7 @@ def files(self, value: t.Optional[EntityFiles]) -> None: def file_parameters(self) -> t.Mapping[str, str]: """Return file parameters. - :returns: application file parameters + :return: the file parameters """ return self._file_parameters @@ -151,7 +169,7 @@ def file_parameters(self) -> t.Mapping[str, str]: def file_parameters(self, value: t.Mapping[str, str]) -> None: """Set the file parameters. - :param value: file parameters + :param value: the file parameters """ self._file_parameters = copy.deepcopy(value) @@ -159,7 +177,7 @@ def file_parameters(self, value: t.Mapping[str, str]) -> None: def incoming_entities(self) -> t.List[SmartSimEntity]: """Return incoming entities. - :returns: incoming entities + :return: incoming entities """ return self._incoming_entities @@ -244,7 +262,7 @@ def attach_generator_files( def attached_files_table(self) -> str: """Return a list of attached files as a plain text table - :returns: String version of table + :return: String version of table """ if not self.files: return "No file attached to this application." diff --git a/smartsim/entity/ensemble.py b/smartsim/entity/ensemble.py index 261f22d65..191730df7 100644 --- a/smartsim/entity/ensemble.py +++ b/smartsim/entity/ensemble.py @@ -83,7 +83,7 @@ def __init__( copy.deepcopy(exe_arg_parameters) if exe_arg_parameters else {} ) """The parameters and values to be used when configuring entities""" - self._files = copy.deepcopy(files) if files else EntityFiles() + self._files = copy.deepcopy(files) if files else None """The files to be copied, symlinked, and/or configured prior to execution""" self._file_parameters = ( copy.deepcopy(file_parameters) if file_parameters else {} @@ -98,25 +98,25 @@ def __init__( @property def exe(self) -> str: - """Return executable to run. + """Return the attached executable. - :returns: application executable to run + :return: the executable """ return self._exe @exe.setter def exe(self, value: str | os.PathLike[str]) -> None: - """Set executable to run. + """Set the executable. - :param value: executable to run + :param value: the executable """ self._exe = os.fspath(value) @property def exe_args(self) -> t.List[str]: - """Return a list of attached executable arguments. + """Return attached list of executable arguments. - :returns: application executable arguments + :return: the executable arguments """ return self._exe_args @@ -124,15 +124,15 @@ def exe_args(self) -> t.List[str]: def exe_args(self, value: t.Sequence[str]) -> None: """Set the executable arguments. - :param value: executable arguments + :param value: the executable arguments """ self._exe_args = list(value) @property def exe_arg_parameters(self) -> t.Mapping[str, t.Sequence[t.Sequence[str]]]: - """Return the executable argument parameters + """Return attached executable argument parameters. - :returns: executable arguments parameters + :return: the executable argument parameters """ return self._exe_arg_parameters @@ -140,35 +140,35 @@ def exe_arg_parameters(self) -> t.Mapping[str, t.Sequence[t.Sequence[str]]]: def exe_arg_parameters( self, value: t.Mapping[str, t.Sequence[t.Sequence[str]]] ) -> None: - """Set the executable arguments. + """Set the executable argument parameters. - :param value: executable arguments + :param value: the executable argument parameters """ self._exe_arg_parameters = copy.deepcopy(value) @property - def files(self) -> EntityFiles: - """Return files to be copied, symlinked, and/or configured prior to - execution. + def files(self) -> t.Union[EntityFiles, None]: + """Return attached EntityFiles object. - :returns: files + :return: the EntityFiles object of files to be copied, symlinked, + and/or configured prior to execution """ return self._files @files.setter - def files(self, value: EntityFiles) -> None: - """Set files to be copied, symlinked, and/or configured prior to - execution. + def files(self, value: t.Optional[EntityFiles]) -> None: + """Set the EntityFiles object. - :param value: files + :param value: the EntityFiles object of files to be copied, symlinked, + and/or configured prior to execution """ self._files = copy.deepcopy(value) @property def file_parameters(self) -> t.Mapping[str, t.Sequence[str]]: - """Return file parameters. + """Return the attached file parameters. - :returns: application file parameters + :return: the file parameters """ return self._file_parameters @@ -176,7 +176,7 @@ def file_parameters(self) -> t.Mapping[str, t.Sequence[str]]: def file_parameters(self, value: t.Mapping[str, t.Sequence[str]]) -> None: """Set the file parameters. - :param value: file parameters + :param value: the file parameters """ self._file_parameters = dict(value) @@ -184,7 +184,7 @@ def file_parameters(self, value: t.Mapping[str, t.Sequence[str]]) -> None: def permutation_strategy(self) -> str | strategies.PermutationStrategyType: """Return the permutation strategy - :return: permutation strategy + :return: the permutation strategy """ return self._permutation_strategy @@ -194,7 +194,7 @@ def permutation_strategy( ) -> None: """Set the permutation strategy - :param value: permutation strategy + :param value: the permutation strategy """ self._permutation_strategy = value @@ -202,7 +202,7 @@ def permutation_strategy( def max_permutations(self) -> int: """Return the maximum permutations - :return: max permutations + :return: the max permutations """ return self._max_permutations @@ -210,29 +210,34 @@ def max_permutations(self) -> int: def max_permutations(self, value: int) -> None: """Set the maximum permutations - :param value: the maxpermutations + :param value: the max permutations """ self._max_permutations = value @property def replicas(self) -> int: - """Return the number of replicas + """Return the number of replicas. - :return: number of replicas + :return: the number of replicas """ return self._replicas @replicas.setter def replicas(self, value: int) -> None: - """Set the number of replicas + """Set the number of replicas. :return: the number of replicas """ self._replicas = value def _create_applications(self) -> tuple[Application, ...]: - """Concretize the ensemble attributes into a collection of - application instances. + """Generate a collection of Application instances based on the Ensembles attributes. + + This method uses a permutation strategy to create various combinations of file + parameters and executable arguments. Each combination is then replicated according + to the specified number of replicas, resulting in a set of Application instances. + + :return: A tuple of Application instances """ permutation_strategy = strategies.resolve(self.permutation_strategy) @@ -255,6 +260,35 @@ def _create_applications(self) -> tuple[Application, ...]: ) def as_jobs(self, settings: LaunchSettings) -> tuple[Job, ...]: + """Expand an Ensemble into a list of deployable Jobs and apply + identical LaunchSettings to each Job. + + The number of Jobs returned is controlled by the Ensemble attributes: + - Ensemble.exe_arg_parameters + - Ensemble.file_parameters + - Ensemble.permutation_strategy + - Ensemble.max_permutations + - Ensemble.replicas + + Consider the example below: + + .. highlight:: python + .. code-block:: python + + # Create LaunchSettings + my_launch_settings = LaunchSettings(...) + + # Initialize the Ensemble + ensemble = Ensemble("my_name", "echo", "hello world", replicas=3) + # Expand Ensemble into Jobs + ensemble_as_jobs = ensemble.as_jobs(my_launch_settings) + + By calling `as_jobs` on `ensemble`, three Jobs are returned because + three replicas were specified. Each Job will have the provided LaunchSettings. + + :param settings: LaunchSettings to apply to each Job + :return: Sequence of Jobs with the provided LaunchSettings + """ apps = self._create_applications() if not apps: raise ValueError("There are no members as part of this ensemble") diff --git a/smartsim/launchable/job.py b/smartsim/launchable/job.py index b7d81bfdc..6ec2bbbc4 100644 --- a/smartsim/launchable/job.py +++ b/smartsim/launchable/job.py @@ -26,7 +26,6 @@ from __future__ import annotations -import os import typing as t from copy import deepcopy @@ -45,11 +44,9 @@ @t.final class Job(BaseJob): """A Job holds a reference to a SmartSimEntity and associated - LaunchSettings prior to launch. It is responsible for turning - the stored entity and launch settings into commands that can be - executed by a launcher. - - Jobs will hold a deep copy of launch settings. + LaunchSettings prior to launch. It is responsible for turning + the stored SmartSimEntity and LaunchSettings into commands that can be + executed by a launcher. Jobs are designed to be started by the Experiment. """ def __init__( @@ -58,47 +55,91 @@ def __init__( launch_settings: LaunchSettings, name: str | None = None, ): + """Initialize a ``Job`` + + Jobs require a SmartSimEntity and a LaunchSettings. Optionally, users may provide + a name. To create a simple Job that echos `Hello World!`, consider the example below: + + .. highlight:: python + .. code-block:: python + + # Create an application that runs the 'echo' command + my_app = Application(name="my_app", exe="echo", exe_args="Hello World!") + # Define the launch settings using SLURM + srun_settings = LaunchSettings(launcher="slurm") + + # Create a Job with the `my_app` and `srun_settings` + my_job = Job(my_app, srun_settings, name="my_job") + + :param entity: the SmartSimEntity object + :param launch_settings: the LaunchSettings object + :param name: the Job name + """ super().__init__() + """Initialize the parent class BaseJob""" self._entity = deepcopy(entity) + """Deepcopy of the SmartSimEntity object""" self._launch_settings = deepcopy(launch_settings) + """Deepcopy of the LaunchSettings object""" self._name = name if name else entity.name + """Name of the Job""" check_name(self._name) @property def name(self) -> str: - """Retrieves the name of the Job.""" + """Return the name of the Job. + + :return: the name of the Job + """ return self._name @name.setter def name(self, name: str) -> None: - """Sets the name of the Job.""" + """Set the name of the Job. + + :param name: the name of the Job + """ check_name(name) logger.debug(f'Overwriting the Job name from "{self._name}" to "{name}"') self._name = name @property def entity(self) -> SmartSimEntity: - """Retrieves the Job entity.""" + """Return the attached entity. + + :return: the attached SmartSimEntity + """ return deepcopy(self._entity) @entity.setter def entity(self, value: SmartSimEntity) -> None: - """Sets the Job entity.""" + """Set the Job entity. + + :param value: the SmartSimEntity + """ self._entity = deepcopy(value) @property def launch_settings(self) -> LaunchSettings: - """Retrieves the Job LaunchSettings.""" + """Return the attached LaunchSettings. + + :return: the attached LaunchSettings + """ return deepcopy(self._launch_settings) @launch_settings.setter def launch_settings(self, value: LaunchSettings) -> None: - """Sets the Job LaunchSettings.""" + """Set the Jobs LaunchSettings. + + :param value: the LaunchSettings + """ self._launch_settings = deepcopy(value) def get_launch_steps(self) -> LaunchCommands: """Return the launch steps corresponding to the internal data. + + :returns: The Jobs launch steps """ # TODO: return JobWarehouseRunner.run(self) raise NotImplementedError diff --git a/smartsim/settings/arguments/batch/lsf.py b/smartsim/settings/arguments/batch/lsf.py index 5e7565afb..23f948bd0 100644 --- a/smartsim/settings/arguments/batch/lsf.py +++ b/smartsim/settings/arguments/batch/lsf.py @@ -30,7 +30,7 @@ from smartsim.log import get_logger -from ...batch_command import SchedulerType +from ...batch_command import BatchSchedulerType from ...common import StringArgument from ..batch_arguments import BatchArguments @@ -38,12 +38,16 @@ class BsubBatchArguments(BatchArguments): + """A class to represent the arguments required for submitting batch + jobs using the bsub command. + """ + def scheduler_str(self) -> str: """Get the string representation of the scheduler :returns: The string representation of the scheduler """ - return SchedulerType.Lsf.value + return BatchSchedulerType.Lsf.value def set_walltime(self, walltime: str) -> None: """Set the walltime @@ -137,7 +141,7 @@ def format_batch_args(self) -> t.List[str]: """ opts = [] - for opt, value in self._scheduler_args.items(): + for opt, value in self._batch_args.items(): prefix = "-" # LSF only uses single dashses @@ -156,4 +160,4 @@ def set(self, key: str, value: str | None) -> None: argument (if applicable), otherwise `None` """ # Store custom arguments in the launcher_args - self._scheduler_args[key] = value + self._batch_args[key] = value diff --git a/smartsim/settings/arguments/batch/pbs.py b/smartsim/settings/arguments/batch/pbs.py index 7f03642df..126207665 100644 --- a/smartsim/settings/arguments/batch/pbs.py +++ b/smartsim/settings/arguments/batch/pbs.py @@ -32,7 +32,7 @@ from smartsim.log import get_logger from ....error import SSConfigError -from ...batch_command import SchedulerType +from ...batch_command import BatchSchedulerType from ...common import StringArgument from ..batch_arguments import BatchArguments @@ -40,12 +40,16 @@ class QsubBatchArguments(BatchArguments): + """A class to represent the arguments required for submitting batch + jobs using the qsub command. + """ + def scheduler_str(self) -> str: """Get the string representation of the scheduler :returns: The string representation of the scheduler """ - return SchedulerType.Pbs.value + return BatchSchedulerType.Pbs.value def set_nodes(self, num_nodes: int) -> None: """Set the number of nodes for this batch job @@ -119,7 +123,7 @@ def format_batch_args(self) -> t.List[str]: :return: batch arguments for `qsub` :raises ValueError: if options are supplied without values """ - opts, batch_arg_copy = self._create_resource_list(self._scheduler_args) + opts, batch_arg_copy = self._create_resource_list(self._batch_args) for opt, value in batch_arg_copy.items(): prefix = "-" if not value: @@ -179,4 +183,4 @@ def set(self, key: str, value: str | None) -> None: :param value: A string representation of the value for the launch argument (if applicable), otherwise `None` """ - self._scheduler_args[key] = value + self._batch_args[key] = value diff --git a/smartsim/settings/arguments/batch/slurm.py b/smartsim/settings/arguments/batch/slurm.py index 7114e947e..26f9cf854 100644 --- a/smartsim/settings/arguments/batch/slurm.py +++ b/smartsim/settings/arguments/batch/slurm.py @@ -31,7 +31,7 @@ from smartsim.log import get_logger -from ...batch_command import SchedulerType +from ...batch_command import BatchSchedulerType from ...common import StringArgument from ..batch_arguments import BatchArguments @@ -39,12 +39,16 @@ class SlurmBatchArguments(BatchArguments): + """A class to represent the arguments required for submitting batch + jobs using the sbatch command. + """ + def scheduler_str(self) -> str: """Get the string representation of the scheduler :returns: The string representation of the scheduler """ - return SchedulerType.Slurm.value + return BatchSchedulerType.Slurm.value def set_walltime(self, walltime: str) -> None: """Set the walltime of the job @@ -127,7 +131,7 @@ def format_batch_args(self) -> t.List[str]: """ opts = [] # TODO add restricted here - for opt, value in self._scheduler_args.items(): + for opt, value in self._batch_args.items(): # attach "-" prefix if argument is 1 character otherwise "--" short_arg = len(opt) == 1 prefix = "-" if short_arg else "--" @@ -149,4 +153,4 @@ def set(self, key: str, value: str | None) -> None: argument (if applicable), otherwise `None` """ # Store custom arguments in the launcher_args - self._scheduler_args[key] = value + self._batch_args[key] = value diff --git a/smartsim/settings/arguments/batch_arguments.py b/smartsim/settings/arguments/batch_arguments.py index a85148697..0fa8d3964 100644 --- a/smartsim/settings/arguments/batch_arguments.py +++ b/smartsim/settings/arguments/batch_arguments.py @@ -44,8 +44,9 @@ class BatchArguments(ABC): the input parameter to a properly formatted launcher argument. """ - def __init__(self, scheduler_args: t.Dict[str, str | None] | None) -> None: - self._scheduler_args = copy.deepcopy(scheduler_args) or {} + def __init__(self, batch_args: t.Dict[str, str | None] | None) -> None: + self._batch_args = copy.deepcopy(batch_args) or {} + """A dictionary of batch arguments""" @abstractmethod def scheduler_str(self) -> str: @@ -104,5 +105,5 @@ def format_batch_args(self) -> t.List[str]: pass def __str__(self) -> str: # pragma: no-cover - string = f"\nScheduler Arguments:\n{fmt_dict(self._scheduler_args)}" + string = f"\nScheduler Arguments:\n{fmt_dict(self._batch_args)}" return string diff --git a/smartsim/settings/arguments/launch_arguments.py b/smartsim/settings/arguments/launch_arguments.py index 0e011339e..6ec741d91 100644 --- a/smartsim/settings/arguments/launch_arguments.py +++ b/smartsim/settings/arguments/launch_arguments.py @@ -50,6 +50,7 @@ def __init__(self, launch_args: t.Dict[str, str | None] | None) -> None: :param launch_args: A mapping of arguments to (optional) values """ self._launch_args = copy.deepcopy(launch_args) or {} + """A dictionary of launch arguments""" @abstractmethod def launcher_str(self) -> str: diff --git a/smartsim/settings/base_settings.py b/smartsim/settings/base_settings.py index 1acd5f605..2e8a87f57 100644 --- a/smartsim/settings/base_settings.py +++ b/smartsim/settings/base_settings.py @@ -23,7 +23,8 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# fmt: off + class BaseSettings: - ... -# fmt: on + """ + A base class for LaunchSettings and BatchSettings. + """ diff --git a/smartsim/settings/batch_command.py b/smartsim/settings/batch_command.py index 8f3b0c89d..a96492d39 100644 --- a/smartsim/settings/batch_command.py +++ b/smartsim/settings/batch_command.py @@ -27,10 +27,8 @@ from enum import Enum -class SchedulerType(Enum): - """Schedulers that are supported by - SmartSim. - """ +class BatchSchedulerType(Enum): + """Schedulers supported by SmartSim.""" Slurm = "slurm" Pbs = "pbs" diff --git a/smartsim/settings/batch_settings.py b/smartsim/settings/batch_settings.py index 10aea8377..734e919ce 100644 --- a/smartsim/settings/batch_settings.py +++ b/smartsim/settings/batch_settings.py @@ -37,77 +37,138 @@ from .arguments.batch.pbs import QsubBatchArguments from .arguments.batch.slurm import SlurmBatchArguments from .base_settings import BaseSettings -from .batch_command import SchedulerType +from .batch_command import BatchSchedulerType from .common import StringArgument logger = get_logger(__name__) class BatchSettings(BaseSettings): + """The BatchSettings class stores scheduler configuration settings and is + used to inject scheduler-specific behavior into a job. + + BatchSettings is designed to be extended by a BatchArguments child class that + corresponds to the scheduler provided during initialization. The supported schedulers + are Slurm, PBS, and LSF. Using the BatchSettings class, users can: + + - Set the scheduler type of a batch job. + - Configure batch arguments and environment variables. + - Access and modify custom batch arguments. + - Update environment variables. + - Retrieve information associated with the ``BatchSettings`` object. + - The scheduler value (BatchSettings.scheduler). + - The derived BatchArguments child class (BatchSettings.batch_args). + - The set environment variables (BatchSettings.env_vars). + - A formatted output of set batch arguments (BatchSettings.format_batch_args). + """ + def __init__( self, - batch_scheduler: t.Union[SchedulerType, str], - scheduler_args: t.Dict[str, t.Union[str, None]] | None = None, + batch_scheduler: t.Union[BatchSchedulerType, str], + batch_args: StringArgument | None = None, env_vars: StringArgument | None = None, ) -> None: + """Initialize a BatchSettings instance. + + The "batch_scheduler" of SmartSim BatchSettings will determine the + child type assigned to the BatchSettings.batch_args attribute. + For example, to configure a job for SLURM batch jobs, assign BatchSettings.batch_scheduler + to "slurm" or BatchSchedulerType.Slurm: + + .. highlight:: python + .. code-block:: python + + sbatch_settings = BatchSettings(batch_scheduler="slurm") + # OR + sbatch_settings = BatchSettings(batch_scheduler=BatchSchedulerType.Slurm) + + This will assign a SlurmBatchArguments object to ``sbatch_settings.batch_args``. + Using the object, users may access the child class functions to set + batch configurations. For example: + + .. highlight:: python + .. code-block:: python + + sbatch_settings.batch_args.set_nodes(5) + sbatch_settings.batch_args.set_cpus_per_task(2) + + To set customized batch arguments, use the `set()` function provided by + the BatchSettings child class. For example: + + .. highlight:: python + .. code-block:: python + + sbatch_settings.batch_args.set(key="nodes", value="6") + + If the key already exists in the existing batch arguments, the value will + be overwritten. + + :param batch_scheduler: The type of scheduler to initialize (e.g., Slurm, PBS, LSF) + :param batch_args: A dictionary of arguments for the scheduler, where the keys + are strings and the values can be either strings or None. This argument is optional + and defaults to None. + :param env_vars: Environment variables for the batch settings, where the keys + are strings and the values can be either strings or None. This argument is + also optional and defaults to None. + :raises ValueError: Raises if the scheduler provided does not exist. + """ try: - self._batch_scheduler = SchedulerType(batch_scheduler) + self._batch_scheduler = BatchSchedulerType(batch_scheduler) + """The scheduler type""" except ValueError: raise ValueError(f"Invalid scheduler type: {batch_scheduler}") from None - self._arguments = self._get_arguments(scheduler_args) + self._arguments = self._get_arguments(batch_args) + """The BatchSettings child class based on scheduler type""" self.env_vars = env_vars or {} - - @property - def scheduler(self) -> str: - """Return the launcher name.""" - return self._batch_scheduler.value + """The environment configuration""" @property def batch_scheduler(self) -> str: - """Return the scheduler name.""" + """Return the scheduler type.""" return self._batch_scheduler.value @property - def scheduler_args(self) -> BatchArguments: - """Return the batch argument translator.""" + def batch_args(self) -> BatchArguments: + """Return the BatchArguments child class.""" return self._arguments @property def env_vars(self) -> StringArgument: """Return an immutable list of attached environment variables.""" - return copy.deepcopy(self._env_vars) + return self._env_vars @env_vars.setter def env_vars(self, value: t.Dict[str, str | None]) -> None: """Set the environment variables.""" self._env_vars = copy.deepcopy(value) - def _get_arguments(self, scheduler_args: StringArgument | None) -> BatchArguments: + def _get_arguments(self, batch_args: StringArgument | None) -> BatchArguments: """Map the Scheduler to the BatchArguments. This method should only be called once during construction. - :param scheduler_args: A mapping of arguments names to values to be + :param schedule_args: A mapping of arguments names to values to be used to initialize the arguments :returns: The appropriate type for the settings instance. + :raises ValueError: An invalid scheduler type was provided. """ - if self._batch_scheduler == SchedulerType.Slurm: - return SlurmBatchArguments(scheduler_args) - elif self._batch_scheduler == SchedulerType.Lsf: - return BsubBatchArguments(scheduler_args) - elif self._batch_scheduler == SchedulerType.Pbs: - return QsubBatchArguments(scheduler_args) + if self._batch_scheduler == BatchSchedulerType.Slurm: + return SlurmBatchArguments(batch_args) + elif self._batch_scheduler == BatchSchedulerType.Lsf: + return BsubBatchArguments(batch_args) + elif self._batch_scheduler == BatchSchedulerType.Pbs: + return QsubBatchArguments(batch_args) else: raise ValueError(f"Invalid scheduler type: {self._batch_scheduler}") def format_batch_args(self) -> t.List[str]: - """Get the formatted batch arguments for a preview + """Get the formatted batch arguments to preview - :return: batch arguments for Sbatch + :return: formatted batch arguments """ return self._arguments.format_batch_args() def __str__(self) -> str: # pragma: no-cover - string = f"\nScheduler: {self.scheduler}{self.scheduler_args}" + string = f"\nBatch Scheduler: {self.batch_scheduler}{self.batch_args}" if self.env_vars: string += f"\nEnvironment variables: \n{fmt_dict(self.env_vars)}" return string diff --git a/smartsim/settings/launch_command.py b/smartsim/settings/launch_command.py index 491f01d86..b848e35e1 100644 --- a/smartsim/settings/launch_command.py +++ b/smartsim/settings/launch_command.py @@ -28,9 +28,7 @@ class LauncherType(Enum): - """Launchers that are supported by - SmartSim. - """ + """Launchers supported by SmartSim.""" Dragon = "dragon" Slurm = "slurm" diff --git a/smartsim/settings/launch_settings.py b/smartsim/settings/launch_settings.py index 6d7da57ca..7b6083022 100644 --- a/smartsim/settings/launch_settings.py +++ b/smartsim/settings/launch_settings.py @@ -52,18 +52,84 @@ class LaunchSettings(BaseSettings): + """The LaunchSettings class stores launcher configuration settings and is + used to inject launcher-specific behavior into a job. + + LaunchSettings is designed to be extended by a LaunchArguments child class that + corresponds to the launcher provided during initialization. The supported launchers + are Dragon, Slurm, PALS, ALPS, Local, Mpiexec, Mpirun, Orterun, and LSF. Using the + LaunchSettings class, users can: + + - Set the launcher type of a job. + - Configure launch arguments and environment variables. + - Access and modify custom launch arguments. + - Update environment variables. + - Retrieve information associated with the ``LaunchSettings`` object. + - The launcher value (LaunchSettings.launcher). + - The derived LaunchSettings child class (LaunchSettings.launch_args). + - The set environment variables (LaunchSettings.env_vars). + """ + def __init__( self, launcher: t.Union[LauncherType, str], launch_args: StringArgument | None = None, env_vars: StringArgument | None = None, ) -> None: + """Initialize a LaunchSettings instance. + + The "launcher" of SmartSim LaunchSettings will determine the + child type assigned to the LaunchSettings.launch_args attribute. + For example, to configure a job for SLURM, assign LaunchSettings.launcher + to "slurm" or LauncherType.Slurm: + + .. highlight:: python + .. code-block:: python + + srun_settings = LaunchSettings(launcher="slurm") + # OR + srun_settings = LaunchSettings(launcher=LauncherType.Slurm) + + This will assign a SlurmLaunchArguments object to ``srun_settings.launch_args``. + Using the object, users may access the child class functions to set + batch configurations. For example: + + .. highlight:: python + .. code-block:: python + + srun_settings.launch_args.set_nodes(5) + srun_settings.launch_args.set_cpus_per_task(2) + + To set customized launch arguments, use the `set()`function provided by + the LaunchSettings child class. For example: + + .. highlight:: python + .. code-block:: python + + srun_settings.launch_args.set(key="nodes", value="6") + + If the key already exists in the existing launch arguments, the value will + be overwritten. + + :param launcher: The type of launcher to initialize (e.g., Dragon, Slurm, + PALS, ALPS, Local, Mpiexec, Mpirun, Orterun, LSF) + :param launch_args: A dictionary of arguments for the launcher, where the keys + are strings and the values can be either strings or None. This argument is optional + and defaults to None. + :param env_vars: Environment variables for the launch settings, where the keys + are strings and the values can be either strings or None. This argument is + also optional and defaults to None. + :raises ValueError: Raises if the launcher provided does not exist. + """ try: self._launcher = LauncherType(launcher) + """The launcher type""" except ValueError: raise ValueError(f"Invalid launcher type: {launcher}") self._arguments = self._get_arguments(launch_args) + """The LaunchSettings child class based on launcher type""" self.env_vars = env_vars or {} + """The environment configuration""" @property def launcher(self) -> str: @@ -89,7 +155,7 @@ def env_vars(self) -> t.Mapping[str, str | None]: :returns: An environment mapping """ - return copy.deepcopy(self._env_vars) + return self._env_vars @env_vars.setter def env_vars(self, value: dict[str, str | None]) -> None: @@ -108,6 +174,7 @@ def _get_arguments(self, launch_args: StringArgument | None) -> LaunchArguments: :param launch_args: A mapping of arguments names to values to be used to initialize the arguments :returns: The appropriate type for the settings instance. + :raises ValueError: An invalid launcher type was provided. """ if self._launcher == LauncherType.Slurm: return SlurmLaunchArguments(launch_args) diff --git a/tests/temp_tests/test_settings/test_batchSettings.py b/tests/temp_tests/test_settings/test_batchSettings.py index e7fd4b5ff..37fd3a33f 100644 --- a/tests/temp_tests/test_settings/test_batchSettings.py +++ b/tests/temp_tests/test_settings/test_batchSettings.py @@ -26,38 +26,46 @@ import pytest from smartsim.settings import BatchSettings -from smartsim.settings.batch_command import SchedulerType +from smartsim.settings.batch_command import BatchSchedulerType pytestmark = pytest.mark.group_a @pytest.mark.parametrize( - "scheduler_enum", + "scheduler_enum,formatted_batch_args", [ - pytest.param(SchedulerType.Slurm, id="slurm"), - pytest.param(SchedulerType.Pbs, id="dragon"), - pytest.param(SchedulerType.Lsf, id="lsf"), + pytest.param( + BatchSchedulerType.Slurm, ["--launch=var", "--nodes=1"], id="slurm" + ), + pytest.param( + BatchSchedulerType.Pbs, ["-l", "nodes=1", "-launch", "var"], id="pbs" + ), + pytest.param( + BatchSchedulerType.Lsf, ["-launch", "var", "-nnodes", "1"], id="lsf" + ), ], ) -def test_create_scheduler_settings(scheduler_enum): +def test_create_scheduler_settings(scheduler_enum, formatted_batch_args): bs_str = BatchSettings( batch_scheduler=scheduler_enum.value, - scheduler_args={"launch": "var"}, + batch_args={"launch": "var"}, env_vars={"ENV": "VAR"}, ) - print(bs_str) + bs_str.batch_args.set_nodes(1) assert bs_str._batch_scheduler == scheduler_enum - # TODO need to test scheduler_args assert bs_str._env_vars == {"ENV": "VAR"} + print(bs_str.format_batch_args()) + assert bs_str.format_batch_args() == formatted_batch_args bs_enum = BatchSettings( batch_scheduler=scheduler_enum, - scheduler_args={"launch": "var"}, + batch_args={"launch": "var"}, env_vars={"ENV": "VAR"}, ) + bs_enum.batch_args.set_nodes(1) assert bs_enum._batch_scheduler == scheduler_enum - # TODO need to test scheduler_args assert bs_enum._env_vars == {"ENV": "VAR"} + assert bs_enum.format_batch_args() == formatted_batch_args def test_launcher_property(): @@ -68,10 +76,5 @@ def test_launcher_property(): def test_env_vars_property(): bs = BatchSettings(batch_scheduler="slurm", env_vars={"ENV": "VAR"}) assert bs.env_vars == {"ENV": "VAR"} - - -def test_env_vars_property_deep_copy(): - bs = BatchSettings(batch_scheduler="slurm", env_vars={"ENV": "VAR"}) - copy_env_vars = bs.env_vars - copy_env_vars.update({"test": "no_update"}) - assert bs.env_vars == {"ENV": "VAR"} + ref = bs.env_vars + assert ref is bs.env_vars diff --git a/tests/temp_tests/test_settings/test_launchSettings.py b/tests/temp_tests/test_settings/test_launchSettings.py index e06cf2939..3fc5e544a 100644 --- a/tests/temp_tests/test_settings/test_launchSettings.py +++ b/tests/temp_tests/test_settings/test_launchSettings.py @@ -64,13 +64,8 @@ def test_launcher_property(): def test_env_vars_property(): ls = LaunchSettings(launcher="local", env_vars={"ENV": "VAR"}) assert ls.env_vars == {"ENV": "VAR"} - - -def test_env_vars_property_deep_copy(): - ls = LaunchSettings(launcher="local", env_vars={"ENV": "VAR"}) - copy_env_vars = ls.env_vars - copy_env_vars.update({"test": "no_update"}) - assert ls.env_vars == {"ENV": "VAR"} + ref = ls.env_vars + assert ref is ls.env_vars def test_update_env_vars(): diff --git a/tests/temp_tests/test_settings/test_lsfScheduler.py b/tests/temp_tests/test_settings/test_lsfScheduler.py index afb73d45b..5e6b7fd0c 100644 --- a/tests/temp_tests/test_settings/test_lsfScheduler.py +++ b/tests/temp_tests/test_settings/test_lsfScheduler.py @@ -26,15 +26,15 @@ import pytest from smartsim.settings import BatchSettings -from smartsim.settings.batch_command import SchedulerType +from smartsim.settings.batch_command import BatchSchedulerType pytestmark = pytest.mark.group_a def test_scheduler_str(): """Ensure scheduler_str returns appropriate value""" - bs = BatchSettings(batch_scheduler=SchedulerType.Lsf) - assert bs.scheduler_args.scheduler_str() == SchedulerType.Lsf.value + bs = BatchSettings(batch_scheduler=BatchSchedulerType.Lsf) + assert bs.batch_args.scheduler_str() == BatchSchedulerType.Lsf.value @pytest.mark.parametrize( @@ -60,18 +60,18 @@ def test_scheduler_str(): ], ) def test_update_env_initialized(function, value, flag, result): - lsfScheduler = BatchSettings(batch_scheduler=SchedulerType.Lsf) - getattr(lsfScheduler.scheduler_args, function)(*value) - assert lsfScheduler.scheduler_args._scheduler_args[flag] == result + lsfScheduler = BatchSettings(batch_scheduler=BatchSchedulerType.Lsf) + getattr(lsfScheduler.batch_args, function)(*value) + assert lsfScheduler.batch_args._batch_args[flag] == result def test_create_bsub(): batch_args = {"core_isolation": None} lsfScheduler = BatchSettings( - batch_scheduler=SchedulerType.Lsf, scheduler_args=batch_args + batch_scheduler=BatchSchedulerType.Lsf, batch_args=batch_args ) - lsfScheduler.scheduler_args.set_nodes(1) - lsfScheduler.scheduler_args.set_walltime("10:10:10") - lsfScheduler.scheduler_args.set_queue("default") + lsfScheduler.batch_args.set_nodes(1) + lsfScheduler.batch_args.set_walltime("10:10:10") + lsfScheduler.batch_args.set_queue("default") args = lsfScheduler.format_batch_args() assert args == ["-core_isolation", "-nnodes", "1", "-W", "10:10", "-q", "default"] diff --git a/tests/temp_tests/test_settings/test_pbsScheduler.py b/tests/temp_tests/test_settings/test_pbsScheduler.py index 642d115ac..36fde6776 100644 --- a/tests/temp_tests/test_settings/test_pbsScheduler.py +++ b/tests/temp_tests/test_settings/test_pbsScheduler.py @@ -27,15 +27,15 @@ from smartsim.settings import BatchSettings from smartsim.settings.arguments.batch.pbs import QsubBatchArguments -from smartsim.settings.batch_command import SchedulerType +from smartsim.settings.batch_command import BatchSchedulerType pytestmark = pytest.mark.group_a def test_scheduler_str(): """Ensure scheduler_str returns appropriate value""" - bs = BatchSettings(batch_scheduler=SchedulerType.Pbs) - assert bs.scheduler_args.scheduler_str() == SchedulerType.Pbs.value + bs = BatchSettings(batch_scheduler=BatchSchedulerType.Pbs) + assert bs.batch_args.scheduler_str() == BatchSchedulerType.Pbs.value @pytest.mark.parametrize( @@ -61,20 +61,20 @@ def test_scheduler_str(): ], ) def test_create_pbs_batch(function, value, flag, result): - pbsScheduler = BatchSettings(batch_scheduler=SchedulerType.Pbs) - assert isinstance(pbsScheduler.scheduler_args, QsubBatchArguments) - getattr(pbsScheduler.scheduler_args, function)(*value) - assert pbsScheduler.scheduler_args._scheduler_args[flag] == result + pbsScheduler = BatchSettings(batch_scheduler=BatchSchedulerType.Pbs) + assert isinstance(pbsScheduler.batch_args, QsubBatchArguments) + getattr(pbsScheduler.batch_args, function)(*value) + assert pbsScheduler.batch_args._batch_args[flag] == result def test_format_pbs_batch_args(): - pbsScheduler = BatchSettings(batch_scheduler=SchedulerType.Pbs) - pbsScheduler.scheduler_args.set_nodes(1) - pbsScheduler.scheduler_args.set_walltime("10:00:00") - pbsScheduler.scheduler_args.set_queue("default") - pbsScheduler.scheduler_args.set_account("myproject") - pbsScheduler.scheduler_args.set_ncpus(10) - pbsScheduler.scheduler_args.set_hostlist(["host_a", "host_b", "host_c"]) + pbsScheduler = BatchSettings(batch_scheduler=BatchSchedulerType.Pbs) + pbsScheduler.batch_args.set_nodes(1) + pbsScheduler.batch_args.set_walltime("10:00:00") + pbsScheduler.batch_args.set_queue("default") + pbsScheduler.batch_args.set_account("myproject") + pbsScheduler.batch_args.set_ncpus(10) + pbsScheduler.batch_args.set_hostlist(["host_a", "host_b", "host_c"]) args = pbsScheduler.format_batch_args() assert args == [ "-l", diff --git a/tests/temp_tests/test_settings/test_slurmScheduler.py b/tests/temp_tests/test_settings/test_slurmScheduler.py index 94fa213da..8ab489cc8 100644 --- a/tests/temp_tests/test_settings/test_slurmScheduler.py +++ b/tests/temp_tests/test_settings/test_slurmScheduler.py @@ -27,15 +27,15 @@ from smartsim.settings import BatchSettings from smartsim.settings.arguments.batch.slurm import SlurmBatchArguments -from smartsim.settings.batch_command import SchedulerType +from smartsim.settings.batch_command import BatchSchedulerType pytestmark = pytest.mark.group_a -def test_scheduler_str(): +def test_batch_scheduler_str(): """Ensure scheduler_str returns appropriate value""" - bs = BatchSettings(batch_scheduler=SchedulerType.Slurm) - assert bs.scheduler_args.scheduler_str() == SchedulerType.Slurm.value + bs = BatchSettings(batch_scheduler=BatchSchedulerType.Slurm) + assert bs.batch_args.scheduler_str() == BatchSchedulerType.Slurm.value @pytest.mark.parametrize( @@ -74,15 +74,15 @@ def test_scheduler_str(): ], ) def test_sbatch_class_methods(function, value, flag, result): - slurmScheduler = BatchSettings(batch_scheduler=SchedulerType.Slurm) - getattr(slurmScheduler.scheduler_args, function)(*value) - assert slurmScheduler.scheduler_args._scheduler_args[flag] == result + slurmScheduler = BatchSettings(batch_scheduler=BatchSchedulerType.Slurm) + getattr(slurmScheduler.batch_args, function)(*value) + assert slurmScheduler.batch_args._batch_args[flag] == result def test_create_sbatch(): batch_args = {"exclusive": None, "oversubscribe": None} slurmScheduler = BatchSettings( - batch_scheduler=SchedulerType.Slurm, scheduler_args=batch_args + batch_scheduler=BatchSchedulerType.Slurm, batch_args=batch_args ) assert isinstance(slurmScheduler._arguments, SlurmBatchArguments) args = slurmScheduler.format_batch_args() @@ -94,32 +94,32 @@ def test_launch_args_input_mutation(): key0, key1, key2 = "arg0", "arg1", "arg2" val0, val1, val2 = "val0", "val1", "val2" - default_scheduler_args = { + default_batch_args = { key0: val0, key1: val1, key2: val2, } slurmScheduler = BatchSettings( - batch_scheduler=SchedulerType.Slurm, scheduler_args=default_scheduler_args + batch_scheduler=BatchSchedulerType.Slurm, batch_args=default_batch_args ) # Confirm initial values are set - assert slurmScheduler.scheduler_args._scheduler_args[key0] == val0 - assert slurmScheduler.scheduler_args._scheduler_args[key1] == val1 - assert slurmScheduler.scheduler_args._scheduler_args[key2] == val2 + assert slurmScheduler.batch_args._batch_args[key0] == val0 + assert slurmScheduler.batch_args._batch_args[key1] == val1 + assert slurmScheduler.batch_args._batch_args[key2] == val2 # Update our common run arguments val2_upd = f"not-{val2}" - default_scheduler_args[key2] = val2_upd + default_batch_args[key2] = val2_upd # Confirm previously created run settings are not changed - assert slurmScheduler.scheduler_args._scheduler_args[key2] == val2 + assert slurmScheduler.batch_args._batch_args[key2] == val2 def test_sbatch_settings(): - scheduler_args = {"nodes": 1, "time": "10:00:00", "account": "A3123"} + batch_args = {"nodes": 1, "time": "10:00:00", "account": "A3123"} slurmScheduler = BatchSettings( - batch_scheduler=SchedulerType.Slurm, scheduler_args=scheduler_args + batch_scheduler=BatchSchedulerType.Slurm, batch_args=batch_args ) formatted = slurmScheduler.format_batch_args() result = ["--nodes=1", "--time=10:00:00", "--account=A3123"] @@ -127,10 +127,10 @@ def test_sbatch_settings(): def test_sbatch_manual(): - slurmScheduler = BatchSettings(batch_scheduler=SchedulerType.Slurm) - slurmScheduler.scheduler_args.set_nodes(5) - slurmScheduler.scheduler_args.set_account("A3531") - slurmScheduler.scheduler_args.set_walltime("10:00:00") + slurmScheduler = BatchSettings(batch_scheduler=BatchSchedulerType.Slurm) + slurmScheduler.batch_args.set_nodes(5) + slurmScheduler.batch_args.set_account("A3531") + slurmScheduler.batch_args.set_walltime("10:00:00") formatted = slurmScheduler.format_batch_args() result = ["--nodes=5", "--account=A3531", "--time=10:00:00"] assert formatted == result