File size: 1,276 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
"""
Functionality of swapping tensors to/from (NVMe) storage devices.
"""
import subprocess
import shlex
class Job(object):
def __init__(self, cmd_line, output_file=None, work_dir=None):
self.cmd_line = cmd_line
self.output_file = output_file
self.work_dir = work_dir
self.output_fd = None
def cmd(self):
return self.cmd_line
def get_stdout(self):
return self.output_fd
def get_stderr(self):
return self.output_fd
def get_cwd(self):
return self.work_dir
def open_output_file(self):
if self.output_file is not None:
self.output_fd = open(self.output_file, 'w')
def close_output_file(self):
if self.output_fd is not None:
self.output_fd.close()
self.output_fd = None
def run_job(job, verbose=False):
args = shlex.split(' '.join(job.cmd()))
if verbose:
print(f'args = {args}')
job.open_output_file()
proc = subprocess.run(args=args, stdout=job.get_stdout(), stderr=job.get_stderr(), cwd=job.get_cwd())
job.close_output_file()
assert proc.returncode == 0, \
f"This command failed: {job.cmd()}"
|