From f389c9c3f06ec2e3154aa1df41d1191fe04d482a Mon Sep 17 00:00:00 2001
From: zmy <izmy@qq.com>
Date: Tue, 22 Dec 2020 19:47:34 +0800
Subject: [PATCH] a little improvments and update readme

---
 README.md               | 43 +++++++++++++++++++++++++++++++++++++++--
 gpuutil/gpuutil.py      |  6 +++---
 gpuutil/set_redirect.py |  4 ++--
 3 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 46cec07..7f913c6 100644
--- a/README.md
+++ b/README.md
@@ -74,6 +74,45 @@ def auto_set(num, allow_nonfree=True, ask=True, blacklist=[], show=True):
 	# some code here.
 ```
 
+# Use this inside an docker.
+For some reason, codes that running in docker cannot get the correct information about the process that using the gpu. 
+To support that, gpuutil supports read the output command of nvidia-smi and ps from an given file, which should be generated by you from host machine
+To use this in docker, try the following steps:
+1. figure out a way to pass the output of command ```nvidia-smi -q -x``` to the docker that your are currently using, save the output as a text file.
+2. pass the output of a ps-like command to the docker. It is a table-like output, the first line is header, which should at least contains user, pid and command. below is an valid output generated by running ```ps -axo user,pid,command```on host machine:
+```
+USER  PID COMMAND
+root    1 /bin/bash -c bash /etc/init.docker; /usr/sbin/sshd -D
+root    8 sshd: /usr/sbin/sshd -D [listener] 0 of 10-100 startups
+root    9 sshd: user1 [priv]
+user1  19 sshd: user1@pts/0
+user1  20 -zsh
+user1  97 tmux
+user1  98 -zsh
+```
+if your generated output have different name, for example when you are using ```docker top``` instead of ```ps```, the ```COMMAND``` section would be ```CMD```, therefore you need prepare a dict that maps its name to either of ```user, pid, command```, note that its insensitive to upper case.
+3. run the configuration script.
+```shell
+python -m gpuutil.set_redirect -nv path/to/your/nvidia/output -ps /path/to/your/ps/output -pst cmd=command,username=user
+```
+for more information about the script, run ```python -m gpuutil.set_redirect -h```, you will get:
+```
+usage: set_redirect.py [-h] [--nvsmi NVSMI] [--ps PS] [--ps_name_trans PS_NAME_TRANS]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --nvsmi NVSMI, -nv NVSMI
+                        a file indicates real nvidia-smi -q -x output.
+  --ps PS, -ps PS       a file indicates real ps-like output.
+  --ps_name_trans PS_NAME_TRANS, -pst PS_NAME_TRANS
+                        a dict of name trans, format: name1=buildin,name2=buildin, buildin can be choosen from cmd,user,pid
+```
+> some advice:
+> 1. you can use a script that run nvidia-smi and ps command and save their output to a directory, the mount the directory to the docker as readonly.
+> 2. you could consider mount the directory as tmpfs.
+
 ## ps:
-1. you can get more detailed gpu info via accessing gpuutil.GPUStat class, for more information, just look the code.
-2. Since it use ps command to get detailed process info, it can only be used on linux.
+1. You can get more detailed gpu info via accessing gpuutil.GPUStat class, for more information, just look the code.
+2. Since it use ps command to get detailed process info, it can only be used on linux, if you use it on windows, some information might be missing.
+3. If you have any trouble, feel free to open an issue.
+4. The code is straight forward, it's also a good choice to take an look at the code if you got any trouble.
diff --git a/gpuutil/gpuutil.py b/gpuutil/gpuutil.py
index 48ef71f..c5a9980 100644
--- a/gpuutil/gpuutil.py
+++ b/gpuutil/gpuutil.py
@@ -183,10 +183,10 @@ def get_basic_process_info_by_file(filepath, col_name_trans=None):
     interested = {
         'user': None,
         'pid': None,
-        'cmd': None
+        'command': None
     }
     if col_name_trans is None:
-        col_name_trans = {'command': 'cmd'}
+        col_name_trans = {'cmd': 'command'}
     for i, word in enumerate(header):
         word = word.lower()
         if word in col_name_trans:
@@ -198,7 +198,7 @@ def get_basic_process_info_by_file(filepath, col_name_trans=None):
         words = line.split(' ')
         pid = words[interested['pid']]
         user = words[interested['user']]
-        cmd = ' '.join(words[interested['cmd']:])
+        cmd = ' '.join(words[interested['command']:])
         processes[pid] = {
             "user": user,
             "command": cmd
diff --git a/gpuutil/set_redirect.py b/gpuutil/set_redirect.py
index 451c95f..84af30b 100644
--- a/gpuutil/set_redirect.py
+++ b/gpuutil/set_redirect.py
@@ -2,7 +2,7 @@ import argparse
 import os
 import json
 
-availabel_name_trans = ['cmd', 'user', 'pid']
+availabel_name_trans = ['command', 'user', 'pid']
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--nvsmi', '-nv', default=None, type=str, help='a file indicates real nvidia-smi -q -x output.')
@@ -25,7 +25,7 @@ if name_trans is not None:
         if len(item) != 2:
             raise ValueError('there must be a = in nametrans')
         key, value = item
-        if value not in avaliable_name_trans:
+        if value not in availabel_name_trans:
             raise ValueError('given buildin name {0} do not exist, avaliable: {1}'.format(value, ','.join(availabel_name_trans)))
         parsed_name_trans[key] = value