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()
