From df71e4cc855ec0f8e08ee692a421ce1e1effa394 Mon Sep 17 00:00:00 2001
From: Anna Woodard <annawoodard@uchicago.edu>
Date: Fri, 23 Aug 2019 13:14:20 -0500
Subject: [PATCH] Support cpu and mem specification via env (#1207)

* Support cpu and mem specification via env

This commit adds support for respecting specification of how much CPU
and memory should be used by workers via the PARSL_CORES and
PARSL_GB_MEMORY environmental variables (as opposed to inferring it
to be the full node, once a job has started on the node.) After this commit, the config
below will result in 4 workers being started, using a total of 4 cores
and 12 GB memory.

This is a pre-requisite to a full solution to #942.

```
from parsl.config import Config
from parsl.providers import SlurmProvider
from parsl.addresses import address_by_hostname
from parsl.executors import HighThroughputExecutor

config = Config(
    executors=[
        HighThroughputExecutor(
            cores_per_worker=1,
            mem_per_worker=3,
            address=address_by_hostname(),
            provider=SlurmProvider(
                'broadwl',
                cmd_timeout=60,
                nodes_per_block=1,
                init_blocks=1,
                min_blocks=1,
                max_blocks=1,
                scheduler_options='#SBATCH --cpus-per-task=4 --mem-per-cpu=3g,
                worker_init='worker_init='export PARSL_MEMORY_GB=$(expr $SLURM_CPUS_ON_NODE \* $SLURM_MEM_PER_CPU); export PARSL_CORES=$SLURM_CPUS_ON_NODE',
                exclusive=False
            ),
        )
    ],
)
```

* Change PARSL_GB_MEMORY -> PARSL_MEMORY_GB

* Revert changes to launchers

This reversts the launcher changes introduced in
18061d2b5197dd6a94e0bc553fb12d4941d9faf7.
---
 .../executors/high_throughput/process_worker_pool.py  | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/parsl/executors/high_throughput/process_worker_pool.py b/parsl/executors/high_throughput/process_worker_pool.py
index 7249e90a37..f4bbf0088d 100755
--- a/parsl/executors/high_throughput/process_worker_pool.py
+++ b/parsl/executors/high_throughput/process_worker_pool.py
@@ -123,8 +123,15 @@ def __init__(self,
         self.uid = uid
         self.block_id = block_id
 
-        cores_on_node = multiprocessing.cpu_count()
-        available_mem_on_node = round(psutil.virtual_memory().available / (2**30), 1)
+        if os.environ.get('PARSL_CORES'):
+            cores_on_node = int(os.environ['PARSL_CORES'])
+        else:
+            cores_on_node = multiprocessing.cpu_count()
+
+        if os.environ.get('PARSL_MEMORY_GB'):
+            available_mem_on_node = float(os.environ['PARSL_MEMORY_GB'])
+        else:
+            available_mem_on_node = round(psutil.virtual_memory().available / (2**30), 1)
 
         self.max_workers = max_workers
         self.prefetch_capacity = prefetch_capacity