diff --git a/gridmap/job.py b/gridmap/job.py index d0611d8..1b35e9f 100644 --- a/gridmap/job.py +++ b/gridmap/job.py @@ -116,12 +116,13 @@ class Job(object): 'cause_of_death', 'num_resubmits', 'home_address', 'log_stderr_fn', 'log_stdout_fn', 'timestamp', 'host_name', 'heart_beat', 'track_mem', 'track_cpu', 'interpreting_shell', - 'copy_env', 'par_env', 'gpu') + 'copy_env', 'par_env', 'gpu', 'h_vmem', 'h_rt', 'resources') def __init__(self, f, args, kwlist=None, cleanup=True, mem_free="1G", name='gridmap_job', num_slots=1, queue=DEFAULT_QUEUE, interpreting_shell=None, copy_env=True, add_env=None, - par_env=DEFAULT_PAR_ENV, gpu=0): + par_env=DEFAULT_PAR_ENV, gpu=0, h_vmem=None, h_rt=None, + resources=None): """ Initializes a new Job. @@ -154,6 +155,12 @@ def __init__(self, f, args, kwlist=None, cleanup=True, mem_free="1G", :type par_env: str :param gpu: number of GPUs to request :type gpu: int + :param h_vmem: hard virtual memory limit (e.g. "4G") + :type h_vmem: str, optional + :param h_rt: hard runtime limit (e.g. "00:59:00") + :type h_rt: str, optional + :param resources: list of additional custom resources specifications + :type resources: list of str, optional """ self.track_mem = [] self.track_cpu = [] @@ -201,6 +208,9 @@ def _add_env(env_vars): self.working_dir = os.getcwd() self.par_env = par_env self.gpu = gpu + self.h_vmem = h_vmem + self.h_rt = h_rt + self.resources = resources def execute(self): """ @@ -238,6 +248,12 @@ def native_specification(self): ret += " -q {}".format(self.queue) if self.gpu: ret += " -l gpu={}".format(self.gpu) + if self.h_vmem: + ret += " -l h_vmem={}".format(self.h_vmem) + if self.h_rt: + ret += " -l h_rt={}".format(self.h_rt) + if self.resources: + ret += "".join([" -l {}".format(x) for x in self.resources]) return ret @@ -899,7 +915,8 @@ def grid_map(f, args_list, cleanup=True, mem_free="1G", name='gridmap_job', num_slots=1, temp_dir=DEFAULT_TEMP_DIR, white_list=None, queue=DEFAULT_QUEUE, quiet=True, local=False, max_processes=1, interpreting_shell=None, copy_env=True, add_env=None, gpu=0, - completion_mail=False, require_cluster=False, par_env=DEFAULT_PAR_ENV): + h_vmem=None, h_rt=None, resources=None, completion_mail=False, + require_cluster=False, par_env=DEFAULT_PAR_ENV): """ Maps a function onto the cluster. @@ -950,6 +967,12 @@ def grid_map(f, args_list, cleanup=True, mem_free="1G", name='gridmap_job', :type add_env: dict :param gpu: number of GPUs to request :type gpu: int + :param h_vmem: hard virtual memory limit (e.g. "4G") + :type h_vmem: str, optional + :param h_rt: hard runtime limit (e.g. "00:59:00") + :type h_rt: str, optional + :param resources: list of additional custom resources specifications + :type resources: list of str, optional :param par_env: parallel environment to use. :type par_env: str :param completion_mail: whether to send an e-mail upon completion of all @@ -968,7 +991,7 @@ def grid_map(f, args_list, cleanup=True, mem_free="1G", name='gridmap_job', name='{}{}'.format(name, job_num), num_slots=num_slots, queue=queue, interpreting_shell=interpreting_shell, copy_env=copy_env, add_env=add_env, par_env=par_env, - gpu=gpu) + gpu=gpu, h_vmem=h_vmem, h_rt=h_rt, resources=resources) for job_num, args in enumerate(args_list)] # process jobs