Allow watching existing task via cluster task

This commit is contained in:
Joshua Boniface 2023-11-16 03:06:13 -05:00
parent 289049d223
commit 63773a3061
2 changed files with 53 additions and 22 deletions

View File

@ -539,7 +539,15 @@ def cli_cluster_maintenance_off():
###############################################################################
@click.command(name="task", short_help="Show status of worker task.")
@connection_req
@click.argument("job_id", required=False, default=None)
@click.argument("task_id", required=False, default=None)
@click.option(
"--wait/--no-wait",
"wait_flag",
is_flag=True,
default=True,
show_default=True,
help="""Wait or don't wait for task to complete, showing progress. Requires TASK_ID; overrides "-f"/"--format".""",
)
@format_opt(
{
"pretty": cli_provisioner_status_format_pretty,
@ -550,12 +558,26 @@ def cli_cluster_maintenance_off():
"json-pretty": lambda d: jdumps(d, indent=2),
}
)
def cli_cluster_task(job_id, format_function):
def cli_cluster_task(task_id, wait_flag, format_function):
"""
Show the current status of worker task JOB_ID or a list of all active and pending tasks.
Show the current status of worker task TASK_ID or a list of all active and pending tasks.
"""
retcode, retdata = pvc.lib.common.task_status(CLI_CONFIG, job_id)
finish(retcode, retdata, format_function)
if task_id is None:
wait_flag = False
if wait_flag:
# First validate that this is actually a valid task that is running
retcode, retdata = pvc.lib.common.task_status(CLI_CONFIG, None)
if task_id in [i["id"] for i in retdata]:
retmsg = wait_for_celery_task(
CLI_CONFIG, {"task_id": task_id}, start_late=True
)
else:
retmsg = f"No task with ID {task_id} found."
finish(retcode, retmsg)
else:
retcode, retdata = pvc.lib.common.task_status(CLI_CONFIG, task_id)
finish(retcode, retdata, format_function)
###############################################################################

View File

@ -65,36 +65,45 @@ def cli_node_waiter(config, node, state_field, state_value):
echo(config, f" done. [{int(t_end - t_start)}s]")
def wait_for_celery_task(CLI_CONFIG, task_detail):
def wait_for_celery_task(CLI_CONFIG, task_detail, start_late=False):
"""
Wait for a Celery task to complete
"""
task_id = task_detail["task_id"]
run_on = task_detail["run_on"]
echo(CLI_CONFIG, f"Task ID: {task_id} assigned to node {run_on}")
echo(CLI_CONFIG, "")
if not start_late:
run_on = task_detail["run_on"]
# Wait for the task to start
echo(CLI_CONFIG, "Waiting for task to start...", newline=False)
while True:
sleep(0.5)
echo(CLI_CONFIG, f"Task ID: {task_id} assigned to node {run_on}")
echo(CLI_CONFIG, "")
# Wait for the task to start
echo(CLI_CONFIG, "Waiting for task to start...", newline=False)
while True:
sleep(0.5)
task_status = pvc.lib.common.task_status(
CLI_CONFIG, task_id=task_id, is_watching=True
)
if task_status.get("state") != "PENDING":
break
echo(CLI_CONFIG, ".", newline=False)
echo(CLI_CONFIG, " done.")
echo(CLI_CONFIG, "")
echo(
CLI_CONFIG,
task_status.get("status") + ":",
)
else:
task_status = pvc.lib.common.task_status(
CLI_CONFIG, task_id=task_id, is_watching=True
)
if task_status.get("state") != "PENDING":
break
echo(CLI_CONFIG, ".", newline=False)
echo(CLI_CONFIG, " done.")
echo(CLI_CONFIG, "")
echo(CLI_CONFIG, f"Watching existing task {task_id}:")
# Start following the task state, updating progress as we go
total_task = task_status.get("total")
echo(
CLI_CONFIG,
task_status.get("status") + ":",
)
with progressbar(length=total_task, show_eta=False) as bar:
last_task = 0
maxlen = 21