def test_write_with_thread_safe_writer_multithread(self):
f1 = "test/thread_safe_writer/test_file_3.log"
f2 = "test/thread_safe_writer/test_file_4.log"
list_to_write = [i for i in range(10000)]
with open(f1, "w") as write_fp:
for data in list_to_write:
write_fp.write(str(data) + "\n")
file_writer = ThreadSafeWriter(f2, "w")
with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
futures = [
executor.submit(file_writer.write,
str(data) + "\n") for data in list_to_write
]
concurrent.futures.wait(futures)
propagate_exceptions(futures)
file_writer.close()
fp1 = open(f1, "r")
fp2 = open(f2, "r")
f1_lines = fp1.readlines()
f2_lines = fp2.readlines()
fp1.close()
fp2.close()
# since it is multi thread writing to the same file, the order is not guaranteed.
# hence we test the content equality by sorting and then comparing.
assert (not filecmp.cmp(f1, f2))
assert (f1_lines.sort() == f2_lines.sort())
os.remove(f1)
os.remove(f2)
def __init__(self, checkpoint_file):
"""
:param checkpoint_file: file to read / write object keys for checkpointing
"""
self._checkpoint_file = checkpoint_file
self._checkpoint_key_map = {}
self._checkpoint_file_append_fp = ThreadSafeWriter(
checkpoint_file, 'a')
self._restore_from_checkpoint_file()
def __init__(self, checkpoint_file):
"""
:param checkpoint_file: file to read / write object keys for checkpointing
"""
self._checkpoint_file = checkpoint_file
# By using ThreadSafeWriter, checkpointer is also thread-safe.
self._checkpoint_file_append_fp = ThreadSafeWriter(
checkpoint_file, 'a')
self._checkpoint_key_set = set()
self._restore_from_checkpoint_file()
class CheckpointKeySet(AbstractCheckpointKeySet):
"""Deals with checkpoint read and write."""
def __init__(self, checkpoint_file):
"""
:param checkpoint_file: file to read / write object keys for checkpointing
"""
self._checkpoint_file = checkpoint_file
# By using ThreadSafeWriter, checkpointer is also thread-safe.
self._checkpoint_file_append_fp = ThreadSafeWriter(
checkpoint_file, 'a')
self._checkpoint_key_set = set()
self._restore_from_checkpoint_file()
def write(self, key):
"""Writes key into checkpoint file. This also flushes data after write to prevent loss of checkpoint data on
system crash.
CAUTION: Make sure to persist your data before calling write. There is risk of data loss if your data is not persisted
and checkpointed on system crash.
"""
if key not in self._checkpoint_key_set:
self._checkpoint_file_append_fp.write(str(key) + "\n")
def contains(self, key):
"""Checks if key exists in the checkpoint set"""
exists = key in self._checkpoint_key_set
if exists:
logging.info(f"{key} found in checkpoint")
return exists
def _restore_from_checkpoint_file(self):
"""Reads all checkpoint keys from checkpoint_file into a set at initialization."""
if os.path.exists(self._checkpoint_file):
with open(self._checkpoint_file, 'r') as read_fp:
for key in read_fp:
self._checkpoint_key_set.add(key.rstrip('\n'))
def __del__(self):
self._checkpoint_file_append_fp.close()
def import_mlflow_experiments(self,
log_file='mlflow_experiments.log',
id_map_file='mlflow_experiments_id_map.log',
log_dir=None,
num_parallel=4):
mlflow_experiments_dir = log_dir if log_dir else self.export_dir
experiments_logfile = mlflow_experiments_dir + log_file
experiments_id_map_file = mlflow_experiments_dir + id_map_file
error_logger = logging_utils.get_error_logger(
wmconstants.WM_IMPORT, wmconstants.MLFLOW_EXPERIMENT_OBJECT,
self.export_dir)
mlflow_experiments_checkpointer = self._checkpoint_service.get_checkpoint_key_set(
wmconstants.WM_IMPORT, wmconstants.MLFLOW_EXPERIMENT_OBJECT)
start = timer()
id_map_thread_safe_writer = ThreadSafeWriter(experiments_id_map_file,
'a')
try:
with open(experiments_logfile, 'r') as fp:
with ThreadPoolExecutor(max_workers=num_parallel) as executor:
futures = [
executor.submit(self._create_experiment,
experiment_str,
id_map_thread_safe_writer,
mlflow_experiments_checkpointer,
error_logger) for experiment_str in fp
]
concurrent.futures.wait(futures,
return_when="FIRST_EXCEPTION")
propagate_exceptions(futures)
finally:
id_map_thread_safe_writer.close()
end = timer()
logging.info("Complete MLflow Experiments Import Time: " +
str(timedelta(seconds=end - start)))
def export_mlflow_experiments_acls(
self,
experiment_log='mlflow_experiments.log',
acl_log_file='mlflow_experiments_acls.log',
num_parallel=4):
"""
Export all experiments' permissions of already exported experiment objects logged in experiment_log file.
:return: writes the result to acl_log_file
"""
experiments_logfile = self.export_dir + experiment_log
acl_log_file_writer = ThreadSafeWriter(self.export_dir + acl_log_file,
'a')
error_logger = logging_utils.get_error_logger(
wmconstants.WM_EXPORT,
wmconstants.MLFLOW_EXPERIMENT_PERMISSION_OBJECT,
self.get_export_dir())
checkpoint_key_set = self._checkpoint_service.get_checkpoint_key_set(
wmconstants.WM_EXPORT,
wmconstants.MLFLOW_EXPERIMENT_PERMISSION_OBJECT)
start = timer()
try:
with open(experiments_logfile, 'r') as fp:
with ThreadPoolExecutor(max_workers=num_parallel) as executor:
futures = [
executor.submit(self._get_mlflow_experiment_acls,
acl_log_file_writer, experiment_str,
checkpoint_key_set, error_logger)
for experiment_str in fp
]
concurrent.futures.wait(futures,
return_when="FIRST_EXCEPTION")
propagate_exceptions(futures)
finally:
acl_log_file_writer.close()
end = timer()
logging.info("Complete MLflow Experiments Permissions Export Time: " +
str(timedelta(seconds=end - start)))
def log_all_workspace_acls(self,
workspace_log_file='user_workspace.log',
dir_log_file='user_dirs.log',
num_parallel=4):
"""
loop through all notebooks and directories to store their associated ACLs
:param workspace_log_file: input file for user notebook listing
:param dir_log_file: input file for user directory listing
"""
# define log file names for notebooks, folders, and libraries
logging.info("Exporting the notebook permissions")
start = timer()
acl_notebooks_error_logger = logging_utils.get_error_logger(
wmconstants.WM_EXPORT, wmconstants.WORKSPACE_NOTEBOOK_ACL_OBJECT,
self.get_export_dir())
acl_notebooks_writer = ThreadSafeWriter(
self.get_export_dir() + "acl_notebooks.log", "w")
try:
self.log_acl_to_file('notebooks', workspace_log_file,
acl_notebooks_writer,
acl_notebooks_error_logger, num_parallel)
finally:
acl_notebooks_writer.close()
end = timer()
logging.info("Complete Notebook ACLs Export Time: " +
str(timedelta(seconds=end - start)))
logging.info("Exporting the directories permissions")
start = timer()
acl_directory_error_logger = logging_utils.get_error_logger(
wmconstants.WM_EXPORT, wmconstants.WORKSPACE_DIRECTORY_ACL_OBJECT,
self.get_export_dir())
acl_directory_writer = ThreadSafeWriter(
self.get_export_dir() + "acl_directories.log", "w")
try:
self.log_acl_to_file('directories', dir_log_file,
acl_directory_writer,
acl_directory_error_logger, num_parallel)
finally:
acl_directory_writer.close()
end = timer()
logging.info("Complete Directories ACLs Export Time: " +
str(timedelta(seconds=end - start)))
def test_write_with_thread_safe_writer(self):
f1 = "test/thread_safe_writer/test_file_1.log"
f2 = "test/thread_safe_writer/test_file_2.log"
list_to_write = [i for i in range(1000)]
with open(f1, "w") as write_fp:
for data in list_to_write:
write_fp.write(str(data) + "\n")
file_writer = ThreadSafeWriter(f2, "w")
for data in list_to_write:
file_writer.write(str(data) + "\n")
file_writer.close()
assert (filecmp.cmp(f1, f2))
os.remove(f1)
os.remove(f2)
def export_top_level_folders(self):
ls_tld = self.get_top_level_folders()
logged_nb_count = 0
workspace_log_writer = ThreadSafeWriter(
self.get_export_dir() + 'user_workspace.log', "a")
libs_log_writer = ThreadSafeWriter(
self.get_export_dir() + 'libraries.log', "a")
dir_log_writer = ThreadSafeWriter(
self.get_export_dir() + 'user_dirs.log', "a")
checkpoint_item_log_set = self._checkpoint_service.get_checkpoint_key_set(
wmconstants.WM_EXPORT, wmconstants.WORKSPACE_ITEM_LOG_OBJECT)
try:
for tld_obj in ls_tld:
# obj has 3 keys, object_type, path, object_id
tld_path = tld_obj.get('path')
log_count = self.log_all_workspace_items(
tld_path, workspace_log_writer, libs_log_writer,
dir_log_writer, checkpoint_item_log_set)
logged_nb_count += log_count
finally:
workspace_log_writer.close()
libs_log_writer.close()
dir_log_writer.close()
dl_nb_count = self.download_notebooks()
print(f'Total logged notebooks: {logged_nb_count}')
print(f'Total Downloaded notebooks: {dl_nb_count}')
def log_all_workspace_items_entry(self,
ws_path='/',
workspace_log_file='user_workspace.log',
libs_log_file='libraries.log',
dir_log_file='user_dirs.log',
exclude_prefixes=[]):
logging.info(
f"Skip all paths with the following prefixes: {exclude_prefixes}")
workspace_log_writer = ThreadSafeWriter(
self.get_export_dir() + workspace_log_file, "a")
libs_log_writer = ThreadSafeWriter(
self.get_export_dir() + libs_log_file, "a")
dir_log_writer = ThreadSafeWriter(self.get_export_dir() + dir_log_file,
"a")
checkpoint_item_log_set = self._checkpoint_service.get_checkpoint_key_set(
wmconstants.WM_EXPORT, wmconstants.WORKSPACE_ITEM_LOG_OBJECT)
try:
num_nbs = self.log_all_workspace_items(
ws_path=ws_path,
workspace_log_writer=workspace_log_writer,
libs_log_writer=libs_log_writer,
dir_log_writer=dir_log_writer,
checkpoint_set=checkpoint_item_log_set,
exclude_prefixes=exclude_prefixes)
finally:
workspace_log_writer.close()
libs_log_writer.close()
dir_log_writer.close()
return num_nbs
def export_user_home(self, username, local_export_dir, num_parallel=4):
"""
Export the provided user's home directory
:param username: user's home directory to export
:param local_export_dir: folder location to do single user exports
:return: None
"""
original_export_dir = self.get_export_dir()
user_export_dir = self.get_export_dir() + local_export_dir
user_root = '/Users/' + username.rstrip().lstrip()
self.set_export_dir(user_export_dir + '/{0}/'.format(username))
print("Export path: {0}".format(self.get_export_dir()))
os.makedirs(self.get_export_dir(), exist_ok=True)
workspace_log_writer = ThreadSafeWriter(
self.get_export_dir() + 'user_workspace.log', "a")
libs_log_writer = ThreadSafeWriter(
self.get_export_dir() + 'libraries.log', "a")
dir_log_writer = ThreadSafeWriter(
self.get_export_dir() + 'user_dirs.log', "a")
checkpoint_item_log_set = self._checkpoint_service.get_checkpoint_key_set(
wmconstants.WM_EXPORT, wmconstants.WORKSPACE_ITEM_LOG_OBJECT)
try:
num_of_nbs = self.log_all_workspace_items(user_root,
workspace_log_writer,
libs_log_writer,
dir_log_writer,
checkpoint_item_log_set)
finally:
workspace_log_writer.close()
libs_log_writer.close()
dir_log_writer.close()
if num_of_nbs == 0:
raise ValueError(
'User does not have any notebooks in this path. Please verify the case of the email'
)
num_of_nbs_dl = self.download_notebooks(ws_dir='user_artifacts/')
print(f"Total notebooks logged: {num_of_nbs}")
print(f"Total notebooks downloaded: {num_of_nbs_dl}")
if num_of_nbs != num_of_nbs_dl:
print(
f"Notebooks logged != downloaded. Check the failed download file at: {user_export_dir}"
)
print(f"Exporting the notebook permissions for {username}")
acl_notebooks_writer = ThreadSafeWriter("acl_notebooks.log", "w")
acl_notebooks_error_logger = logging_utils.get_error_logger(
wmconstants.WM_EXPORT, wmconstants.WORKSPACE_NOTEBOOK_ACL_OBJECT,
self.get_export_dir())
try:
self.log_acl_to_file('notebooks', 'user_workspace.log',
acl_notebooks_writer,
acl_notebooks_error_logger, num_parallel)
finally:
acl_notebooks_writer.close()
print(f"Exporting the directories permissions for {username}")
acl_directories_writer = ThreadSafeWriter("acl_directories.log", "w")
acl_directories_error_logger = logging_utils.get_error_logger(
wmconstants.WM_EXPORT, wmconstants.WORKSPACE_DIRECTORY_ACL_OBJECT,
self.get_export_dir())
try:
self.log_acl_to_file('directories', 'user_dirs.log',
acl_directories_writer,
acl_directories_error_logger, num_parallel)
finally:
acl_directories_writer.close()
# reset the original export dir for other calls to this method using the same client
self.set_export_dir(original_export_dir)
class CheckpointKeyMap(AbstractCheckpointKeyMap):
"""Deals with checkpoint read and write. Unlike CheckpointKeySet, it also saves the value.
Useful when the corresponding value is needed.
"""
def __init__(self, checkpoint_file):
"""
:param checkpoint_file: file to read / write object keys for checkpointing
"""
self._checkpoint_file = checkpoint_file
self._checkpoint_key_map = {}
self._checkpoint_file_append_fp = ThreadSafeWriter(
checkpoint_file, 'a')
self._restore_from_checkpoint_file()
def write(self, key, value):
if key not in self._checkpoint_key_map or "IN_USE_BY" in self._checkpoint_key_map[
key]:
self._checkpoint_key_map[key] = value
self._checkpoint_file_append_fp.write(
json.dumps({
"key": str(key),
"value": str(value)
}) + "\n")
def check_contains_otherwise_mark_in_use(self, key):
"""
If the key_map does not have the key value yet, mark the key to be IN_USE_BY_$THREAD_ID, and return False.
If the key_map has the key value,
if the value is "IN_USE_BY_XXX" wait for the result to be ready and return True (self.contains(key))
if the value is not "IN_USE_BY_XXX" return True (self.contains(key))
This method is thread-safe since map.setdefault(key, value) is thread-safe
The thread that calls this method and gets False (meaning, the value wasn't there and thus this thread is using
this key) must set the value of this key subsequently to make sure other threads do not wait for this key
forever.
"""
in_use_str = f"IN_USE_BY_{threading.get_ident()}"
# setdefault is thread safe, so only one thread can successfully set the value for the key.
result = self._checkpoint_key_map.setdefault(key, in_use_str)
if result == in_use_str:
return False
else:
while key in self._checkpoint_key_map and "IN_USE_BY" in self._checkpoint_key_map[
key]:
logging.info(f"Waiting for {key} result to be available..")
time.sleep(5)
return self.contains(key)
def contains(self, key):
exists = key in self._checkpoint_key_map
if exists:
logging.info(f"{key} found in checkpoint")
return exists
def remove(self, key):
if key in self._checkpoint_key_map:
self._checkpoint_key_map.pop(key)
def get(self, key):
return self._checkpoint_key_map[key]
def get_file_path(self):
return self._checkpoint_file
def _restore_from_checkpoint_file(self):
if os.path.exists(self._checkpoint_file):
with open(self._checkpoint_file, 'r') as read_fp:
for single_key_value_map_str in read_fp:
single_key_value_map = json.loads(single_key_value_map_str)
self._checkpoint_key_map[single_key_value_map[
"key"]] = single_key_value_map["value"]
def __del__(self):
self._checkpoint_file_append_fp.close()