import logging
import os
import shutil
import subprocess
from tempfile import TemporaryDirectory
from sqlalchemy import literal
from sqlalchemy.sql import and_, select
from firex_keeper.db_model import firex_tasks
from firex_keeper.keeper_helper import FireXTreeTask
from firex_keeper.persist import (
get_db_file,
get_db_manager,
get_keeper_complete_file_path,
get_keeper_query_ready_file_path,
task_by_uuid_exp,
)
from firexapp.common import wait_until
from firexapp.events.model import (
FireXTask,
RunStates,
TaskColumn,
get_chain_exception_child_uuid,
is_chain_exception,
)
logger = logging.getLogger(__name__)
INCOMPLETE_RUNSTATES = {
s.to_celery_event_type() for s in RunStates if not s.is_complete()
}
REVOKED_RUNSTATES = {
s.to_celery_event_type() for s in [RunStates.REVOKED, RunStates.REVOKE_COMPLETED]
}
[docs]
class FireXTaskQueryException(Exception):
pass
def _task_col_eq(task_col, val):
return firex_tasks.c[task_col.value] == val
def _wait_and_query(
logs_dir, query, db_file_query_ready_timeout, **kwargs
) -> list[FireXTask]:
wait_on_keeper_query_ready(logs_dir, db_file_query_ready_timeout)
with get_db_manager(logs_dir) as db_manager:
return db_manager.query_tasks(query, **kwargs)
def _unlink_if_exists(path: str) -> None:
try:
if os.path.isfile(path):
os.unlink(path)
except OSError:
pass
def _rsync_keeper_db(rsync_bin: str, src: str, dest: str, cwd: str) -> bool:
"""Copy one keeper DB with rsync.
``src`` may be ``host:path``. rsync uses SSH only for that form.
``-e ssh`` is ignored for a plain local path.
"""
logger.info("Copying keeper DB for local query using rsync %s to %s", src, dest)
try:
subprocess.check_call(
[
rsync_bin,
"--whole-file",
"-e",
"ssh -o BatchMode=yes -o StrictHostKeyChecking=no",
src,
dest,
],
cwd=cwd,
)
except (OSError, subprocess.SubprocessError) as e:
logger.warning("rsync keeper copy failed: %s", e)
_unlink_if_exists(dest)
return False
if os.path.isfile(dest) and os.path.getsize(dest) > 0:
logger.info("Copied keeper DB for local query using rsync %s to %s", src, dest)
return True
logger.warning(
"rsync keeper copy succeeded but dest missing or empty: %s",
dest,
)
_unlink_if_exists(dest)
return False
def _copy_keeper_db_for_local_query(
existing_db_file: str,
new_tmp_db_file: str,
tmp_cwd: str,
rsync_src_host: str | None = None,
) -> None:
"""Copy keeper SQLite to a temp path for local reads.
``rsync_src_host`` is an SSH host that should read ``existing_db_file``.
The rsync source is ``host:existing_db_file`` when a host is given, and
the local database path otherwise. Then sqlite3 .backup, then shutil.copyfile.
"""
rsync_bin = shutil.which("rsync")
if rsync_bin:
src = (
f"{rsync_src_host}:{existing_db_file}"
if rsync_src_host
else existing_db_file
)
if _rsync_keeper_db(rsync_bin, src, new_tmp_db_file, tmp_cwd):
return
if rsync_src_host and _rsync_keeper_db(
rsync_bin, existing_db_file, new_tmp_db_file, tmp_cwd
):
return
sqlite_bin = "/bin/sqlite3"
if os.path.isfile(sqlite_bin):
logger.info(
"Copying keeper DB for local query using sqlite3 .backup from %s",
existing_db_file,
)
try:
subprocess.check_call(
[sqlite_bin, existing_db_file, f".backup {new_tmp_db_file}"],
cwd=tmp_cwd,
)
logger.info("Copied keeper DB for local query using sqlite3 .backup")
return
except (OSError, subprocess.SubprocessError) as e:
logger.warning("sqlite3 .backup keeper copy failed: %s", e)
_unlink_if_exists(new_tmp_db_file)
logger.info(
"Copying keeper DB for local query using shutil.copyfile from %s",
existing_db_file,
)
shutil.copyfile(existing_db_file, new_tmp_db_file)
logger.info("Copied keeper DB for local query using shutil.copyfile")
def _query_tasks(
logs_dir,
query,
db_file_query_ready_timeout=15,
copy_before_query=False,
rsync_src_host: str | None = None,
**kwargs,
) -> list[FireXTask]:
if copy_before_query:
tmp_base_dir = "/dev/shm" if os.path.isdir("/dev/shm") else None
with TemporaryDirectory(dir=tmp_base_dir) as temp_log_dir:
existing_db_file = get_db_file(logs_dir, new=False)
new_tmp_db_file = get_db_file(temp_log_dir, new=True)
_copy_keeper_db_for_local_query(
existing_db_file,
new_tmp_db_file,
temp_log_dir,
rsync_src_host=rsync_src_host,
)
query_results = _wait_and_query(
temp_log_dir, query, db_file_query_ready_timeout, **kwargs
)
else:
query_results = _wait_and_query(
logs_dir, query, db_file_query_ready_timeout, **kwargs
)
return query_results
def all_tasks(logs_dir, **kwargs) -> list[FireXTask]:
return _query_tasks(logs_dir, True, **kwargs)
def tasks_by_name(logs_dir, name, **kwargs) -> list[FireXTask]:
if "." in name:
col = TaskColumn.LONG_NAME
else:
col = TaskColumn.NAME
return _query_tasks(logs_dir, _task_col_eq(col, name), **kwargs)
def single_task_by_name(logs_dir, name, **kwargs) -> FireXTask:
tasks = _query_tasks(logs_dir, _task_col_eq(TaskColumn.NAME, name), **kwargs)
if len(tasks) != 1:
raise FireXTaskQueryException(
f"Required exactly one task named '{name}', found {len(tasks)}"
)
return tasks[0]
def task_by_uuid(
logs_dir, uuid, wait_for_exp_exist=None, max_wait=3, **kwargs
) -> FireXTask:
if wait_for_exp_exist is None:
wait_for_exp_exist = task_by_uuid_exp(uuid)
tasks = _query_tasks(
logs_dir,
_task_col_eq(TaskColumn.UUID, uuid),
wait_for_exp_exist=wait_for_exp_exist,
max_wait=max_wait,
**kwargs,
)
if not tasks:
raise FireXTaskQueryException(f"Found no task with UUID {uuid}")
return tasks[0]
def task_by_name_and_arg_pred(logs_dir, name, arg, pred) -> list[FireXTask]:
tasks_with_name = tasks_by_name(logs_dir, name)
return [
t
for t in tasks_with_name
if arg in t.firex_bound_args and pred(t.firex_bound_args[arg])
]
def task_by_name_and_arg_value(logs_dir, name, arg, value) -> list[FireXTask]:
pred = lambda arg_value: arg_value == value
return task_by_name_and_arg_pred(logs_dir, name, arg, pred)
def failed_tasks(logs_dir, **kwargs) -> list[FireXTask]:
return _query_tasks(
logs_dir, _task_col_eq(TaskColumn.STATE, RunStates.FAILED.value), **kwargs
)
def revoked_tasks(logs_dir, **kwargs) -> list[FireXTask]:
return _query_tasks(
logs_dir, firex_tasks.c[TaskColumn.STATE.value].in_(REVOKED_RUNSTATES), **kwargs
)
def running_tasks(logs_dir, **kwargs) -> list[FireXTask]:
return _query_tasks(
logs_dir,
firex_tasks.c[TaskColumn.STATE.value].in_(INCOMPLETE_RUNSTATES),
**kwargs,
)
def running_not_blocked_tasks(logs_dir, **kwargs) -> list[FireXTask]:
return _query_tasks(
logs_dir,
and_(
firex_tasks.c[TaskColumn.STATE.value].in_(INCOMPLETE_RUNSTATES),
firex_tasks.c[TaskColumn.STATE.value] != RunStates.BLOCKED.value,
),
**kwargs,
)
def failed_by_tasks(logs_dir, failed_uuid: str, **kwargs) -> list[FireXTask]:
# TODO: make this work with copy_before_query without copying twice,
# or tune page size to make NFS queries faster.
# assert is_failed(task_by_uuid(logs_dir, failed_uuid, **kwargs)), \
# f'Task {failed_uuid} did not fail.'
return _query_tasks(
logs_dir,
firex_tasks.c[TaskColumn.EXCEPTION_CAUSE_UUID.value] == failed_uuid,
**kwargs,
)
def _child_ids_by_parent_id(tasks_by_uuid):
child_uuids_by_parent_id = {u: [] for u in tasks_by_uuid}
for t in tasks_by_uuid.values():
# TODO: what if a child is entered in to the DB before its parent? Ignore for now.
if t.parent_id in child_uuids_by_parent_id:
child_uuids_by_parent_id[t.parent_id].append(t.uuid)
return child_uuids_by_parent_id
def _get_tree_tasks_by_uuid(root_uuid, tasks_by_uuid):
if root_uuid is None:
root_uuid = next(
(t.uuid for t in tasks_by_uuid.values() if t.parent_id is None), None
)
# FIXME: handle multiple roots?
if root_uuid is None:
raise FireXTaskQueryException("Found no root task with null parent_id.")
child_ids_by_parent_id = _child_ids_by_parent_id(tasks_by_uuid)
tree_tasks_by_uuid = {}
if root_uuid in tasks_by_uuid:
uuids_to_add = [root_uuid]
while uuids_to_add:
cur_task_uuid = uuids_to_add.pop()
cur_task = tasks_by_uuid[cur_task_uuid]
parent_tree_task = tree_tasks_by_uuid.get(cur_task.parent_id, None)
cur_tree_task = FireXTreeTask(
**{**cur_task._asdict(), "children": [], "parent": parent_tree_task}
)
if parent_tree_task:
parent_tree_task.children.append(cur_tree_task)
tree_tasks_by_uuid[cur_tree_task.uuid] = cur_tree_task
uuids_to_add += child_ids_by_parent_id[cur_tree_task.uuid]
return tree_tasks_by_uuid
def _create_task_tree(logs_dir, root_uuid=None, **kwargs) -> FireXTreeTask | None:
with get_db_manager(logs_dir) as db_manager:
if root_uuid is None:
root_uuid = db_manager.query_single_run_metadata().root_uuid
descendant_task_uuids = select(literal(root_uuid).label("uuid")).cte(
"descendant_task_uuids", recursive=True
)
descendant_uuids = select(firex_tasks.c.uuid).join(
descendant_task_uuids,
firex_tasks.c.parent_id == descendant_task_uuids.c.uuid,
)
# Approx SQL, uses index on parent_id.
# WITH RECURSIVE descendant_task_uuids(uuid) AS (
# SELECT :param_1 AS uuid
# UNION
# SELECT firex_tasks.uuid AS uuid
# FROM firex_tasks
# JOIN descendant_task_uuids ON firex_tasks.parent_id = descendant_task_uuids.uuid
# )
# SELECT firex_tasks.*
# FROM firex_tasks
# WHERE firex_tasks.uuid IN (
# SELECT uuid FROM descendant_task_uuids
# )
final_stmt = select(firex_tasks).where(
firex_tasks.c.uuid.in_(
select(descendant_task_uuids.union(descendant_uuids))
)
)
root_and_descendant_tasks = db_manager.query_tasks(final_stmt, **kwargs)
return _get_tree_tasks_by_uuid(
root_uuid,
{t.uuid: t for t in root_and_descendant_tasks},
).get(root_uuid) # root_uuid might not be in task table.
def task_tree_to_task(task_tree: FireXTreeTask) -> FireXTask:
task_tree_dict = task_tree._asdict()
task_tree_dict.pop("children")
task_tree_dict.pop("parent")
return FireXTask(**task_tree_dict)
def flatten_tree(task_tree: FireXTreeTask) -> list[FireXTreeTask]:
flat_tasks = []
to_check = [task_tree]
while to_check:
cur_task = to_check.pop()
to_check += cur_task.children
flat_tasks.append(cur_task)
return flat_tasks
def get_descendants(logs_dir, uuid) -> list[FireXTreeTask]:
# TODO: historically a FireXTreeTask was returned because
# the graph needed to be created within the application anyways.
# Now that sqlite recursive query is managing the tree,
# it's no longer necessary to create the graph in memory
# unless the caller actually needs it, which many currently don't
# (i.e. most call sites don't use "children" or "parent")
# Therefore an optimization is possible to query descendants
# but return a FireXTask, skipping application-side tree creation.
subtree = _create_task_tree(logs_dir, root_uuid=uuid)
# None if UUID isn't found.
if subtree is None:
return []
return [t for t in flatten_tree(subtree) if t.uuid != uuid]
def ancestor_by_long_name(
logs_dir, uuid, ancestor_long_name, **kwargs
) -> FireXTreeTask:
tasks_by_uuid = {t.uuid: t for t in all_tasks(logs_dir, **kwargs)}
# TODO: could avoid fetching all tasks by using sqlite recursive query.
tree_tasks_by_uuid = _get_tree_tasks_by_uuid(None, tasks_by_uuid)
if uuid in tree_tasks_by_uuid:
tree_task = tree_tasks_by_uuid[uuid]
while tree_task.parent and tree_task.parent.long_name != ancestor_long_name:
tree_task = tree_task.parent
if tree_task.parent.long_name == ancestor_long_name:
return tree_task.parent
return None
def find_task_causing_chain_exception(task: FireXTreeTask):
assert task.exception, f"Expected exception, received: {task.exception}"
if not is_chain_exception(task) or not task.children:
return task
causing_uuid = get_chain_exception_child_uuid(task)
causing_child = [c for c in task.children if c.uuid == causing_uuid]
# Note a chain interrupted exception can be caused by a non-descendant task via stitch_chains.
if not causing_child:
return task
causing_child = causing_child[0]
if not is_chain_exception(causing_child):
return causing_child
return find_task_causing_chain_exception(causing_child)
def wait_on_keeper_query_ready(logs_dir: str, timeout: int = 10):
if os.path.isfile(get_keeper_complete_file_path(logs_dir)):
return True
return wait_until(
os.path.isfile, timeout, 0.5, get_keeper_query_ready_file_path(logs_dir)
)
def wait_on_keeper_complete(logs_dir, timeout=30) -> bool:
return wait_until(
os.path.isfile, timeout, 1, get_keeper_complete_file_path(logs_dir)
)