class Srun(JobManager): """A class for the srun job manager command.""" def __init__(self, job, path="", subprocess=False): """Create a Srun object. Args: job (ExecutableCommand): command object to manage. path (str, optional): path to location of command binary file. Defaults to "". subprocess (bool, optional): whether the command is run as a subprocess. Defaults to False. """ super(Srun, self).__init__("/run/srun", "srun", job, path, subprocess) self.label = FormattedParameter("--label", True) self.mpi = FormattedParameter("--mpi={}", "pmi2") self.export = FormattedParameter("--export={}", "ALL") self.ntasks = FormattedParameter("--ntasks={}", None) self.distribution = FormattedParameter("--distribution={}", None) self.nodefile = FormattedParameter("--nodefile={}", None) self.nodelist = FormattedParameter("--nodelist={}", None) self.ntasks_per_node = FormattedParameter("--ntasks-per-node={}", None) self.nodes = FormattedParameter("--nodes={}", None) self.reservation = FormattedParameter("--reservation={}", None) self.partition = FormattedParameter("--partition={}", None) self.output = FormattedParameter("--output={}", None) def assign_hosts(self, hosts, path=None, slots=None): """Assign the hosts to use with the command (-f). Args: hosts (list): list of hosts to specify in the hostfile path (str, optional): hostfile path. Defaults to None. slots (int, optional): number of slots per host to specify in the hostfile. Defaults to None. """ self._hosts = hosts kwargs = {"hostlist": self._hosts, "slots": None} if path is not None: kwargs["path"] = path self.nodefile.value = write_host_file(**kwargs) self.ntasks_per_node.value = slots def assign_processes(self, processes): """Assign the number of processes per node (--ntasks). Args: processes (int): number of processes per node """ self.ntasks.value = processes self.distribution.value = "cyclic" def assign_environment(self, env_vars, append=False): """Assign or add environment variables to the command. Args: env_vars (EnvironmentVariables): the environment variables to use assign or add to the command append (bool): whether to assign (False) or append (True) the specified environment variables """ if append and self.export.value is not None: # Convert the current list of environmental variable assignments # into an EnvironmentVariables (dict) object. Then update the # dictionary keys with the specified values or add new key value # pairs to the dictionary. Finally convert the updated dictionary # back to a string for the parameter assignment. original = EnvironmentVariables({ item.split("=")[0]: item.split("=")[1] if "=" in item else None for item in self.export.value.split(",") }) original.update(env_vars) self.export.value = ",".join(original.get_list()) else: # Overwrite the environmental variable assignment self.export.value = ",".join(env_vars.get_list()) def assign_environment_default(self, env_vars): """Assign the default environment variables for the command. Args: env_vars (EnvironmentVariables): the environment variables to assign as the default """ self.export.update_default(env_vars.get_list())
class Orterun(JobManager): """A class for the orterun job manager command.""" def __init__(self, job, subprocess=False): """Create a Orterun object. Args: job (ExecutableCommand): command object to manage. subprocess (bool, optional): whether the command is run as a subprocess. Defaults to False. """ load_mpi("openmpi") path = os.path.dirname(find_executable("orterun")) super(Orterun, self).__init__("/run/orterun/*", "orterun", job, path, subprocess) # Default mca values to avoid queue pair errors mca_default = { "btl_openib_warn_default_gid_prefix": "0", "btl": "tcp,self", "oob": "tcp", "pml": "ob1", } self.hostfile = FormattedParameter("--hostfile {}", None) self.processes = FormattedParameter("--np {}", 1) self.display_map = FormattedParameter("--display-map", False) self.map_by = FormattedParameter("--map-by {}", "node") self.export = FormattedParameter("-x {}", None) self.enable_recovery = FormattedParameter("--enable-recovery", True) self.report_uri = FormattedParameter("--report-uri {}", None) self.allow_run_as_root = FormattedParameter("--allow-run-as-root", None) self.mca = FormattedParameter("--mca {}", mca_default) self.pprnode = FormattedParameter("--map-by ppr:{}:node", None) self.tag_output = FormattedParameter("--tag-output", True) self.ompi_server = FormattedParameter("--ompi-server {}", None) self.working_dir = FormattedParameter("-wdir {}", None) def assign_hosts(self, hosts, path=None, slots=None): """Assign the hosts to use with the command (--hostfile). Args: hosts (list): list of hosts to specify in the hostfile path (str, optional): hostfile path. Defaults to None. slots (int, optional): number of slots per host to specify in the hostfile. Defaults to None. """ self._hosts = hosts kwargs = {"hostlist": self._hosts, "slots": slots} if path is not None: kwargs["path"] = path self.hostfile.value = write_host_file(**kwargs) def assign_processes(self, processes): """Assign the number of processes per node (-np). Args: processes (int): number of processes per node """ self.processes.value = processes def assign_environment(self, env_vars, append=False): """Assign or add environment variables to the command. Args: env_vars (EnvironmentVariables): the environment variables to use assign or add to the command append (bool): whether to assign (False) or append (True) the specified environment variables """ if append and self.export.value is not None: # Convert the current list of environmental variable assignments # into an EnvironmentVariables (dict) object. Then update the # dictionary keys with the specified values or add new key value # pairs to the dictionary. Finally convert the updated dictionary # back to a list for the parameter assignment. original = EnvironmentVariables({ item.split("=")[0]: item.split("=")[1] if "=" in item else None for item in self.export.value }) original.update(env_vars) self.export.value = original.get_list() else: # Overwrite the environmental variable assignment self.export.value = env_vars.get_list() def assign_environment_default(self, env_vars): """Assign the default environment variables for the command. Args: env_vars (EnvironmentVariables): the environment variables to assign as the default """ self.export.update_default(env_vars.get_list()) def run(self): """Run the orterun command. Raises: CommandFailure: if there is an error running the command """ load_mpi("openmpi") return super(Orterun, self).run()
class Mpirun(JobManager): """A class for the mpirun job manager command.""" def __init__(self, job, subprocess=False, mpitype="openmpi"): """Create a Mpirun object. Args: job (ExecutableCommand): command object to manage. subprocess (bool, optional): whether the command is run as a subprocess. Defaults to False. """ load_mpi(mpitype) path = os.path.dirname(find_executable("mpirun")) super(Mpirun, self).__init__("/run/mpirun", "mpirun", job, path, subprocess) mca_default = None if mpitype == "openmpi": # Default mca values to avoid queue pair errors w/ OpenMPI mca_default = { "btl_openib_warn_default_gid_prefix": "0", "btl": "tcp,self", "oob": "tcp", "pml": "ob1", } self.hostfile = FormattedParameter("-hostfile {}", None) self.processes = FormattedParameter("-np {}", 1) self.ppn = FormattedParameter("-ppn {}", None) self.envlist = FormattedParameter("-envlist {}", None) self.mca = FormattedParameter("--mca {}", mca_default) self.working_dir = FormattedParameter("-wdir {}", None) self.mpitype = mpitype def assign_hosts(self, hosts, path=None, slots=None): """Assign the hosts to use with the command (-f). Args: hosts (list): list of hosts to specify in the hostfile path (str, optional): hostfile path. Defaults to None. slots (int, optional): number of slots per host to specify in the hostfile. Defaults to None. """ self._hosts = hosts kwargs = {"hostlist": self._hosts, "slots": slots} if path is not None: kwargs["path"] = path self.hostfile.value = write_host_file(**kwargs) def assign_processes(self, processes): """Assign the number of processes per node (-np). Args: processes (int): number of processes per node """ self.processes.value = processes def assign_environment(self, env_vars, append=False): """Assign or add environment variables to the command. Args: env_vars (EnvironmentVariables): the environment variables to use assign or add to the command append (bool): whether to assign (False) or append (True) the specified environment variables """ # Pass the environment variables via the process.run method env argument if append and self.env is not None: # Update the existing dictionary with the new values self.env.update(env_vars) else: # Overwrite/create the dictionary of environment variables self.env = EnvironmentVariables(env_vars) def assign_environment_default(self, env_vars): """Assign the default environment variables for the command. Args: env_vars (EnvironmentVariables): the environment variables to assign as the default """ self.envlist.update_default(env_vars.get_list()) def run(self): """Run the mpirun command. Raises: CommandFailure: if there is an error running the command """ load_mpi(self.mpitype) return super(Mpirun, self).run()
class Mpirun(JobManager): """A class for the mpirun job manager command.""" def __init__(self, job, subprocess=False, mpitype="openmpi"): """Create a Mpirun object. Args: job (ExecutableCommand): command object to manage. subprocess (bool, optional): whether the command is run as a subprocess. Defaults to False. """ load_mpi(mpitype) path = os.path.dirname(find_executable("mpirun")) super(Mpirun, self).__init__( "/run/mpirun", "mpirun", job, path, subprocess) self.hostfile = FormattedParameter("-hostfile {}", None) self.processes = FormattedParameter("-np {}", 1) self.ppn = FormattedParameter("-ppn {}", None) self.envlist = FormattedParameter("-envlist {}", None) self.mpitype = mpitype # deprecated: Use assign_[hosts|processes|environment]() methods instead def setup_command(self, env, hostfile, processes): """Set up the mpirun command with common inputs. Args: env (EnvironmentVariables): the environment variables to use with the launch command hostfile (str): file defining host names and slots processes (int): number of host processes """ # Setup the env for the job to export with the mpirun command self._pre_command = env.get_export_str() # Setup the orterun command self.hostfile.value = hostfile self.processes.value = processes def assign_hosts(self, hosts, path=None, slots=None): """Assign the hosts to use with the command (-f). Args: hosts (list): list of hosts to specify in the hostfile path (str, optional): hostfile path. Defaults to None. slots (int, optional): number of slots per host to specify in the hostfile. Defaults to None. """ kwargs = {"hostlist": hosts, "slots": slots} if path is not None: kwargs["path"] = path self.hostfile.value = write_host_file(**kwargs) def assign_processes(self, processes): """Assign the number of processes per node (-np). Args: processes (int): number of processes per node """ self.processes.value = processes def assign_environment(self, env_vars, append=False): """Assign or add environment variables to the command. Args: env_vars (EnvironmentVariables): the environment variables to use assign or add to the command append (bool): whether to assign (False) or append (True) the specified environment variables """ if append and self.envlist.value is not None: # Convert the current list of environmental variable assignments # into an EnvironmentVariables (dict) object. Then update the # dictionary keys with the specified values or add new key value # pairs to the dictionary. Finally convert the updated dictionary # back to a string for the parameter assignment. original = EnvironmentVariables({ item.split("=")[0]: item.split("=")[1] if "=" in item else None for item in self.envlist.value.split(",")}) original.update(env_vars) self.envlist.value = ",".join(original.get_list()) else: # Overwrite the environmental variable assignment self.envlist.value = ",".join(env_vars.get_list()) def assign_environment_default(self, env_vars): """Assign the default environment variables for the command. Args: env_vars (EnvironmentVariables): the environment variables to assign as the default """ self.envlist.update_default(env_vars.get_list()) def run(self): """Run the mpirun command. Raises: CommandFailure: if there is an error running the command """ load_mpi(self.mpitype) return super(Mpirun, self).run()