🐛 fix an bug when parsing from file

README.md

@@ -6,7 +6,7 @@ A naive tool for observing gpu status and auto set visible gpu in python code.
 
 1. install the package.
 ```shell
-pip install https://git.zmy.pub/zmyme/gpuutil/archive/v0.0.3.tar.gz
+pip install https://git.zmy.pub/zmyme/gpuutil/archive/v0.0.4.tar.gz
 ```
 
 2. for observing gpu status, just input
@@ -74,6 +74,46 @@ 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.
+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.
+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.

gpuutil/__main__.py

@@ -1,6 +1,5 @@
-from gpuutil import GPUStat
+from gpuutil import GPUStat, loaddict, savedict
 import sys
-import json
 import argparse
 import os
 
@@ -18,13 +17,11 @@ def load_config():
     configpath = os.path.join(home_dir, '.gpuutil.conf')
     if not os.path.isfile(configpath):
         return {}
-    with open(configpath, 'r', encoding='utf-8') as f:
+    return loaddict(configpath)
-        return json.load(f)
 def save_config(config):
     home_dir = os.path.expanduser('~')
     configdir = os.path.join(home_dir, '.gpuutil.conf')
-    with open(configdir, 'w+', encoding='utf-8') as f:
+    savedict(configdir, config)
-        json.dump(config, f, ensure_ascii=False, indent=4)
 
 # style format: |c|l:15|r|c:14rl:13|
 def parse_style(style):

gpuutil/gpuutil.py

@@ -10,6 +10,31 @@ import platform
 
 osname = platform.system()
 
+def loadfile(path):
+    with open(path, 'r', encoding='utf-8') as f:
+        return f.read()
+def savefile(path, content):
+    with open(path, 'w+', encoding='utf-8') as f:
+        return f.write(content)
+def loaddict(path):
+    content = loadfile(path)
+    content = content.strip()
+    if len(content) != 0:
+        return json.loads(content)
+    else:
+        return {}
+def savedict(path, dictionary):
+    content = json.dumps(dictionary, indent=4, ensure_ascii=False)
+    savefile(path, content)
+def clean_split(line, delemeter=' '):
+    words = line.split(delemeter)
+    words = [w.strip() for w in words]
+    words = [w for w in words if w != '']
+    return words
+
+def exe_cmd(command):
+    pipe = os.popen(command)
+    return pipe.read()
 
 def xml2dict(node):
     node_dict = {}
@@ -25,10 +50,8 @@ def xml2dict(node):
             node_dict[child.tag].append(xml2dict(child))
     return node_dict
 
-def parse_nvsmi_info(command='nvidia-smi -q -x'):
+def parse_nvsmi_info(nvsmixml):
-    pipe = os.popen(command)
+    tree = ET.fromstring(nvsmixml)
-    xml = pipe.read()
-    tree = ET.fromstring(xml)
     return xml2dict(tree)
 
 def parse_gpu_info(stat):
@@ -140,7 +163,7 @@ def get_basic_process_info_linux():
     lines = output.split('\n')[1:]
     processes = {}
     for line in lines:
-        words = [p for p in line.split(' ') if p != '']
+        words = clean_split(line)
         if len(words) < 3:
             continue
         username = words[0]
@@ -168,6 +191,37 @@ def get_basic_process_info_windows():
         }
     return processes
 
+def get_basic_process_info_by_file(filepath, col_name_trans=None):
+    # suppose cmd is always at the last, and the previous lines have no space.
+    content = loadfile(filepath)
+    lines = content.split('\n')
+    header = clean_split(lines[0])
+    interested = {
+        'user': None,
+        'pid': None,
+        'command': None
+    }
+    if col_name_trans is None:
+        col_name_trans = {'cmd': 'command'}
+    for i, word in enumerate(header):
+        word = word.lower()
+        if word in col_name_trans:
+            word = col_name_trans[word]
+        if word in interested:
+            interested[word] = i
+    processes = {}
+    for line in lines[1:]:
+        words = clean_split(line)
+        pid = words[interested['pid']]
+        user = words[interested['user']]
+        cmd = ' '.join(words[interested['command']:])
+        processes[pid] = {
+            "user": user,
+            "command": cmd
+        }
+    return processes
+
+
 def draw_table(table, rowsty=None, colsty=None, colsz = None):
     def justify(s, align, width):
         if align == 'c':
@@ -267,13 +321,35 @@ class GPUStat():
         self.cuda_version = ''
         self.attached_gpus = ''
         self.driver_version = ''
+        self.nvsmi_source = None
+        self.ps_source = None
+        self.ps_name_trans = None
+        self.load_configure()
+    def load_configure(self):
+        configuration_path = os.path.expanduser('~/.gpuutil.conf')
+        if os.path.isfile(configuration_path):
+            configuration = loaddict(configuration_path)
+            if 'redirect' in configuration:
+                if 'nvsmi_src' in configuration['redirect']:
+                    self.nvsmi_source = configuration['redirect']['nvsmi_src']
+                if 'ps_src' in configuration['redirect']:
+                    self.ps_source = configuration['redirect']['ps_src']
+                if 'ps_name_trans' in configuration['redirect']:
+                    self.ps_name_trans = configuration['redirect']['ps_name_trans']
+
+            
     def get_process_info(self):
+        if self.ps_source is not None:
+            return get_basic_process_info_by_file(self.ps_source, self.ps_name_trans)
         if osname == 'Windows':
             return get_basic_process_info_windows()
         elif osname == 'Linux':
             return get_basic_process_info_linux()
     def parse(self):
-        self.raw_info = parse_nvsmi_info('nvidia-smi -q -x')
+        if self.nvsmi_source is None:
+            self.raw_info = parse_nvsmi_info(exe_cmd('nvidia-smi -q -x'))
+        else:
+            self.raw_info = parse_nvsmi_info(loadfile(self.nvsmi_source))
         self.detailed_info = {}
         for key, value in self.raw_info.items():
             if key != 'gpu':

gpuutil/set_redirect.py

@@ -0,0 +1,42 @@
+import argparse
+import os
+from gpuutil import loaddict, savedict
+
+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.')
+parser.add_argument('--ps', '-ps', default=None, type=str, help='a file indicates real ps-like output.')
+parser.add_argument('--ps_name_trans', '-pst', default=None, type=str, help='a dict of name trans, \
+                                                                            format: name1=buildin,name2=buildin, \
+                                                                            buildin can be choosen from {0}'.format(','.join(availabel_name_trans)))
+
+args = parser.parse_args()
+
+# lets chech the pst.
+parsed_name_trans = {}
+name_trans = args.ps_name_trans
+if name_trans is not None:
+    name_trans = name_trans.split(',')
+    name_trans = [t.strip() for t in name_trans]
+    name_trans = [t for t in name_trans if t!='']
+    for item in name_trans:
+        item = item.split('=', maxsplit=1)
+        if len(item) != 2:
+            raise ValueError('there must be a = in nametrans')
+        key, value = item
+        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
+
+config_file = os.path.expanduser('~/.gpuutil.conf')
+configuration = {}
+if os.path.isfile(config_file):
+    configuration = loaddict(config_file)
+configuration['redirect'] = {
+    "nvsmi_src": args.nvsmi,
+    "ps_src": args.ps,
+    "ps_name_trans": parsed_name_trans
+}
+
+savedict(config_file, configuration)

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name = 'gpuutil',
-    version = '0.0.3',
+    version = '0.0.4',
     keywords='gpu utils',
     description = 'A tool for observing gpu stat and auto set visible gpu in python code.',
     license = 'MIT License',