# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. import ctypes import os import sys import shlex import tarfile import time from datetime import datetime from subprocess import Popen import psutil from .log_utils import LogType, RemoteLogger, StdOutputType, nni_log from .commands import CommandType trial_output_path_name = ".nni" class Trial: def __init__(self, args, data): self.process = None self.data = data self.args = args self.command_channel = args.command_channel self.trial_syslogger_stdout = None global NNI_TRIAL_JOB_ID self.id = data["trialId"] if self.id is None: raise Exception("trial_id is not found in %s" % data) os.environ['NNI_TRIAL_JOB_ID'] = self.id NNI_TRIAL_JOB_ID = self.id # for multiple nodes. If it's None, it means single node. self.node_id = args.node_id if self.node_id is None: self.name = self.id else: self.name = "%s_%s" % (self.id, self.node_id) def run(self): # redirect trial's stdout and stderr to syslog self.trial_syslogger_stdout = RemoteLogger(self.args.nnimanager_ip, self.args.nnimanager_port, 'trial', StdOutputType.Stdout, self.args.log_collection, self.id, self.args.command_channel) nni_log(LogType.Info, "%s: start to run trial" % self.name) trial_working_dir = os.path.realpath(os.path.join(os.curdir, "..", "..", "trials", self.id)) self.trial_output_dir = os.path.join(trial_working_dir, trial_output_path_name) trial_code_dir = os.path.join(trial_working_dir, "code") trial_nnioutput_dir = os.path.join(trial_working_dir, "nnioutput") environ = os.environ.copy() environ['NNI_TRIAL_SEQ_ID'] = str(self.data["sequenceId"]) environ['NNI_OUTPUT_DIR'] = os.path.join(trial_working_dir, "nnioutput") environ['NNI_SYS_DIR'] = trial_working_dir self.working_dir = trial_working_dir # prepare code and parameters prepared_flag_file_name = os.path.join(trial_working_dir, "trial_prepared") if not os.path.exists(trial_working_dir): os.makedirs(trial_working_dir, exist_ok=True) os.makedirs(self.trial_output_dir, exist_ok=True) os.makedirs(trial_nnioutput_dir, exist_ok=True) # prepare code os.makedirs(trial_code_dir, exist_ok=True) with tarfile.open(os.path.join("..", "nni-code.tar.gz"), "r:gz") as tar: tar.extractall(trial_code_dir) # save parameters nni_log(LogType.Info, '%s: saving parameter %s' % (self.name, self.data["parameter"]["value"])) parameter_file_name = os.path.join(trial_working_dir, "parameter.cfg") with open(parameter_file_name, "w") as parameter_file: parameter_file.write(self.data["parameter"]["value"]) # ready flag with open(prepared_flag_file_name, "w") as prepared_flag_file: prepared_flag_file.write("%s" % (int(datetime.now().timestamp() * 1000))) # make sure code prepared by other node. if self.node_id is not None: while True: if os.path.exists(prepared_flag_file_name): break time.sleep(0.1) trial_command = self.args.trial_command gpuIndices = self.data.get("gpuIndices") if (gpuIndices is not None): if sys.platform == "win32": trial_command = 'set CUDA_VISIBLE_DEVICES="%s " && call %s' % (gpuIndices, trial_command) else: trial_command = 'CUDA_VISIBLE_DEVICES="%s " %s' % (gpuIndices, trial_command) self.log_pipe_stdout = self.trial_syslogger_stdout.get_pipelog_reader() self.process = Popen(trial_command, shell=True, stdout=self.log_pipe_stdout, stderr=self.log_pipe_stdout, cwd=trial_code_dir, env=dict(environ)) nni_log(LogType.Info, '{0}: spawns a subprocess (pid {1}) to run command: {2}'. format(self.name, self.process.pid, shlex.split(trial_command))) def save_parameter_file(self, command_data): parameters = command_data["parameters"] file_index = int(parameters["index"]) if file_index == 0: parameter_file_name = "parameter.cfg" else: parameter_file_name = "parameter_{}.cfg".format(file_index) parameter_file_name = os.path.join(self.working_dir, parameter_file_name) with open(parameter_file_name, "w") as parameter_file: nni_log(LogType.Info, '%s: saving parameter %s' % (self.name, parameters["value"])) parameter_file.write(parameters["value"]) def is_running(self): if (self.process is None): return False retCode = self.process.poll() # child worker process exits and all stdout data is read if retCode is not None and self.log_pipe_stdout.set_process_exit() and self.log_pipe_stdout.is_read_completed == True: # In Windows, the retCode -1 is 4294967295. It's larger than c_long, and raise OverflowError. # So covert it to int32. retCode = ctypes.c_long(retCode).value nni_log(LogType.Info, '{0}: subprocess terminated. Exit code is {1}.'.format(self.name, retCode)) end_time = int(datetime.now().timestamp() * 1000) end_message = { "code": retCode, "time": end_time, "trial": self.id, } self.command_channel.send(CommandType.TrialEnd, end_message) self.cleanup() return False else: return True def kill(self, trial_id=None): if trial_id == self.id or trial_id is None: if self.process is not None: try: nni_log(LogType.Info, "%s: killing trial" % self.name) for child in psutil.Process(self.process.pid).children(True): child.kill() self.process.kill() except psutil.NoSuchProcess: nni_log(LogType.Info, "kill trial %s failed: %s does not exist!" % (trial_id, self.process.pid)) except Exception as ex: nni_log(LogType.Error, "kill trial %s failed: %s " % (trial_id, str(ex))) self.cleanup() def cleanup(self): nni_log(LogType.Info, "%s: clean up trial" % self.name) self.process = None if self.log_pipe_stdout is not None: self.log_pipe_stdout.set_process_exit() self.log_pipe_stdout = None if self.trial_syslogger_stdout is not None: self.trial_syslogger_stdout.close() self.trial_syslogger_stdout = None