125 lines
3.6 KiB
Python
125 lines
3.6 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from argparse import ArgumentParser
|
|
|
|
from evalscope.cli.base import CLICommand
|
|
|
|
current_path = os.path.dirname(os.path.abspath(__file__))
|
|
print(current_path)
|
|
root_path = os.path.dirname(current_path)
|
|
print(root_path)
|
|
|
|
|
|
def subparser_func(args):
|
|
""" Function which will be called for a specific sub parser.
|
|
"""
|
|
return PerfServerCMD(args)
|
|
|
|
|
|
def add_perf_args(parser):
|
|
parser.add_argument('--server-command', required=True, type=str, help='The start server command.')
|
|
parser.add_argument(
|
|
'--logdir',
|
|
required=True,
|
|
type=str,
|
|
help='The monitor log save dir, tensorboard start at this path for display!')
|
|
parser.add_argument('--host', type=str, default='0.0.0.0', help='The tensorboard host')
|
|
parser.add_argument('--tensorboard-port', type=str, default='6006', help='The tensorboard port')
|
|
|
|
|
|
def async_run_command_with_popen(cmd):
|
|
sub_process = subprocess.Popen(
|
|
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1, universal_newlines=True, encoding='utf8')
|
|
return sub_process
|
|
|
|
|
|
def start_monitor(args):
|
|
cmd = ['python', '%s/perf/monitor.py' % root_path, '--logdir', args.logdir]
|
|
print(cmd)
|
|
p = async_run_command_with_popen(cmd)
|
|
os.set_blocking(p.stdout.fileno(), False)
|
|
return p
|
|
|
|
|
|
def start_tensorboard(args):
|
|
cmd = ['tensorboard', '--logdir', args.logdir, '--host', args.host, '--port', args.tensorboard_port]
|
|
p = async_run_command_with_popen(cmd)
|
|
os.set_blocking(p.stdout.fileno(), False)
|
|
return p
|
|
|
|
|
|
def start_server(args):
|
|
cmd = args.server_command
|
|
print(cmd)
|
|
sub_process = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
bufsize=1,
|
|
shell=True,
|
|
universal_newlines=True,
|
|
encoding='utf8')
|
|
|
|
os.set_blocking(sub_process.stdout.fileno(), False)
|
|
return sub_process
|
|
|
|
|
|
def wait_for_workers(workers):
|
|
while True:
|
|
for idx, worker in enumerate(workers):
|
|
if worker is None:
|
|
continue
|
|
# check worker is completed.
|
|
if worker.poll() is None:
|
|
for line in iter(worker.stdout.readline, ''):
|
|
if line != '':
|
|
sys.stdout.write(line)
|
|
else:
|
|
break
|
|
else:
|
|
print('Worker %s completed!' % idx)
|
|
for line in iter(worker.stdout.readline, ''):
|
|
if line != '':
|
|
sys.stdout.write(line)
|
|
else:
|
|
break
|
|
workers[idx] = None
|
|
|
|
is_all_completed = True
|
|
for idx, worker in enumerate(workers):
|
|
if worker is not None:
|
|
is_all_completed = False
|
|
break
|
|
|
|
if is_all_completed:
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
|
|
class PerfServerCMD(CLICommand):
|
|
name = 'server'
|
|
|
|
def __init__(self, args):
|
|
self.args = args
|
|
|
|
@staticmethod
|
|
def define_args(parsers: ArgumentParser):
|
|
""" define args for create pipeline template command.
|
|
"""
|
|
parser = parsers.add_parser(PerfServerCMD.name)
|
|
add_perf_args(parser)
|
|
parser.set_defaults(func=subparser_func)
|
|
|
|
def execute(self):
|
|
# start monitor
|
|
p_monitor = start_monitor(self.args)
|
|
# start tensorboard
|
|
p_tensorboard = start_tensorboard(self.args)
|
|
# start server
|
|
p_server = start_server(self.args)
|
|
|
|
wait_for_workers([p_monitor, p_tensorboard, p_server])
|