def abort_environment_build(environment_build_uuid, is_running=False):
"""Aborts an environment build.
Aborts an environment build by setting its state to ABORTED and
sending a REVOKE and ABORT command to celery.
Args:
is_running:
environment_build_uuid: uuid of the environment build to abort
Returns:
"""
filter_by = {
"build_uuid": environment_build_uuid,
}
status_update = {"status": "ABORTED"}
celery_app = make_celery(current_app)
# Make use of both constructs (revoke, abort) so we cover both a
# task that is pending and a task which is running.
celery_app.control.revoke(environment_build_uuid, timeout=1.0)
if is_running:
res = AbortableAsyncResult(environment_build_uuid, app=celery_app)
# It is responsibility of the task to terminate by reading it's
# aborted status.
res.abort()
update_status_db(
status_update,
model=models.EnvironmentBuild,
filter_by=filter_by,
)
def _collateral(self, task_id: str):
celery = make_celery(current_app)
celery.send_task(
"app.core.tasks.build_jupyter",
task_id=task_id,
)
def process_notification_deliveries(app) -> None:
with app.app_context():
app.logger.debug("Sending process notifications deliveries task.")
celery = make_celery(app)
res = celery.send_task(
name="app.core.tasks.process_notifications_deliveries")
res.forget()
def _set_celery_worker_parallelism_at_runtime(worker: str,
current_parallelism: int,
new_parallelism: int) -> bool:
"""Set the parallelism of a celery worker at runtime.
Args:
worker: Name of the worker.
current_parallelism: Current parallelism level.
new_parallelism: New parallelism level.
Returns:
True if the parallelism level could be changed, False otherwise.
Only allows to increase parallelism, the reason is that celery
won't gracefully decrease the parallelism level if it's not
possible because processes are busy with a task.
"""
if current_parallelism is None or new_parallelism is None:
return False
if new_parallelism < current_parallelism:
return False
if current_parallelism == new_parallelism:
return True
# We don't query the celery-worker and rely on arguments because the
# worker might take some time to spawn new processes, leading to
# race conditions.
celery = make_celery(current_app)
worker = f"celery@{worker}"
celery.control.pool_grow(new_parallelism - current_parallelism, [worker])
return True
def delete(self, experiment_uuid):
"""Stops an experiment given its UUID.
However, it will not delete any corresponding database entries,
it will update the status of corresponding objects to "REVOKED".
"""
experiment = models.Experiment.query.get_or_404(
experiment_uuid,
description="Experiment not found",
)
run_uuids = [run.run_uuid for run in experiment.pipeline_runs]
# Revokes all pipeline runs and waits for a reply for 1.0s.
celery = make_celery(current_app)
celery.control.revoke(run_uuids, timeout=1.0)
# Update the status of the run and step entries to "REVOKED".
models.NonInteractiveRun.query.filter_by(
experiment_uuid=experiment_uuid).update({"status": "REVOKED"})
models.NonInteractiveRunPipelineStep.query.filter_by(
experiment_uuid=experiment_uuid).update({"status": "REVOKED"})
db.session.commit()
return {"message": "Experiment termination was successful"}, 200
def post(self):
"""Starts a new (interactive) pipeline run."""
post_data = request.get_json()
post_data["run_config"]["run_endpoint"] = "runs"
pipeline = construct_pipeline(**post_data)
# Create Celery object with the Flask context and construct the
# kwargs for the job.
celery = make_celery(current_app)
celery_job_kwargs = {
"pipeline_description": pipeline.to_dict(),
"project_uuid": post_data["project_uuid"],
"run_config": post_data["run_config"],
}
# Start the run as a background task on Celery. Due to circular
# imports we send the task by name instead of importing the
# function directly.
res = celery.send_task("app.core.tasks.run_partial",
kwargs=celery_job_kwargs)
# NOTE: this is only if a backend is configured. The task does
# not return anything. Therefore we can forget its result and
# make sure that the Celery backend releases recourses (for
# storing and transmitting results) associated to the task.
# Uncomment the line below if applicable.
# res.forget()
# NOTE: we are setting the status of the run ourselves without
# using the option of celery to get the status of tasks. This
# way we do not have to configure a backend (where the default
# of "rpc://" does not give the results we would want).
run = {
"run_uuid": res.id,
"pipeline_uuid": pipeline.properties["uuid"],
"project_uuid": post_data["project_uuid"],
"status": "PENDING",
}
db.session.add(models.InteractiveRun(**run))
# Set an initial value for the status of the pipline steps that
# will be run.
step_uuids = [s.properties["uuid"] for s in pipeline.steps]
pipeline_steps = []
for step_uuid in step_uuids:
pipeline_steps.append(
models.InteractiveRunPipelineStep(
**{
"run_uuid": res.id,
"step_uuid": step_uuid,
"status": "PENDING"
}))
db.session.bulk_save_objects(pipeline_steps)
db.session.commit()
run["pipeline_steps"] = pipeline_steps
return run, 201
def post(self):
"""Start a new run."""
post_data = request.get_json()
# Construct pipeline.
pipeline = construct_pipeline(**post_data)
# Create Celery object with the Flask context and construct the
# kwargs for the job.
celery = make_celery(current_app)
celery_job_kwargs = {
'pipeline_description': pipeline.to_dict(),
'run_config': post_data['run_config']
}
# Start the run as a background task on Celery. Due to circular
# imports we send the task by name instead of importing the
# function directly.
res = celery.send_task('app.core.runners.run_partial',
kwargs=celery_job_kwargs)
# NOTE: this is only if a backend is configured. The task does
# not return anything. Therefore we can forget its result and
# make sure that the Celery backend releases recourses (for
# storing and transmitting results) associated to the task.
# Uncomment the line below if applicable.
# res.forget()
# NOTE: we are setting the status of the run ourselves without
# using the option of celery to get the status of tasks. This
# way we do not have to configure a backend (where the default
# of "rpc://" does not give the results we would want).
run = {
'run_uuid': res.id,
'pipeline_uuid': pipeline.properties['uuid'],
'status': 'PENDING',
}
db.session.add(models.Run(**run))
# Set an initial value for the status of the pipline steps that
# will be run.
step_uuids = [s.properties['uuid'] for s in pipeline.steps]
step_statuses = []
for step_uuid in step_uuids:
step_statuses.append(
models.StepStatus(
**{
'run_uuid': res.id,
'step_uuid': step_uuid,
'status': 'PENDING'
}))
db.session.bulk_save_objects(step_statuses)
db.session.commit()
run['step_statuses'] = step_statuses
return run, 201
def _collateral(self, run_uuids: List[str]):
# Aborts and revokes all pipeline runs and waits for a reply for
# 1.0s.
celery = make_celery(current_app)
celery.control.revoke(run_uuids, timeout=1.0)
for run_uuid in run_uuids:
res = AbortableAsyncResult(run_uuid, app=celery)
# It is responsibility of the task to terminate by reading
# its aborted status.
res.abort()
def _collateral(
self,
project_uuid: str,
task_id: str,
pipeline: Pipeline,
run_config: Dict[str, Any],
env_variables: Dict[str, Any],
**kwargs,
):
# Get docker ids of images to use and make it so that the images
# will not be deleted in case they become outdated by an
# environment rebuild.
try:
env_uuid_docker_id_mappings = lock_environment_images_for_run(
task_id,
project_uuid,
pipeline.get_environments(),
)
except errors.ImageNotFound as e:
msg = (
"Pipeline references environments that do not exist in the"
f" project, the following environments do not exist: [{e}].\n\n"
"Please make sure all pipeline steps are assigned an"
" environment that exists in the project."
)
raise errors.ImageNotFound(msg)
# Create Celery object with the Flask context and construct the
# kwargs for the job.
celery = make_celery(current_app)
run_config["env_uuid_docker_id_mappings"] = env_uuid_docker_id_mappings
run_config["user_env_variables"] = env_variables
celery_job_kwargs = {
"pipeline_definition": pipeline.to_dict(),
"project_uuid": project_uuid,
"run_config": run_config,
}
# Start the run as a background task on Celery. Due to circular
# imports we send the task by name instead of importing the
# function directly.
res = celery.send_task(
"app.core.tasks.run_pipeline",
kwargs=celery_job_kwargs,
task_id=task_id,
)
# NOTE: this is only if a backend is configured. The task does
# not return anything. Therefore we can forget its result and
# make sure that the Celery backend releases recourses (for
# storing and transmitting results) associated to the task.
# Uncomment the line below if applicable.
res.forget()
def _collateral(self, run_uuid: Optional[str]):
"""Revoke the pipeline run celery task"""
# If there run status was not STARTED/PENDING then there is
# nothing to abort/revoke.
if not run_uuid:
return
celery_app = make_celery(current_app)
res = AbortableAsyncResult(run_uuid, app=celery_app)
# It is responsibility of the task to terminate by reading it's
# aborted status.
res.abort()
celery_app.control.revoke(run_uuid)
def _collateral(self, environment_build_uuid: Optional[str]):
if not environment_build_uuid:
return
celery_app = make_celery(current_app)
# Make use of both constructs (revoke, abort) so we cover both a
# task that is pending and a task which is running.
celery_app.control.revoke(environment_build_uuid, timeout=1.0)
res = AbortableAsyncResult(environment_build_uuid, app=celery_app)
# It is responsibility of the task to terminate by reading it's
# aborted status.
res.abort()
def _collateral(self, task_id: str, project_uuid: str,
environment_uuid: str, project_path: str):
celery = make_celery(current_app)
celery_job_kwargs = {
"project_uuid": project_uuid,
"environment_uuid": environment_uuid,
"project_path": project_path,
}
celery.send_task(
"app.core.tasks.build_environment",
kwargs=celery_job_kwargs,
task_id=task_id,
)
def stop_experiment(experiment_uuid) -> bool:
"""Stop an experiment.
Args:
experiment_uuid:
Returns:
True if the experiment exists and was stopped, false
if it did not exist or if it was already completed.
"""
experiment = models.Experiment.query.filter_by(
experiment_uuid=experiment_uuid).one_or_none()
if experiment is None:
return False
run_uuids = [
run.run_uuid for run in experiment.pipeline_runs
if run.status in ["PENDING", "STARTED"]
]
if len(run_uuids) == 0:
return False
# Aborts and revokes all pipeline runs and waits for a
# reply for 1.0s.
celery = make_celery(current_app)
celery.control.revoke(run_uuids, timeout=1.0)
# TODO: possibly set status of steps and Run to "ABORTED"
# note that a race condition would be present since the task
# will try to set the status as well
for run_uuid in run_uuids:
res = AbortableAsyncResult(run_uuid, app=celery)
# it is responsibility of the task to terminate by reading \
# it's aborted status
res.abort()
filter_by = {"run_uuid": run_uuid}
status_update = {"status": "ABORTED"}
update_status_db(status_update,
model=models.NonInteractivePipelineRun,
filter_by=filter_by)
update_status_db(status_update,
model=models.PipelineRunStep,
filter_by=filter_by)
db.session.commit()
return True
def _collateral(
self,
project_uuid: str,
task_id: str,
pipeline: Pipeline,
run_config: Dict[str, Any],
env_variables: Dict[str, Any],
env_uuid_to_image: Dict[str, str],
**kwargs,
):
# Create Celery object with the Flask context and construct the
# kwargs for the job.
celery = make_celery(current_app)
run_config["env_uuid_to_image"] = env_uuid_to_image
run_config["user_env_variables"] = env_variables
run_config["session_uuid"] = (
project_uuid[:18] + pipeline.properties["uuid"][:18]
)
run_config["session_type"] = "interactive"
celery_job_kwargs = {
"pipeline_definition": pipeline.to_dict(),
"run_config": run_config,
"session_uuid": run_config["session_uuid"],
}
# Start the run as a background task on Celery. Due to circular
# imports we send the task by name instead of importing the
# function directly.
res = celery.send_task(
"app.core.tasks.run_pipeline",
kwargs=celery_job_kwargs,
task_id=task_id,
)
# NOTE: this is only if a backend is configured. The task does
# not return anything. Therefore we can forget its result and
# make sure that the Celery backend releases recourses (for
# storing and transmitting results) associated to the task.
# Uncomment the line below if applicable.
res.forget()
def process_images_for_deletion(app) -> None:
"""Processes built images to find inactive ones.
Goes through env images and marks the inactive ones for removal,
moreover, if necessary, queues the celery task in charge of removing
images from the registry.
Note: this function should probably be moved in the scheduler module
to be consistent with the scheduler module of the webserver, said
module would need a bit of a refactoring.
"""
with app.app_context():
environments.mark_env_images_that_can_be_removed()
utils.mark_custom_jupyter_images_to_be_removed()
db.session.commit()
# Don't queue the task if there are build tasks going, this is a
# "cheap" way to avoid piling up multiple garbage collection
# tasks while a build is ongoing.
if db.session.query(
db.session.query(models.EnvironmentImageBuild).filter(
models.EnvironmentImageBuild.status.in_(
["PENDING", "STARTED"])).exists()).scalar():
app.logger.info("Ongoing build, not queueing registry gc task.")
return
if db.session.query(
db.session.query(models.JupyterImageBuild).filter(
models.JupyterImageBuild.status.in_(
["PENDING", "STARTED"])).exists()).scalar():
app.logger.info("Ongoing build, not queueing registry gc task.")
return
celery = make_celery(app)
app.logger.info("Sending registry garbage collection task.")
res = celery.send_task(
name="app.core.tasks.registry_garbage_collection")
res.forget()
def stop_pipeline_run(run_uuid) -> bool:
"""Stop a pipeline run.
The run will cancelled if not running yet, otherwise
it will be aborted.
Args:
run_uuid:
Returns:
True if a cancellation was issued to the run, false if the
run did not exist or was not PENDING/STARTED.
"""
interactive_run = models.InteractiveRun.query.filter(
models.InteractiveRun.status.in_(["PENDING", "STARTED"]),
models.InteractiveRun.run_uuid == run_uuid,
).one_or_none()
non_interactive_run = models.NonInteractiveRun.query.filter(
models.NonInteractiveRun.status.in_(["PENDING", "STARTED"]),
models.NonInteractiveRun.run_uuid == run_uuid,
).one_or_none()
if interactive_run is None and non_interactive_run is None:
return False
celery_app = make_celery(current_app)
res = AbortableAsyncResult(run_uuid, app=celery_app)
# it is responsibility of the task to terminate by reading
# it's aborted status
res.abort()
celery_app.control.revoke(run_uuid)
# TODO: possibly set status of steps and Run to "ABORTED"
# note that a race condition would be present since the
# task will try to set the status as well
return True
def post(self):
"""Queues a new experiment."""
# TODO: possibly use marshal() on the post_data
# https://flask-restplus.readthedocs.io/en/stable/api.html#flask_restplus.marshal
# to make sure the default values etc. are filled in.
post_data = request.get_json()
# TODO: maybe we can expect a datetime (in the schema) so we
# do not have to parse it here.
# https://flask-restplus.readthedocs.io/en/stable/api.html#flask_restplus.fields.DateTime
scheduled_start = post_data["scheduled_start"]
scheduled_start = datetime.fromisoformat(scheduled_start)
pipeline_runs = []
pipeline_run_spec = post_data["pipeline_run_spec"]
for pipeline_description, id_ in zip(
post_data["pipeline_descriptions"],
post_data["pipeline_run_ids"]):
pipeline_run_spec["pipeline_description"] = pipeline_description
pipeline = construct_pipeline(**post_data["pipeline_run_spec"])
# TODO: This can be made more efficient, since the pipeline
# is the same for all pipeline runs. The only
# difference is the parameters. So all the jobs could
# be created in batch.
# Create Celery object with the Flask context and construct the
# kwargs for the job.
celery = make_celery(current_app)
celery_job_kwargs = {
"experiment_uuid": post_data["experiment_uuid"],
"project_uuid": post_data["project_uuid"],
"pipeline_description": pipeline.to_dict(),
"run_config": pipeline_run_spec["run_config"],
}
# Start the run as a background task on Celery. Due to circular
# imports we send the task by name instead of importing the
# function directly.
res = celery.send_task(
"app.core.tasks.start_non_interactive_pipeline_run",
eta=scheduled_start,
kwargs=celery_job_kwargs,
)
# NOTE: this is only if a backend is configured. The task does
# not return anything. Therefore we can forget its result and
# make sure that the Celery backend releases recourses (for
# storing and transmitting results) associated to the task.
# Uncomment the line below if applicable.
res.forget()
non_interactive_run = {
"experiment_uuid": post_data["experiment_uuid"],
"run_uuid": res.id,
"pipeline_run_id": id_,
"pipeline_uuid": pipeline.properties["uuid"],
"project_uuid": post_data["project_uuid"],
"status": "PENDING",
}
db.session.add(models.NonInteractiveRun(**non_interactive_run))
# TODO: this code is also in `namespace_runs`. Could
# potentially be put in a function for modularity.
# Set an initial value for the status of the pipline steps that
# will be run.
step_uuids = [s.properties["uuid"] for s in pipeline.steps]
pipeline_steps = []
for step_uuid in step_uuids:
pipeline_steps.append(
models.NonInteractiveRunPipelineStep(
**{
"experiment_uuid": post_data["experiment_uuid"],
"run_uuid": res.id,
"step_uuid": step_uuid,
"status": "PENDING",
}))
db.session.bulk_save_objects(pipeline_steps)
db.session.commit()
non_interactive_run["pipeline_steps"] = pipeline_steps
pipeline_runs.append(non_interactive_run)
experiment = {
"experiment_uuid": post_data["experiment_uuid"],
"project_uuid": post_data["project_uuid"],
"pipeline_uuid": post_data["pipeline_uuid"],
"scheduled_start": scheduled_start,
"total_number_of_pipeline_runs": len(pipeline_runs),
}
db.session.add(models.Experiment(**experiment))
db.session.commit()
experiment["pipeline_runs"] = pipeline_runs
return experiment, 201
def _collateral(
self,
job: Dict[str, Any],
pipeline_run_spec: Dict[str, Any],
tasks_to_launch: Tuple[str, Pipeline],
):
# Safety check in case the job has no runs.
if not tasks_to_launch:
return
# Get docker ids of images to use and make it so that the
# images will not be deleted in case they become outdate by an
# an environment rebuild. Compute it only once because this way
# we are guaranteed that the mappings will be the same for all
# runs, having a new environment build terminate while
# submitting the different runs won't affect the job.
try:
env_uuid_docker_id_mappings = lock_environment_images_for_run(
# first (task_id, pipeline) -> task id.
tasks_to_launch[0][0],
job["project_uuid"],
# first (task_id, pipeline) -> pipeline.
tasks_to_launch[0][1].get_environments(),
)
except errors.ImageNotFound as e:
raise errors.ImageNotFound(
"Pipeline was referencing environments for "
f"which an image does not exist, {e}")
for task_id, _ in tasks_to_launch[1:]:
image_mappings = [
models.PipelineRunImageMapping(
**{
"run_uuid": task_id,
"orchest_environment_uuid": env_uuid,
"docker_img_id": docker_id,
})
for env_uuid, docker_id in env_uuid_docker_id_mappings.items()
]
db.session.bulk_save_objects(image_mappings)
db.session.commit()
# Launch each task through celery.
celery = make_celery(current_app)
for task_id, pipeline in tasks_to_launch:
run_config = pipeline_run_spec["run_config"]
run_config[
"env_uuid_docker_id_mappings"] = env_uuid_docker_id_mappings
celery_job_kwargs = {
"job_uuid": job["job_uuid"],
"project_uuid": job["project_uuid"],
"pipeline_definition": pipeline.to_dict(),
"run_config": run_config,
}
# Due to circular imports we use the task name instead of
# importing the function directly.
task_args = {
"name": "app.core.tasks.start_non_interactive_pipeline_run",
"eta": job["scheduled_start"],
"kwargs": celery_job_kwargs,
"task_id": task_id,
}
res = celery.send_task(**task_args)
# NOTE: this is only if a backend is configured. The task
# does not return anything. Therefore we can forget its
# result and make sure that the Celery backend releases
# recourses (for storing and transmitting results)
# associated to the task. Uncomment the line below if
# applicable.
res.forget()
from config import CONFIG_CLASS
from app import create_app
from app.celery_app import make_celery
from app.connections import docker_client
from app.core.environment_builds import build_environment_task
from app.core.pipelines import Pipeline, PipelineDefinition
from app.core.sessions import launch_noninteractive_session
logger = get_task_logger(__name__)
# TODO: create_app is called twice, meaning create_all (create
# databases) is called twice, which means celery-worker needs the
# /userdir bind to access the DB which is probably not a good idea.
# create_all should only be called once per app right?
celery = make_celery(create_app(CONFIG_CLASS, use_db=False), use_backend_db=True)
# This will not work yet, because Celery does not yet support asyncio
# tasks. In Celery 5.0 however this should be possible.
# https://stackoverflow.com/questions/39815771/how-to-combine-celery-with-asyncio
class APITask(Task):
"""
Idea:
Make the aiohttp.ClientSession persistent. Then we get:
"So if you’re making several requests to the same host, the
underlying TCP connection will be reused, which can result in a
significant performance increase."
def post(self):
"""Queues a new experiment."""
# TODO: possibly use marshal() on the post_data. Note that we
# have moved over to using flask_restx
# https://flask-restx.readthedocs.io/en/stable/api.html#flask_restx.marshal
# to make sure the default values etc. are filled in.
post_data = request.get_json()
# TODO: maybe we can expect a datetime (in the schema) so we
# do not have to parse it here. Again note that we are now
# using flask_restx
# https://flask-restx.readthedocs.io/en/stable/api.html#flask_restx.fields.DateTime
scheduled_start = post_data["scheduled_start"]
scheduled_start = datetime.fromisoformat(scheduled_start)
experiment = {
"experiment_uuid":
post_data["experiment_uuid"],
"project_uuid":
post_data["project_uuid"],
"pipeline_uuid":
post_data["pipeline_uuid"],
"scheduled_start":
scheduled_start,
"total_number_of_pipeline_runs":
len(post_data["pipeline_definitions"]),
}
db.session.add(models.Experiment(**experiment))
db.session.commit()
pipeline_runs = []
pipeline_run_spec = post_data["pipeline_run_spec"]
env_uuid_docker_id_mappings = None
# this way we write the entire exp to db, but avoid
# launching any run (celery task) if we detected a problem
experiment_creation_error_messages = []
tasks_to_launch = []
# TODO: This can be made more efficient, since the pipeline
# is the same for all pipeline runs. The only
# difference is the parameters. So all the jobs could
# be created in batch.
for pipeline_definition, id_ in zip(post_data["pipeline_definitions"],
post_data["pipeline_run_ids"]):
pipeline_run_spec["pipeline_definition"] = pipeline_definition
pipeline = construct_pipeline(**post_data["pipeline_run_spec"])
# specify the task_id beforehand to avoid race conditions
# between the task and its presence in the db
task_id = str(uuid.uuid4())
non_interactive_run = {
"experiment_uuid": post_data["experiment_uuid"],
"run_uuid": task_id,
"pipeline_run_id": id_,
"pipeline_uuid": pipeline.properties["uuid"],
"project_uuid": post_data["project_uuid"],
"status": "PENDING",
}
db.session.add(
models.NonInteractivePipelineRun(**non_interactive_run))
# need to flush because otherwise the bulk insertion of
# pipeline steps will lead to foreign key errors
# https://docs.sqlalchemy.org/en/13/orm/persistence_techniques.html#bulk-operations-caveats
db.session.flush()
# TODO: this code is also in `namespace_runs`. Could
# potentially be put in a function for modularity.
# Set an initial value for the status of the pipeline
# steps that will be run.
step_uuids = [s.properties["uuid"] for s in pipeline.steps]
pipeline_steps = []
for step_uuid in step_uuids:
pipeline_steps.append(
models.PipelineRunStep(
**{
"run_uuid": task_id,
"step_uuid": step_uuid,
"status": "PENDING",
}))
db.session.bulk_save_objects(pipeline_steps)
db.session.commit()
non_interactive_run["pipeline_steps"] = pipeline_steps
pipeline_runs.append(non_interactive_run)
# get docker ids of images to use and make it so that the
# images will not be deleted in case they become
# outdated by an environment rebuild
# compute it only once because this way we are guaranteed
# that the mappings will be the same for all runs, having
# a new environment build terminate while submitting the
# different runs won't affect the experiment
if env_uuid_docker_id_mappings is None:
try:
env_uuid_docker_id_mappings = lock_environment_images_for_run(
task_id,
post_data["project_uuid"],
pipeline.get_environments(),
)
except errors.ImageNotFound as e:
experiment_creation_error_messages.append(
f"Pipeline was referencing environments for "
f"which an image does not exist, {e}")
else:
image_mappings = [
models.PipelineRunImageMapping(
**{
"run_uuid": task_id,
"orchest_environment_uuid": env_uuid,
"docker_img_id": docker_id,
}) for env_uuid, docker_id in
env_uuid_docker_id_mappings.items()
]
db.session.bulk_save_objects(image_mappings)
db.session.commit()
if len(experiment_creation_error_messages) == 0:
# prepare the args for the task
run_config = pipeline_run_spec["run_config"]
run_config[
"env_uuid_docker_id_mappings"] = env_uuid_docker_id_mappings
celery_job_kwargs = {
"experiment_uuid": post_data["experiment_uuid"],
"project_uuid": post_data["project_uuid"],
"pipeline_definition": pipeline.to_dict(),
"run_config": run_config,
}
# Due to circular imports we use the task name instead
# of importing the function directly.
tasks_to_launch.append({
"name":
"app.core.tasks.start_non_interactive_pipeline_run",
"eta": scheduled_start,
"kwargs": celery_job_kwargs,
"task_id": task_id,
})
experiment["pipeline_runs"] = pipeline_runs
if len(experiment_creation_error_messages) == 0:
# Create Celery object with the Flask context
celery = make_celery(current_app)
for task in tasks_to_launch:
res = celery.send_task(**task)
# NOTE: this is only if a backend is configured.
# The task does not return anything. Therefore we can
# forget its result and make sure that the Celery
# backend releases recourses (for storing and
# transmitting results) associated to the task.
# Uncomment the line below if applicable.
res.forget()
return experiment, 201
else:
current_app.logger.error(
"\n".join(experiment_creation_error_messages))
# simple way to update both in memory objects
# and the db while avoiding multiple update statements
# (1 for each object)
for pipeline_run in experiment["pipeline_runs"]:
pipeline_run.status = "SUCCESS"
for step in pipeline_run["pipeline_steps"]:
step.status = "FAILURE"
models.PipelineRunStep.query.filter_by(
run_uuid=pipeline_run["run_uuid"]).update(
{"status": "FAILURE"})
models.NonInteractivePipelineRun.query.filter_by(
experiment_uuid=post_data["experiment_uuid"]).update(
{"status": "SUCCESS"})
db.session.commit()
return {
"message":
("Failed to create experiment because not all referenced"
"environments are available.")
}, 500
from app.core import environments, notifications, registry
from app.core.environment_image_builds import build_environment_image_task
from app.core.jupyter_image_builds import build_jupyter_image_task
from app.core.pipelines import Pipeline, run_pipeline_workflow
from app.core.sessions import launch_noninteractive_session
from app.types import PipelineDefinition, RunConfig
from config import CONFIG_CLASS
logger = get_task_logger(__name__)
# TODO: create_app is called twice, meaning create_all (create
# databases) is called twice, which means celery-worker needs the
# /userdir bind to access the DB which is probably not a good idea.
# create_all should only be called once per app right?
application = create_app(CONFIG_CLASS, use_db=True, register_api=False)
celery = make_celery(application, use_backend_db=True)
# This will not work yet, because Celery does not yet support asyncio
# tasks. In Celery 5.0 however this should be possible.
# https://stackoverflow.com/questions/39815771/how-to-combine-celery-with-asyncio
class APITask(Task):
"""
Idea:
Make the aiohttp.ClientSession persistent. Then we get:
"So if you’re making several requests to the same host, the
underlying TCP connection will be reused, which can result in a
significant performance increase."
def post(self):
"""Queues a list of environment builds.
Only unique requests are considered, meaning that a request
containing duplicate environment_build_requests will produce an
environment build only for each unique
environment_build_request. Note that requesting an
environment_build for an environment (identified by
project_uuid, environment_uuid, project_path) will REVOKE/ABORT
any other active (queued or actually started) environment build
for that environment. This implies that only an environment
build can be active (queued or actually started) for a given
environment.
"""
# keep only unique requests
post_data = request.get_json()
builds_requests = post_data["environment_build_requests"]
builds_requests = set([(req["project_uuid"], req["environment_uuid"],
req["project_path"])
for req in builds_requests])
builds_requests = [{
"project_uuid": req[0],
"environment_uuid": req[1],
"project_path": req[2],
} for req in builds_requests]
defined_builds = []
celery = make_celery(current_app)
# Start a celery task for each unique environment build request.
for build_request in builds_requests:
# Check if a build for this project/environment is
# PENDING/STARTED.
builds = models.EnvironmentBuild.query.filter(
models.EnvironmentBuild.project_uuid ==
build_request["project_uuid"],
models.EnvironmentBuild.environment_uuid ==
build_request["environment_uuid"],
models.EnvironmentBuild.project_path ==
build_request["project_path"],
or_(
models.EnvironmentBuild.status == "PENDING",
models.EnvironmentBuild.status == "STARTED",
),
).all()
for build in builds:
abort_environment_build(build.build_uuid,
build.status == "STARTED")
# We specify the task id beforehand so that we can commit to
# the db before actually launching the task, since the task
# might make some calls to the orchest-api referring to
# itself (e.g. a status update), and thus expecting to find
# itself in the db. This way we avoid race conditions.
task_id = str(uuid.uuid4())
# TODO: verify if forget has the same effect of
# ignore_result=True because ignore_result cannot be used
# with abortable tasks
# https://stackoverflow.com/questions/9034091/how-to-check-task-status-in-celery
# task.forget()
environment_build = {
"build_uuid":
task_id,
"project_uuid":
build_request["project_uuid"],
"environment_uuid":
build_request["environment_uuid"],
"project_path":
build_request["project_path"],
"requested_time":
datetime.fromisoformat(datetime.utcnow().isoformat()),
"status":
"PENDING",
}
defined_builds.append(environment_build)
db.session.add(models.EnvironmentBuild(**environment_build))
db.session.commit()
# could probably do without this...
celery_job_kwargs = {
"project_uuid": build_request["project_uuid"],
"environment_uuid": build_request["environment_uuid"],
"project_path": build_request["project_path"],
}
celery.send_task(
"app.core.tasks.build_environment",
kwargs=celery_job_kwargs,
task_id=task_id,
)
return {"environment_builds": defined_builds}
def post(self):
"""Queues a new experiment."""
# TODO: possibly use marshal() on the post_data
# https://flask-restplus.readthedocs.io/en/stable/api.html#flask_restplus.marshal
# to make sure the default values etc. are filled in.
post_data = request.get_json()
# TODO: maybe we can expect a datetime (in the schema) so we
# do not have to parse it here.
# https://flask-restplus.readthedocs.io/en/stable/api.html#flask_restplus.fields.DateTime
scheduled_start = post_data['scheduled_start']
scheduled_start = datetime.fromisoformat(scheduled_start)
pipeline_runs = []
pipeline_run_spec = post_data['pipeline_run_spec']
for pipeline_description, id_ in zip(
post_data['pipeline_descriptions'],
post_data['pipeline_run_ids']):
pipeline_run_spec['pipeline_description'] = pipeline_description
pipeline = construct_pipeline(**post_data['pipeline_run_spec'])
# TODO: This can be made more efficient, since the pipeline
# is the same for all pipeline runs. The only
# difference is the parameters. So all the jobs could
# be created in batch.
# Create Celery object with the Flask context and construct the
# kwargs for the job.
celery = make_celery(current_app)
celery_job_kwargs = {
'experiment_uuid': post_data['experiment_uuid'],
'pipeline_description': pipeline.to_dict(),
'run_config': pipeline_run_spec['run_config'],
}
# Start the run as a background task on Celery. Due to circular
# imports we send the task by name instead of importing the
# function directly.
res = celery.send_task(
'app.core.tasks.start_non_interactive_pipeline_run',
eta=scheduled_start,
kwargs=celery_job_kwargs)
non_interactive_run = {
'experiment_uuid': post_data['experiment_uuid'],
'run_uuid': res.id,
'pipeline_run_id': id_,
'pipeline_uuid': pipeline.properties['uuid'],
'status': 'PENDING',
}
db.session.add(models.NonInteractiveRun(**non_interactive_run))
# TODO: this code is also in `namespace_runs`. Could
# potentially be put in a function for modularity.
# Set an initial value for the status of the pipline steps that
# will be run.
step_uuids = [s.properties['uuid'] for s in pipeline.steps]
pipeline_steps = []
for step_uuid in step_uuids:
pipeline_steps.append(
models.NonInteractiveRunPipelineStep(
**{
'experiment_uuid': post_data['experiment_uuid'],
'run_uuid': res.id,
'step_uuid': step_uuid,
'status': 'PENDING'
}))
db.session.bulk_save_objects(pipeline_steps)
db.session.commit()
non_interactive_run['pipeline_steps'] = pipeline_steps
pipeline_runs.append(non_interactive_run)
experiment = {
'experiment_uuid': post_data['experiment_uuid'],
'pipeline_uuid': post_data['pipeline_uuid'],
'scheduled_start': scheduled_start,
}
db.session.add(models.Experiment(**experiment))
db.session.commit()
experiment['pipeline_runs'] = pipeline_runs
return experiment, 201
from app import create_app
from app.celery_app import make_celery
from app.connections import docker_client
from app.core.pipelines import Pipeline, PipelineDefinition
from app.core.sessions import launch_noninteractive_session
from app.core.environment_builds import build_environment_task
from config import CONFIG_CLASS
logger = get_task_logger(__name__)
# TODO: create_app is called twice, meaning create_all (create
# databases) is called twice, which means celery-worker needs the
# /userdir bind to access the DB which is probably not a good idea.
# create_all should only be called once per app right?
celery = make_celery(create_app(CONFIG_CLASS, use_db=False))
# This will not work yet, because Celery does not yet support asyncio
# tasks. In Celery 5.0 however this should be possible.
# https://stackoverflow.com/questions/39815771/how-to-combine-celery-with-asyncio
class APITask(Task):
"""
Idea:
Make the aiohttp.ClientSession persistent. Then we get:
"So if you’re making several requests to the same host, the
underlying TCP connection will be reused, which can result in a
significant performance increase."
def post(self):
"""Starts a new (interactive) pipeline run."""
post_data = request.get_json()
post_data["run_config"]["run_endpoint"] = "runs"
pipeline = construct_pipeline(**post_data)
# specify the task_id beforehand to avoid race conditions
# between the task and its presence in the db
task_id = str(uuid.uuid4())
# NOTE: we are setting the status of the run ourselves without
# using the option of celery to get the status of tasks. This
# way we do not have to configure a backend (where the default
# of "rpc://" does not give the results we would want).
run = {
"run_uuid": task_id,
"pipeline_uuid": pipeline.properties["uuid"],
"project_uuid": post_data["project_uuid"],
"status": "PENDING",
}
db.session.add(models.InteractiveRun(**run))
# Set an initial value for the status of the pipeline steps that
# will be run.
step_uuids = [s.properties["uuid"] for s in pipeline.steps]
pipeline_steps = []
for step_uuid in step_uuids:
pipeline_steps.append(
models.InteractiveRunPipelineStep(
**{
"run_uuid": task_id,
"step_uuid": step_uuid,
"status": "PENDING"
}))
db.session.bulk_save_objects(pipeline_steps)
db.session.commit()
run["pipeline_steps"] = pipeline_steps
# get docker ids of images to use and make it so that the images
# will not be deleted in case they become outdated by an
# environment rebuild
try:
env_uuid_docker_id_mappings = lock_environment_images_for_run(
task_id,
post_data["project_uuid"],
pipeline.get_environments(),
is_interactive=True,
)
except errors.ImageNotFound as e:
logging.error(f"Pipeline was referencing environments for "
f"which an image does not exist, {e}")
# simple way to update both in memory objects
# and the db while avoiding multiple update statements
# (1 for each object)
# TODO: make it so that the client does not rely
# on SUCCESS as a status
run["status"] = "SUCCESS"
for step in run["pipeline_steps"]:
step.status = "FAILURE"
models.InteractiveRun.query.filter_by(run_uuid=task_id).update(
{"status": "SUCCESS"})
models.InteractiveRunPipelineStep.query.filter_by(
run_uuid=task_id).update({"status": "FAILURE"})
db.session.commit()
return {
"message":
"Failed to start interactive run because not all referenced environments are available."
}, 500
# Create Celery object with the Flask context and construct the
# kwargs for the job.
celery = make_celery(current_app)
run_config = post_data["run_config"]
run_config["env_uuid_docker_id_mappings"] = env_uuid_docker_id_mappings
celery_job_kwargs = {
"pipeline_definition": pipeline.to_dict(),
"project_uuid": post_data["project_uuid"],
"run_config": run_config,
}
# Start the run as a background task on Celery. Due to circular
# imports we send the task by name instead of importing the
# function directly.
res = celery.send_task("app.core.tasks.run_pipeline",
kwargs=celery_job_kwargs,
task_id=task_id)
# NOTE: this is only if a backend is configured. The task does
# not return anything. Therefore we can forget its result and
# make sure that the Celery backend releases recourses (for
# storing and transmitting results) associated to the task.
# Uncomment the line below if applicable.
res.forget()
return marshal(run, schema.interactive_run), 201
def _get_worker_parallelism(worker: str) -> int:
celery = make_celery(current_app)
worker = f"celery@{worker}"
stats = celery.control.inspect([worker]).stats()
return len(stats[worker]["pool"]["processes"])
import asyncio
from typing import Dict, Union
import aiohttp
from celery import Task
from app import create_app
from app.celery_app import make_celery
from app.utils import Pipeline, PipelineDescription
from config import CONFIG_CLASS
celery = make_celery(create_app(CONFIG_CLASS))
# This will not work yet, because Celery does not yet support asyncio
# tasks. In Celery 5.0 however this should be possible.
# https://stackoverflow.com/questions/39815771/how-to-combine-celery-with-asyncio
class APITask(Task):
"""
Idea:
Make the aiohttp.ClientSession persistent. Then we get:
"So if you’re making several requests to the same host, the
underlying TCP connection will be reused, which can result in a
significant performance increase."
Recources:
https://docs.celeryproject.org/en/master/userguide/tasks.html#instantiation
"""
_session = None