Repository: Lightning-AI/pytorch-lightning
Pull request: Fix min_max gpu memory logging bug
## Before submitting

- Was this discussed/approved via a Github issue? Yes. #452 
- Did you read the [contributor guideline](https://github.com/williamFalcon/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)? Yes.
- Did you make sure to update the docs?  N/A
- Did you write any new necessary tests? Yes.

## What does this PR do?
Fixes #452 

Additionally, this PR updates & simplifies some code.
1. Use `subprocess.run` as officially recommended by python docs (refer to https://docs.python.org/3/library/subprocess.html#using-the-subprocess-module)
2. Clean up code around gpu memory function.
3. Increase test coverage by adding test for get_memory_profile.

## Did you have fun?
Yes :)


Problem statement:
min_max log_gpu_memory option bug
**Describe the bug**
Setting `log_gpu_memory='min_max'` in `Trainer` leads to the following bug.

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 347, in fit
    self.single_gpu_train(model)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/dp_mixin.py", line 79, in single_gpu_train
    self.run_pretrain_routine(model)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 467, in run_pretrain_routine
    self.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/train_loop_mixin.py", line 60, in train
    self.run_training_epoch()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/train_loop_mixin.py", line 126, in run_training_epoch
    self.log_metrics(batch_step_metrics, grad_norm_dic)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/logging_mixin.py", line 20, in log_metrics
    mem_map = memory.get_memory_profile(self.log_gpu_memory)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/root_module/memory.py", line 205, in get_memory_profile
    for k, v in memory_map:
ValueError: too many values to unpack (expected 2)
```


**To Reproduce**
On current master, execute the following.
```
    trainer = Trainer(
        ...
        log_gpu_memory='min_max',
        ...
    )
    trainer.fit(model)
```

**Expected behavior**
Log the min/max utilization of gpu memory, as `min_max` option is documented.

**Desktop (please complete the following information):**
 - OS: Ubuntu 18.04 
 - Version: Current master

I am working on this issue. Will submit a PR soon.

Hints from issue discussion:

Commit message: #452 Add test for get_memory_profile

Patch to review:
diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py
index d16d2659405cf..3b968bb3bfdeb 100644
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@@ -198,19 +198,10 @@ def get_memory_profile(mode):
     memory_map = get_gpu_memory_map()
 
     if mode == 'min_max':
-        min_mem = 1000000
-        min_k = None
-        max_mem = 0
-        max_k = None
-        for k, v in memory_map:
-            if v > max_mem:
-                max_mem = v
-                max_k = k
-            if v < min_mem:
-                min_mem = v
-                min_k = k
-
-        memory_map = {min_k: min_mem, max_k: max_mem}
+        min_index, min_memory = min(memory_map.items(), key=lambda item: item[1])
+        max_index, max_memory = max(memory_map.items(), key=lambda item: item[1])
+
+        memory_map = {min_index: min_memory, max_index: max_memory}
 
     return memory_map
 
@@ -224,17 +215,18 @@ def get_gpu_memory_map():
         Keys are device ids as integers.
         Values are memory usage as integers in MB.
     """
-    result = subprocess.check_output(
+    result = subprocess.run(
         [
-            'nvidia-smi', '--query-gpu=memory.used',
-            '--format=csv,nounits,noheader'
-        ], encoding='utf-8')
+            'nvidia-smi',
+            '--query-gpu=memory.used',
+            '--format=csv,nounits,noheader',
+        ],
+        encoding='utf-8',
+        capture_output=True,
+        check=True)
     # Convert lines into a dictionary
-    gpu_memory = [int(x) for x in result.strip().split('\n')]
-    gpu_memory_map = {}
-    for k, v in zip(range(len(gpu_memory)), gpu_memory):
-        k = f'gpu_{k}'
-        gpu_memory_map[k] = v
+    gpu_memory = [int(x) for x in result.stdout.strip().split('\n')]
+    gpu_memory_map = {f'gpu_{index}': memory for index, memory in enumerate(gpu_memory)}
     return gpu_memory_map
 
 
diff --git a/tests/test_gpu_models.py b/tests/test_gpu_models.py
index 118a3f2a521d3..17f8a8fd41b5d 100644
--- a/tests/test_gpu_models.py
+++ b/tests/test_gpu_models.py
@@ -224,7 +224,7 @@ def test_multi_gpu_model_dp():
     testing_utils.run_gpu_model_test(trainer_options, model, hparams)
 
     # test memory helper functions
-    memory.get_gpu_memory_map()
+    memory.get_memory_profile('min_max')
 
 
 def test_ddp_sampler_error():


Code review comments:
diff_hunk: @@ -224,17 +215,18 @@ def get_gpu_memory_map():
         Keys are device ids as integers.
         Values are memory usage as integers in MB.
     """
-    result = subprocess.check_output(
+    result = subprocess.run(
         [
-            'nvidia-smi', '--query-gpu=memory.used',
-            '--format=csv,nounits,noheader'
-        ], encoding='utf-8')
+            'nvidia-smi',
+            '--query-gpu=memory.used',
+            '--format=csv,nounits,noheader',
+        ],
+        encoding='utf-8',
+        capture_output=True,
+        check=True)
     # Convert lines into a dictionary
-    gpu_memory = [int(x) for x in result.strip().split('\n')]
-    gpu_memory_map = {}
-    for k, v in zip(range(len(gpu_memory)), gpu_memory):
-        k = f'gpu_{k}'
-        gpu_memory_map[k] = v
+    gpu_memory = [int(x) for x in result.stdout.strip().split('\n')]
path: pytorch_lightning/root_module/memory.py
text: @user1:
instead of `\n` consider using `os.sep`

@author:
Thanks for your feedback. I used `os.linesep` instead.