def record_kv_data(self, stream, kv_data, **kwargs):
"""Add key-value data to the stream."""
logging.info(f'Received new kvdata on stream {stream}')
U.f_mkdir(self.kvstream_folder)
with open(os.path.join(self.kvstream_folder, stream) + '.pkl',
'wb') as f:
d = dict(kv=kv_data, stream=stream, **kwargs)
pickle.dump(d, f)
def register_profile(self, offset, data, dst_dir_name, fname, done):
"""TODO: If multi threaded client, add filelock support here."""
del dst_dir_name # unused
U.f_mkdir(self.profile_folder)
self._stream_to_file(offset, data,
os.path.join(self.profile_folder, fname), done)
if done:
logging.info("Received new profile which is saved at %s/%s",
self.checkpoint_folder, fname)
def _save_profile(self, options, run_metadata):
tl = timeline.Timeline(run_metadata.step_stats)
ctf = tl.generate_chrome_trace_format()
export_path = os.path.join(TEMP_FOLDER, str(uuid.uuid4()))
U.f_mkdir(export_path)
with open(os.path.join(export_path, 'timeline.json'), 'w') as f:
f.write(ctf)
file_uploader = self._get_file_uploader()
file_uploader.send('register_profile',
src_fname=os.path.join(export_path,
'timeline.json'),
dst_fname='timeline.json',
dst_dir_name='') # dst_dir_name is unused.
U.f_remove(export_path)
def _register_configs(self, agent_config, env_config, sess_config,
**kwargs):
config_folder = self.config.configs_folder
U.f_mkdir(config_folder)
U.pretty_dump(agent_config,
os.path.join(config_folder, 'agent_config.json'))
U.pretty_dump(env_config, os.path.join(config_folder,
'env_config.json'))
U.pretty_dump(sess_config,
os.path.join(config_folder, 'sess_config.json'))
U.pretty_dump(kwargs, os.path.join(config_folder, 'misc_config.json'))
def register_checkpoint(self, offset, data, dst_dir_name, fname, done):
"""TODO: If multi threaded client, add filelock support here."""
U.f_mkdir(os.path.join(self.checkpoint_folder, dst_dir_name))
self._stream_to_file(
offset, data,
os.path.join(self.checkpoint_folder, dst_dir_name, fname), done)
if done:
with open(os.path.join(self.checkpoint_folder, 'info.txt'),
'a') as f:
print('{"dst_dir_name": "%s", "time":%d}' %
(dst_dir_name, int(time.time())),
file=f)
logging.info("Received new checkpoint which is saved at %s/%s/%s",
self.checkpoint_folder, dst_dir_name, fname)
self.enforce_checkpoint_policy()
def _create_ckpt(self):
# save the model
export_path = os.path.join(TEMP_FOLDER, str(uuid.uuid4()))
U.f_mkdir(export_path)
logging.info('Using %s folder for checkpointing ' % export_path)
with self._graph.as_default():
self._saver.save(self.sess,
export_path + '/learner',
global_step=self.global_step,
write_meta_graph=False)
file_uploader = self._get_file_uploader()
for fname in os.listdir(export_path):
if fname.endswith('.meta'):
continue
file_uploader.send('register_checkpoint',
src_fname=os.path.join(export_path, fname),
dst_fname=fname,
dst_dir_name='%d/' % self.global_step)
U.f_remove(export_path)
def _register_src(self):
src_folder = self.config.src_folder
U.f_mkdir(src_folder)
# this import requires fetching git executable which might not be available
# on all systems.
from git import Repo
repo = Repo('./')
commit = repo.head.commit
src = dict(
branch_name=repo.active_branch.name,
commit_summary=commit.summary,
commit_id=str(commit),
commit_datetime=commit.committed_datetime.strftime(
'%Y-%m-%d %H:%M:%S UTC'),
)
U.pretty_dump(src, os.path.join(src_folder, 'git_info.txt'))
with open(os.path.join(src_folder, 'git_diff.txt'), 'w') as f:
f.write(repo.git.diff(repo.head.commit.tree))
os.system('conda list > %s' %
os.path.join(src_folder, 'conda_env_list.txt'))
U.compress_tar('./liaison/', os.path.join(src_folder,
'liaison.tar.gz'))
def save_file(self, fname, data, **kwargs):
fname = f'{self.config.vis_files_folder}/{fname}'
U.f_mkdir(os.path.dirname(fname))
with open(fname, 'wb') as f:
f.write(data)
def register_metagraph(self, offset, data, _, fname, done):
"""TODO: If multi threaded client, add filelock support here."""
U.f_mkdir(self.checkpoint_folder)
self._stream_to_file(offset, data,
os.path.join(self.checkpoint_folder, fname), done)
def register_commands(self, **cmds):
U.f_mkdir(self.config.cmd_folder)
U.pretty_dump(cmds, os.path.join(self.config.cmd_folder, 'cmds.txt'))
def _register_cmd(self):
cmd_folder = self.config.cmd_folder
U.f_mkdir(cmd_folder)
with open(os.path.join(cmd_folder, 'cmd.txt'), 'w') as f:
f.write('\n'.join(sys.argv))