def fake_run_manager(mocker, run=None, cloud=True, rm_class=wandb.run_manager.RunManager): # NOTE: This will create a run directory so make sure it's called in an isolated file system # We have an optional rm_class object because we mock it above so we need it before it's mocked api = InternalApi(load_settings=False) api.set_setting('project', 'testing') if wandb.run is None: wandb.run = run or Run() wandb.config = wandb.run.config wandb.run._api = api wandb.run._mkdir() wandb.run.socket = wandb_socket.Server() api.set_current_run_id(wandb.run.id) mocker.patch('wandb.apis.internal.FileStreamApi') api._file_stream_api = mocker.MagicMock() run_manager = rm_class(wandb.run, cloud=cloud, port=wandb.run.socket.port) class FakeProc(object): def poll(self): return None def exit(self, code=0): return None run_manager.proc = FakeProc() run_manager._meta = mocker.MagicMock() run_manager._stdout_tee = mocker.MagicMock() run_manager._stderr_tee = mocker.MagicMock() run_manager._output_log = mocker.MagicMock() run_manager._stdout_stream = mocker.MagicMock() run_manager._stderr_stream = mocker.MagicMock() run_manager.mirror_stdout_stderr = mocker.MagicMock() run_manager.unmirror_stdout_stderr = mocker.MagicMock() socket_thread = threading.Thread(target=wandb.run.socket.listen) socket_thread.start() run_manager._socket.ready() thread = threading.Thread(target=run_manager._sync_etc) thread.daemon = True thread.start() def test_shutdown(): if wandb.run and wandb.run.socket: wandb.run.socket.done() # TODO: is this needed? socket_thread.join() thread.join() run_manager.test_shutdown = test_shutdown run_manager._unblock_file_observer() run_manager._file_pusher._push_function = mocker.MagicMock() return run_manager
def _init_headless(run, cloud=True): global join global _user_process_finished_called environ = dict(os.environ) run.set_environment(environ) server = wandb_socket.Server() run.socket = server hooks = ExitHooks() hooks.hook() if sys.platform == "win32": # PTYs don't work in windows so we use pipes. stdout_master_fd, stdout_slave_fd = os.pipe() stderr_master_fd, stderr_slave_fd = os.pipe() else: stdout_master_fd, stdout_slave_fd = io_wrap.wandb_pty(resize=False) stderr_master_fd, stderr_slave_fd = io_wrap.wandb_pty(resize=False) headless_args = { 'command': 'headless', 'pid': os.getpid(), 'stdout_master_fd': stdout_master_fd, 'stderr_master_fd': stderr_master_fd, 'cloud': cloud, 'port': server.port } internal_cli_path = os.path.join( os.path.dirname(__file__), 'internal_cli.py') if six.PY2: # TODO(adrian): close_fds=False is bad for security. we set # it so we can pass the PTY FDs to the wandb process. We # should use subprocess32, which has pass_fds. popen_kwargs = {'close_fds': False} else: popen_kwargs = {'pass_fds': [stdout_master_fd, stderr_master_fd]} # TODO(adrian): ensure we use *exactly* the same python interpreter # TODO(adrian): make wandb the foreground process so we don't give # up terminal control until syncing is finished. # https://stackoverflow.com/questions/30476971/is-the-child-process-in-foreground-or-background-on-fork-in-c wandb_process = subprocess.Popen([sys.executable, internal_cli_path, json.dumps( headless_args)], env=environ, **popen_kwargs) termlog('Started W&B process version {} with PID {}'.format( __version__, wandb_process.pid)) os.close(stdout_master_fd) os.close(stderr_master_fd) # Listen on the socket waiting for the wandb process to be ready try: success, message = server.listen(30) except KeyboardInterrupt: success = False else: if not success: termerror('W&B process (PID {}) did not respond'.format( wandb_process.pid)) if not success: wandb_process.kill() for i in range(20): time.sleep(0.1) if wandb_process.poll() is not None: break if wandb_process.poll() is None: termerror('Failed to kill wandb process, PID {}'.format( wandb_process.pid)) # TODO attempt to upload a debug log path = GLOBAL_LOG_FNAME.replace(os.getcwd()+os.sep, "") raise LaunchError( "W&B process failed to launch, see: {}".format(path)) stdout_slave = os.fdopen(stdout_slave_fd, 'wb') stderr_slave = os.fdopen(stderr_slave_fd, 'wb') stdout_redirector = io_wrap.FileRedirector(sys.stdout, stdout_slave) stderr_redirector = io_wrap.FileRedirector(sys.stderr, stderr_slave) # TODO(adrian): we should register this right after starting the wandb process to # make sure we shut down the W&B process eg. if there's an exception in the code # above atexit.register(_user_process_finished, server, hooks, wandb_process, stdout_redirector, stderr_redirector) def _wandb_join(): _user_process_finished(server, hooks, wandb_process, stdout_redirector, stderr_redirector) join = _wandb_join _user_process_finished_called = False # redirect output last of all so we don't miss out on error messages stdout_redirector.redirect() if not env.is_debug(): stderr_redirector.redirect()
def _init_headless(run, cloud=True): global join global _user_process_finished_called program = util.get_program() if program: os.environ[env.PROGRAM] = os.getenv(env.PROGRAM) or program environ = dict(os.environ) run.set_environment(environ) server = wandb_socket.Server() run.socket = server hooks = ExitHooks() hooks.hook() if platform.system() == "Windows": try: import win32api # Make sure we are not ignoring CTRL_C_EVENT # https://docs.microsoft.com/en-us/windows/console/setconsolectrlhandler # https://stackoverflow.com/questions/1364173/stopping-python-using-ctrlc win32api.SetConsoleCtrlHandler(None, False) except ImportError: termerror( "Install the win32api library with `pip install pypiwin32`") # PTYs don't work in windows so we create these unused pipes and # mirror stdout to run.dir/output.log. There should be a way to make # pipes work, but I haven't figured it out. See links in compat/windows stdout_master_fd, stdout_slave_fd = os.pipe() stderr_master_fd, stderr_slave_fd = os.pipe() else: stdout_master_fd, stdout_slave_fd = io_wrap.wandb_pty(resize=False) stderr_master_fd, stderr_slave_fd = io_wrap.wandb_pty(resize=False) headless_args = { 'command': 'headless', 'pid': os.getpid(), 'stdout_master_fd': stdout_master_fd, 'stderr_master_fd': stderr_master_fd, 'cloud': cloud, 'port': server.port } internal_cli_path = os.path.join(os.path.dirname(__file__), 'internal_cli.py') if six.PY2 or platform.system() == "Windows": # TODO(adrian): close_fds=False is bad for security. we set # it so we can pass the PTY FDs to the wandb process. We # should use subprocess32, which has pass_fds. popen_kwargs = {'close_fds': False} else: popen_kwargs = {'pass_fds': [stdout_master_fd, stderr_master_fd]} # TODO(adrian): ensure we use *exactly* the same python interpreter # TODO(adrian): make wandb the foreground process so we don't give # up terminal control until syncing is finished. # https://stackoverflow.com/questions/30476971/is-the-child-process-in-foreground-or-background-on-fork-in-c wandb_process = subprocess.Popen( [sys.executable, internal_cli_path, json.dumps(headless_args)], env=environ, **popen_kwargs) termlog('Tracking run with wandb version {}'.format(__version__)) os.close(stdout_master_fd) os.close(stderr_master_fd) # Listen on the socket waiting for the wandb process to be ready try: success, _ = server.listen(30) except KeyboardInterrupt: success = False else: if not success: termerror('W&B process (PID {}) did not respond'.format( wandb_process.pid)) if not success: wandb_process.kill() for _ in range(20): time.sleep(0.1) if wandb_process.poll() is not None: break if wandb_process.poll() is None: termerror('Failed to kill wandb process, PID {}'.format( wandb_process.pid)) # TODO attempt to upload a debug log path = GLOBAL_LOG_FNAME.replace(os.getcwd() + os.sep, "") raise LaunchError("W&B process failed to launch, see: {}".format(path)) if platform.system() == "Windows": output = open(os.path.join(run.dir, "output.log"), "wb") stdout_redirector = io_wrap.WindowsRedirector(sys.stdout, output) stderr_redirector = io_wrap.WindowsRedirector(sys.stderr, output) else: stdout_slave = os.fdopen(stdout_slave_fd, 'wb') stderr_slave = os.fdopen(stderr_slave_fd, 'wb') try: stdout_redirector = io_wrap.FileRedirector(sys.stdout, stdout_slave) stderr_redirector = io_wrap.FileRedirector(sys.stderr, stderr_slave) except (ValueError, AttributeError): # stdout / err aren't files output = open(os.path.join(run.dir, "output.log"), "wb") stdout_redirector = io_wrap.WindowsRedirector(sys.stdout, output) stderr_redirector = io_wrap.WindowsRedirector(sys.stderr, output) # TODO(adrian): we should register this right after starting the wandb process to # make sure we shut down the W&B process eg. if there's an exception in the code # above atexit.register(_user_process_finished, server, hooks, wandb_process, stdout_redirector, stderr_redirector) def _wandb_join(exit_code=None): global _global_run_stack shutdown_async_log_thread() run.close_files() if exit_code is not None: hooks.exit_code = exit_code _user_process_finished(server, hooks, wandb_process, stdout_redirector, stderr_redirector) if len(_global_run_stack) > 0: _global_run_stack.pop() join = _wandb_join _user_process_finished_called = False # redirect output last of all so we don't miss out on error messages stdout_redirector.redirect() if not env.is_debug(): stderr_redirector.redirect()
def _init_headless(api, run, job_type, cloud=True): if 'WANDB_DESCRIPTION' in os.environ: run.description = os.environ['WANDB_DESCRIPTION'] env = dict(os.environ) run.set_environment(env) server = wandb_socket.Server() run.socket = server hooks = ExitHooks() hooks.hook() stdout_master_fd, stdout_slave_fd = pty.openpty() stderr_master_fd, stderr_slave_fd = pty.openpty() # raw mode so carriage returns etc. don't get added by the terminal driver tty.setraw(stdout_master_fd) tty.setraw(stderr_master_fd) headless_args = { 'command': 'headless', 'pid': os.getpid(), 'stdout_master_fd': stdout_master_fd, 'stderr_master_fd': stderr_master_fd, 'cloud': cloud, 'job_type': job_type, 'port': server.port } internal_cli_path = os.path.join(os.path.dirname(__file__), 'internal_cli.py') if six.PY2: # TODO(adrian): close_fds=False is bad for security. we set # it so we can pass the PTY FDs to the wandb process. We # should use subprocess32, which has pass_fds. popen_kwargs = {'close_fds': False} else: popen_kwargs = {'pass_fds': [stdout_master_fd, stderr_master_fd]} # TODO(adrian): make wandb the foreground process so we don't give # up terminal control until syncing is finished. # https://stackoverflow.com/questions/30476971/is-the-child-process-in-foreground-or-background-on-fork-in-c subprocess.Popen([ '/usr/bin/env', 'python', internal_cli_path, json.dumps(headless_args) ], env=env, **popen_kwargs) os.close(stdout_master_fd) os.close(stderr_master_fd) stdout_slave = os.fdopen(stdout_slave_fd, 'wb') stderr_slave = os.fdopen(stderr_slave_fd, 'wb') stdout_redirector = io_wrap.FileRedirector(sys.stdout, stdout_slave) stderr_redirector = io_wrap.FileRedirector(sys.stderr, stderr_slave) stdout_redirector.redirect() if os.environ.get('WANDB_DEBUG') != 'true': stderr_redirector.redirect() # Listen on the socket waiting for the wandb process to be ready success, message = server.listen(30) if not success: print('wandb Error: Failed to start') sys.exit(1) run.storage_id = message['storage_id'] def done(): server.done(hooks.exit_code) logger.info("Waiting for wandb process to finish") server.listen() atexit.register(done)
def _init_headless(run, job_type, cloud=True): run.description = env.get_description(run.description) environ = dict(os.environ) run.set_environment(environ) server = wandb_socket.Server() run.socket = server hooks = ExitHooks() hooks.hook() if sys.platform == "win32": # PTYs don't work in windows so we use pipes. stdout_master_fd, stdout_slave_fd = os.pipe() stderr_master_fd, stderr_slave_fd = os.pipe() else: stdout_master_fd, stdout_slave_fd = pty.openpty() stderr_master_fd, stderr_slave_fd = pty.openpty() # raw mode so carriage returns etc. don't get added by the terminal driver tty.setraw(stdout_master_fd) tty.setraw(stderr_master_fd) headless_args = { 'command': 'headless', 'pid': os.getpid(), 'stdout_master_fd': stdout_master_fd, 'stderr_master_fd': stderr_master_fd, 'cloud': cloud, 'job_type': job_type, 'port': server.port } internal_cli_path = os.path.join(os.path.dirname(__file__), 'internal_cli.py') if six.PY2: # TODO(adrian): close_fds=False is bad for security. we set # it so we can pass the PTY FDs to the wandb process. We # should use subprocess32, which has pass_fds. popen_kwargs = {'close_fds': False} else: popen_kwargs = {'pass_fds': [stdout_master_fd, stderr_master_fd]} # TODO(adrian): make wandb the foreground process so we don't give # up terminal control until syncing is finished. # https://stackoverflow.com/questions/30476971/is-the-child-process-in-foreground-or-background-on-fork-in-c wandb_process = subprocess.Popen([ '/usr/bin/env', 'python', internal_cli_path, json.dumps(headless_args) ], env=environ, **popen_kwargs) termlog('Started W&B process with PID {}'.format(wandb_process.pid)) os.close(stdout_master_fd) os.close(stderr_master_fd) # Listen on the socket waiting for the wandb process to be ready try: success, message = server.listen(30) except KeyboardInterrupt: success = False else: if not success: termerror('W&B process (PID {}) did not respond'.format( wandb_process.pid)) if not success: wandb_process.kill() for i in range(20): time.sleep(0.1) if wandb_process.poll() is not None: break if wandb_process.poll() is None: termerror('Failed to kill wandb process, PID {}'.format( wandb_process.pid)) sys.exit(1) stdout_slave = os.fdopen(stdout_slave_fd, 'wb') stderr_slave = os.fdopen(stderr_slave_fd, 'wb') stdout_redirector = io_wrap.FileRedirector(sys.stdout, stdout_slave) stderr_redirector = io_wrap.FileRedirector(sys.stderr, stderr_slave) # TODO(adrian): we should register this right after starting the wandb process to # make sure we shut down the W&B process eg. if there's an exception in the code # above atexit.register(_user_process_finished, server, hooks, wandb_process, stdout_redirector, stderr_redirector) # redirect output last of all so we don't miss out on error messages stdout_redirector.redirect() if not env.get_debug(): stderr_redirector.redirect()