Format command template given available gpu ids.
(gpu_ids: List[int])
| 22 | |
| 23 | |
| 24 | def get_command_template(gpu_ids: List[int]) -> str: |
| 25 | """Format command template given available gpu ids.""" |
| 26 | if is_npu_available(): |
| 27 | tmpl = 'ASCEND_RT_VISIBLE_DEVICES=' + ','.join(str(i) for i in gpu_ids) |
| 28 | tmpl += ' {task_cmd}' |
| 29 | elif sys.platform == 'win32': # Always return win32 for Windows |
| 30 | # use command in Windows format |
| 31 | tmpl = 'set CUDA_VISIBLE_DEVICES=' + ','.join(str(i) for i in gpu_ids) |
| 32 | tmpl += ' & {task_cmd}' |
| 33 | else: |
| 34 | tmpl = 'CUDA_VISIBLE_DEVICES=' + ','.join(str(i) for i in gpu_ids) |
| 35 | tmpl += ' {task_cmd}' |
| 36 | return tmpl |
| 37 | |
| 38 | |
| 39 | @RUNNERS.register_module() |