Commit 8f74c5d

committed

gpu_count refactor

1 parent 684b9b0 commit 8f74c5dCopy full SHA for 8f74c5d

File tree

+11

-10

lines changed

+11

-10

lines changed

Lines changed: 2 additions & 10 deletions

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@`
`21`	`21`	`from ...utils.profiler import get_system_total_memory_gb`
`22`	`22`	`from ..engine import MapNode`
`23`	`23`	`from .base import DistributedPluginBase`
	`24`	`+from .tools import gpu_count`
`24`	`25`
`25`	`26`	`try:`
`26`	`27`	`from textwrap import indent`
`@@ -132,7 +133,7 @@ def __init__(self, plugin_args=None):`
`132`	`133`	`self.raise_insufficient = self.plugin_args.get("raise_insufficient", True)`
`133`	`134`
`134`	`135`	`# GPU found on system`
`135`		`- self.n_gpus_visible = MultiProcPlugin.gpu_count()`
	`136`	`+ self.n_gpus_visible = gpu_count()`
`136`	`137`	`# proc per GPU set by user`
`137`	`138`	`self.n_gpu_procs = self.plugin_args.get('n_gpu_procs', self.n_gpus_visible)`
`138`	`139`
`@@ -423,12 +424,3 @@ def _sort_jobs(self, jobids, scheduler="tsort"):`
`423`	`424`	`)`
`424`	`425`	`return jobids`
`425`	`426`
`426`		`- @staticmethod`
`427`		`- def gpu_count():`
`428`		`- n_gpus = 1`
`429`		`- try:`
`430`		`- import GPUtil`
`431`		`-`
`432`		`- return len(GPUtil.getGPUs())`
`433`		`- except ImportError:`
`434`		`- return n_gpus`

Lines changed: 9 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -178,3 +178,12 @@ def create_pyscript(node, updatehash=False, store_exception=True):`
`178`	`178`	`with open(pyscript, "w") as fp:`
`179`	`179`	`fp.writelines(cmdstr)`
`180`	`180`	`return pyscript`
	`181`	`+`
	`182`	`+def gpu_count():`
	`183`	`+ n_gpus = 1`
	`184`	`+ try:`
	`185`	`+ import GPUtil`
	`186`	`+ except ImportError:`
	`187`	`+ return 1`
	`188`	`+ else:`
	`189`	`+ return len(GPUtil.getGPUs())`

Comments

(0)