"""
Experiment manager for the ALAB.
It is responsible for parsing the incoming experiment requests into many
tasks and samples and mark the finished tasks in the database when it is
done.
"""
import logging
import time
from contextlib import contextmanager
from typing import Any
from bson import ObjectId # type: ignore
from alab_management.utils.time import utc_now
from .config import AlabOSConfig
from .experiment_view import CompletedExperimentView, ExperimentStatus, ExperimentView
from .logger import DBLogger
from .sample_view import SampleView
from .task_view import TaskStatus, TaskView
from .utils.graph_ops import Graph
cli_logger = logging.getLogger(__name__)
[docs]
class ExperimentManager:
"""
Experiment manager read experiments from the experiment collection
and submit the experiment to executor and flag the completed experiments.
"""
def __init__(self):
self.experiment_view = ExperimentView()
self.task_view = TaskView()
self.sample_view = SampleView()
self.logger = DBLogger(task_id=None)
config = AlabOSConfig()
self.__copy_to_completed_db = (
"mongodb_completed" in config
) # if this is not defined in the config, assume it this feature is not being used.
if self.__copy_to_completed_db:
self.completed_experiment_view = CompletedExperimentView()
self._pause_handling_experiments = False
[docs]
def run(self):
"""Start the event loop."""
self.logger.system_log(
level="DEBUG",
log_data={
"logged_by": self.__class__.__name__,
"type": "ExperimentManagerStarted",
},
)
while True:
self._loop()
time.sleep(1)
def _loop(self):
if not self._pause_handling_experiments:
self.handle_pending_experiments()
self.mark_completed_experiments()
[docs]
def handle_pending_experiments(self):
"""
This method will scan the database to find out if there are
any pending experiments and submit it to task database.
"""
from alab_management.device_rpc_status import (
dismiss_lab_starting_user_input,
ensure_lab_starting_user_input,
get_lab_readiness,
)
readiness = get_lab_readiness()
if not readiness["lab_ready"]:
ensure_lab_starting_user_input()
cli_logger.info(
"Holding PENDING experiments until lab devices are ready "
"(%s).",
readiness.get("lab_ready_label"),
)
return
dismiss_lab_starting_user_input()
pending_experiments = self.experiment_view.get_experiments_with_status(
ExperimentStatus.PENDING
)
for experiment in pending_experiments:
started = self._handle_pending_experiment(experiment=experiment)
if not started:
continue
self.logger.system_log(
level="DEBUG",
log_data={
"logged_by": self.__class__.__name__,
"type": "ExperimentStarted",
"exp_id": experiment["_id"],
},
)
def _handle_pending_experiment(self, experiment: dict[str, Any]) -> bool:
"""Expand a PENDING experiment into samples/tasks.
Returns True if the experiment was claimed and left RUNNING. Returns False
if it was skipped (cycle / claim lost to cancel) or aborted because a
cancel landed while tasks were being created.
"""
samples: list[dict[str, Any]] = experiment["samples"]
tasks: list[dict[str, Any]] = experiment["tasks"]
exp_id = experiment["_id"]
# check if there is any cycle in the graph
reversed_edges = {i: task["prev_tasks"] for i, task in enumerate(tasks)}
task_graph = Graph(
list(range(len(tasks))),
# reverse reserved edges to get right directions of edges
{
i: [
j
for j, children in reversed_edges.items()
for child in children
if child == i
]
for i in list(range(len(tasks)))
},
)
if task_graph.has_cycle():
self.experiment_view.try_update_experiment_status(
exp_id,
ExperimentStatus.ERROR,
from_statuses=[ExperimentStatus.PENDING],
)
cli_logger.info(f"Experiment ({exp_id}) has a cycle in the graph.")
return False
# Claim before creating tasks so a cancel that already set CANCELLED is
# not overwritten by an unconditional RUNNING write at the end.
if not self.experiment_view.claim_pending_experiment(exp_id):
cli_logger.info(
f"Experiment ({exp_id}) is no longer PENDING; skipping start "
"(likely cancelled)."
)
return False
# create samples in the sample database (or reuse existing when flagged)
sample_ids: dict[str, ObjectId] = {}
for sample in samples:
existing_id = sample.get("sample_id", None)
if sample.get("reuse_existing") and existing_id is not None:
oid = (
existing_id
if isinstance(existing_id, ObjectId)
else ObjectId(existing_id)
)
if not self.sample_view.exists(sample_id=oid):
self.experiment_view.try_update_experiment_status(
exp_id,
ExperimentStatus.ERROR,
from_statuses=[ExperimentStatus.RUNNING],
)
cli_logger.info(
f"Experiment ({exp_id}) reuse_existing sample "
f"{oid} was not found."
)
return False
sample_ids[sample["name"]] = oid
else:
sample_ids[sample["name"]] = self.sample_view.create_sample(
sample["name"],
sample_id=existing_id,
tags=sample.get("tags", []),
metadata=sample.get("metadata", {}),
)
# create tasks in the task database
task_ids = []
for task in tasks:
task_samples = [
{"name": samplename, "sample_id": sample_ids[samplename]}
for samplename in task["samples"]
]
task_ids.append(
self.task_view.create_task(
task_type=task["type"],
parameters=task["parameters"],
samples=task_samples,
task_id=task.get("task_id", None),
)
)
# change the content of graph's vertices
task_graph.vertices = task_ids
# add dependency to each task
for task_id in task_ids:
self.task_view.update_task_dependency(
task_id,
next_tasks=task_graph.get_children(task_id),
prev_tasks=task_graph.get_parents(task_id),
)
self.task_view.try_to_mark_task_ready(task_id)
# write back the assign task & sample ids so Cancel can find them
self.experiment_view.update_sample_task_id(
exp_id=exp_id,
sample_ids=list(sample_ids.values()),
task_ids=task_ids,
)
# If Cancel landed after the claim (and possibly before task ids were
# written), stop the tasks we just created and leave CANCELLED alone.
current = self.experiment_view.get_experiment(exp_id)
if current is None or current.get("status") == ExperimentStatus.CANCELLED.name:
self._cancel_tasks_created_for_aborted_start(task_ids)
cli_logger.info(
f"Experiment ({exp_id}) was cancelled while starting; "
f"aborted {len(task_ids)} newly created task(s)."
)
return False
return True
def _cancel_tasks_created_for_aborted_start(self, task_ids: list[ObjectId]) -> None:
"""Force-cancel tasks created for an experiment that lost the cancel race."""
if not task_ids:
return
from alab_management.experiment_cancel import _force_cancel_experiment_tasks
_force_cancel_experiment_tasks(self.task_view, task_ids, utc_now())
[docs]
def mark_completed_experiments(self):
"""This method will scan the database to mark completed experiments in time."""
running_experiments = self.experiment_view.get_experiments_with_status(
ExperimentStatus.RUNNING
)
for experiment in running_experiments:
task_ids = [task["task_id"] for task in experiment["tasks"]]
# if all the tasks of an experiment have been finished
task_statuses = [
self.task_view.get_status(task_id=task_id) for task_id in task_ids
]
if all(
status
in {
TaskStatus.COMPLETED,
TaskStatus.ERROR,
TaskStatus.CANCELLED,
}
for status in task_statuses
):
finished_status = (
ExperimentStatus.CANCELLED
if any(status is TaskStatus.CANCELLED for status in task_statuses)
else ExperimentStatus.COMPLETED
)
self.experiment_view.update_experiment_status(
exp_id=experiment["_id"], status=finished_status
)
self.logger.system_log(
level="DEBUG",
log_data={
"logged_by": self.__class__.__name__,
"type": "ExperimentCompleted",
"exp_id": experiment["_id"],
},
)
cli_logger.info(f"Experiment ({experiment['_id']}) completed.")
if self.__copy_to_completed_db:
self.completed_experiment_view.save_experiment(experiment["_id"])
cli_logger.info(f"Experiment ({experiment['_id']}) and associated samples/tasks were copied to the completed db.")
self.logger.system_log(
level="DEBUG",
log_data={
"logged_by": self.__class__.__name__,
"type": "ExperimentSavedToCompletedDB",
"exp_id": experiment["_id"],
},
)
from alab_management.lab_reset import prune_archived_unplaced_from_live
prune_archived_unplaced_from_live()
[docs]
@contextmanager
def pause_handling_experiments(self):
"""This method will pause the handle pending experiments."""
try:
self._pause_handling_experiments = True
cli_logger.info("Pausing handling experiments.")
yield
finally:
self._pause_handling_experiments = False
cli_logger.info("Resuming handling experiments.")