Specify Resources in App Decorator

parsl/app/app.py

@@ -21,7 +21,7 @@ class AppBase(metaclass=ABCMeta):
 
     """
 
-    def __init__(self, func, data_flow_kernel=None, walltime=60, executors='all', cache=False):
+    def __init__(self, func, data_flow_kernel=None, walltime=60, executors='all', cache=False, cores=None, disk=None, mem=None):
         """Construct the App object.
 
         Args:
@@ -44,6 +44,9 @@ def __init__(self, func, data_flow_kernel=None, walltime=60, executors='all', ca
         self.data_flow_kernel = data_flow_kernel
         self.status = 'created'
         self.executors = executors
+        self.cores = cores
+        self.disk = disk
+        self.mem = mem
         self.cache = cache
         if not (isinstance(executors, list) or isinstance(executors, str)):
             logger.error("App {} specifies invalid executor option, expects string or list".format(
@@ -77,7 +80,7 @@ def __call__(self, *args, **kwargs):
         pass
 
 
-def App(apptype, data_flow_kernel=None, walltime=60, cache=False, executors='all'):
+def App(apptype, data_flow_kernel=None, walltime=60, cache=False, cores=None, mem=None, disk=None, executors='all'):
     """The App decorator function.
 
     Args:
@@ -114,11 +117,14 @@ def wrapper(f):
                          data_flow_kernel=data_flow_kernel,
                          walltime=walltime,
                          cache=cache,
+                         cores=cores,
+                         mem=mem,
+                         disk=disk,
                          executors=executors)
     return wrapper
 
 
-def python_app(function=None, data_flow_kernel=None, walltime=60, cache=False, executors='all'):
+def python_app(function=None, data_flow_kernel=None, walltime=60, cache=False, cores=None, mem=None, disk=None, executors='all'):
     """Decorator function for making python apps.
 
     Parameters
@@ -146,14 +152,17 @@ def wrapper(f):
                              data_flow_kernel=data_flow_kernel,
                              walltime=walltime,
                              cache=cache,
+                             cores=cores,
+                             mem=mem,
+                             disk=disk,
                              executors=executors)
         return wrapper(func)
     if function is not None:
         return decorator(function)
     return decorator
 
 
-def bash_app(function=None, data_flow_kernel=None, walltime=60, cache=False, executors='all'):
+def bash_app(function=None, data_flow_kernel=None, walltime=60, cache=False, cores=None, mem=None, disk=None, executors='all'):
     """Decorator function for making bash apps.
 
     Parameters
@@ -181,6 +190,9 @@ def wrapper(f):
                            data_flow_kernel=data_flow_kernel,
                            walltime=walltime,
                            cache=cache,
+                           cores=cores,
+                           mem=mem,
+                           disk=disk,
                            executors=executors)
         return wrapper(func)
     if function is not None:

parsl/app/bash.py

@@ -119,8 +119,8 @@ def open_std_fd(fdname):
 
 class BashApp(AppBase):
 
-    def __init__(self, func, data_flow_kernel=None, walltime=60, cache=False, executors='all'):
-        super().__init__(func, data_flow_kernel=data_flow_kernel, walltime=60, executors=executors, cache=cache)
+    def __init__(self, func, data_flow_kernel=None, walltime=60, cache=False, cores=None, mem=None, disk=None, executors='all'):
+        super().__init__(func, data_flow_kernel=data_flow_kernel, walltime=60, executors=executors, cores=cores, mem=mem, disk=disk, cache=cache)
         self.kwargs = {}
 
         # We duplicate the extraction of parameter defaults
@@ -147,6 +147,12 @@ def __call__(self, *args, **kwargs):
         """
         # Update kwargs in the app definition with ones passed in at calltime
         self.kwargs.update(kwargs)
+        if 'parsl_resource_specification' not in self.kwargs:
+            self.kwargs['parsl_resource_specification'] = {
+                'cores': self.cores,
+                'disk': self.disk,
+                'mem': self.mem
+            }
 
         if self.data_flow_kernel is None:
             dfk = DataFlowKernelLoader.dfk()

parsl/app/python.py

@@ -35,12 +35,15 @@ def inject_exception(thread):
 class PythonApp(AppBase):
     """Extends AppBase to cover the Python App."""
 
-    def __init__(self, func, data_flow_kernel=None, walltime=60, cache=False, executors='all'):
+    def __init__(self, func, data_flow_kernel=None, walltime=60, cache=False, cores=None, disk=None, mem=None, executors='all'):
         super().__init__(
             wrap_error(func),
             data_flow_kernel=data_flow_kernel,
             walltime=walltime,
             executors=executors,
+            cores=cores,
+            disk=disk,
+            mem=mem,
             cache=cache
         )
 
@@ -62,6 +65,13 @@ def __call__(self, *args, **kwargs):
         else:
             dfk = self.data_flow_kernel
 
+        if 'parsl_resource_specification' not in kwargs:
+            kwargs['parsl_resource_specification'] = {
+                'cores': self.cores,
+                'mem': self.mem,
+                'disk': self.disk
+            }
+
         walltime = self.kwargs.get('walltime')
         if walltime is not None:
             self.func = timeout(self.func, walltime)

parsl/executors/workqueue/executor.py

@@ -120,6 +120,7 @@ def WorkQueueSubmitThread(task_queue=multiprocessing.Queue(),
             input_files = item["input_files"]
             output_files = item["output_files"]
             std_files = item["std_files"]
+            resource_specs = item["resources"]
 
             full_script_name = workqueue_worker.__file__
             script_name = full_script_name.split("/")[-1]
@@ -182,6 +183,17 @@ def WorkQueueSubmitThread(task_queue=multiprocessing.Queue(),
             for item in std_files:
                 t.specify_file(item[0], item[1], WORK_QUEUE_OUTPUT, cache=item[2])
 
+            # Specify resource requirements for the task
+            if "cores" in resource_specs and resource_specs["cores"] is not None:
+                t.specify_cores(resource_specs["cores"])
+                logger.debug("Task {} using {} cores".format(t.id, resource_specs["cores"]))
+            if "disk" in resource_specs and resource_specs["disk"] is not None:
+                t.specify_disk(resource_specs["disk"])
+                logger.debug("Task {} using {} MB disk".format(t.id, resource_specs["disk"]))
+            if "mem" in resource_specs and resource_specs["mem"] is not None:
+                t.specify_memory(resource_specs["mem"])
+                logger.debug("Task {} using {} MB memory".format(t.id, resource_specs["mem"]))
+
             # Submit the task to the WorkQueue object
             logger.debug("Submitting task {} to WorkQueue".format(parsl_id))
             try:
@@ -548,6 +560,12 @@ def submit(self, func, *args, **kwargs):
         input_files = []
         output_files = []
         std_files = []
+        resources = {}
+
+        # Receive resource specifications from the kwargs, then remove it
+        if 'parsl_resource_specification' in kwargs:
+            resources = kwargs["parsl_resource_specification"]
+            kwargs.pop("parsl_resource_specification")
 
         # Add input files from the "inputs" keyword argument
         func_inputs = kwargs.get("inputs", [])
@@ -631,6 +649,7 @@ def submit(self, func, *args, **kwargs):
         # Create message to put into the message queue
         logger.debug("Placing task {} on message queue".format(task_id))
         msg = {"task_id": task_id,
+               "resources": resources,
                "data_loc": function_data_file,
                "result_loc": function_result_file,
                "input_files": input_files,