def _transaction(
self,
project_uuid: str,
pipeline_uuid: str,
):
session = models.InteractiveSession.query.filter_by(
project_uuid=project_uuid, pipeline_uuid=pipeline_uuid, status="RUNNING"
).one_or_none()
if session is None:
self.collateral_kwargs["container_ids"] = None
self.collateral_kwargs["notebook_server_info"] = None
return False
else:
# Abort interactive run if it was PENDING/STARTED.
run = models.InteractivePipelineRun.query.filter(
models.InteractivePipelineRun.project_uuid == project_uuid,
models.InteractivePipelineRun.pipeline_uuid == pipeline_uuid,
models.InteractivePipelineRun.status.in_(["PENDING", "STARTED"]),
).one_or_none()
if run is not None:
AbortPipelineRun(self.tpe).transaction(run.uuid)
self.collateral_kwargs["container_ids"] = session.container_ids
self.collateral_kwargs[
"notebook_server_info"
] = session.notebook_server_info
return True
def _transaction(self, project_uuid: str):
# Any interactive run related to the project is stopped if
# if necessary, then deleted.
interactive_runs = (models.InteractivePipelineRun.query.filter_by(
project_uuid=project_uuid).filter(
models.InteractivePipelineRun.status.in_(
["PENDING", "STARTED"])).all())
for run in interactive_runs:
AbortPipelineRun(self.tpe).transaction(run.run_uuid)
# Will delete cascade interactive run pipeline step,
# interactive run image mapping.
db.session.delete(run)
# Stop (and delete) any interactive session related to the
# project.
sessions = (models.InteractiveSession.query.filter_by(
project_uuid=project_uuid, ).with_entities(
models.InteractiveSession.project_uuid,
models.InteractiveSession.pipeline_uuid,
).distinct().all())
for session in sessions:
# Stop any interactive session related to the pipeline.
StopInteractiveSession(self.tpe).transaction(
project_uuid, session.pipeline_uuid)
# Any job related to the pipeline is stopped if necessary
# , then deleted.
jobs = (models.Job.query.filter_by(
project_uuid=project_uuid, ).with_entities(
models.Job.job_uuid).all())
for job in jobs:
DeleteJob(job.job_uuid)
# Remove images related to the project.
DeleteProjectEnvironmentImages(self.tpe).transaction(project_uuid)
def _transaction(
self,
project_uuid: str,
pipeline_uuid: str,
):
# The with for update is to avoid a race condition where
# updating the status to STOPPING would overwrite a RUNNING
# status after reading the previous_state as LAUNCHING, which
# would then cause the collateral effect to wait the full time
# before shutting down the session.
session = (models.InteractiveSession.query.with_for_update().
populate_existing().filter_by(
project_uuid=project_uuid,
pipeline_uuid=pipeline_uuid).one_or_none())
if session is None:
self.collateral_kwargs["project_uuid"] = None
self.collateral_kwargs["pipeline_uuid"] = None
self.collateral_kwargs["container_ids"] = None
self.collateral_kwargs["notebook_server_info"] = None
self.collateral_kwargs["previous_state"] = None
return False
else:
# Abort interactive run if it was PENDING/STARTED.
run = models.InteractivePipelineRun.query.filter(
models.InteractivePipelineRun.project_uuid == project_uuid,
models.InteractivePipelineRun.pipeline_uuid == pipeline_uuid,
models.InteractivePipelineRun.status.in_(
["PENDING", "STARTED"]),
).one_or_none()
if run is not None:
AbortPipelineRun(self.tpe).transaction(run.uuid)
previous_state = session.status
session.status = "STOPPING"
self.collateral_kwargs["project_uuid"] = project_uuid
self.collateral_kwargs["pipeline_uuid"] = pipeline_uuid
# This data is kept here instead of querying again in the
# collateral phase because when deleting a project the
# project deletion (in the transactional phase) will cascade
# delete the session, so the collateral phase would not be
# able to find the session by querying the db.
self.collateral_kwargs["container_ids"] = session.container_ids
self.collateral_kwargs[
"notebook_server_info"] = session.notebook_server_info
self.collateral_kwargs["previous_state"] = previous_state
return True
def _transaction(self, project_uuid: str, pipeline_uuid: str):
# Any interactive run related to the pipeline is stopped if
# necessary, then deleted.
interactive_runs = (models.InteractivePipelineRun.query.filter_by(
project_uuid=project_uuid, pipeline_uuid=pipeline_uuid).filter(
models.InteractivePipelineRun.status.in_(
["PENDING", "STARTED"])).all())
for run in interactive_runs:
AbortPipelineRun(self.tpe).transaction(run.run_uuid)
# Will delete cascade: run pipeline step, interactive run
# image mapping,
db.session.delete(run)
# Stop any interactive session related to the pipeline.
StopInteractiveSession(self.tpe).transaction(project_uuid,
pipeline_uuid)
def _transaction(self, project_uuid: str, environment_uuid: str):
# Stop all interactive runs making use of the env.
int_runs = interactive_runs_using_environment(project_uuid,
environment_uuid)
for run in int_runs:
AbortPipelineRun(self.tpe).transaction(run.run_uuid)
# Stop all jobs making use of the environment.
exps = jobs_using_environment(project_uuid, environment_uuid)
for exp in exps:
AbortJob(self.tpe).transaction(exp.job_uuid)
# Cleanup references to the builds and dangling images
# of this environment.
DeleteProjectEnvironmentBuilds(self.tpe).transaction(
project_uuid, environment_uuid)
self.collateral_kwargs["project_uuid"] = project_uuid
self.collateral_kwargs["environment_uuid"] = environment_uuid
def _transaction(
self,
project_uuid: str,
pipeline_uuid: str,
):
session = models.InteractiveSession.query.filter_by(
project_uuid=project_uuid, pipeline_uuid=pipeline_uuid
).one_or_none()
if session is None:
self.collateral_kwargs["project_uuid"] = None
self.collateral_kwargs["pipeline_uuid"] = None
self.collateral_kwargs["container_ids"] = None
self.collateral_kwargs["notebook_server_info"] = None
return False
else:
# Abort interactive run if it was PENDING/STARTED.
run = models.InteractivePipelineRun.query.filter(
models.InteractivePipelineRun.project_uuid == project_uuid,
models.InteractivePipelineRun.pipeline_uuid == pipeline_uuid,
models.InteractivePipelineRun.status.in_(["PENDING", "STARTED"]),
).one_or_none()
if run is not None:
AbortPipelineRun(self.tpe).transaction(run.uuid)
session.status = "STOPPING"
self.collateral_kwargs["project_uuid"] = project_uuid
self.collateral_kwargs["pipeline_uuid"] = pipeline_uuid
# This data is kept here instead of querying again in the
# collateral phase because when deleting a project the
# project deletion (in the transactional phase) will cascade
# delete the session, so the collateral phase would not be
# able to find the session by querying the db.
self.collateral_kwargs["container_ids"] = session.container_ids
self.collateral_kwargs[
"notebook_server_info"
] = session.notebook_server_info
return True
def _transaction(self, project_uuid: str, environment_uuid: str):
# Stop all interactive sessions making use of the env by using
# it as a service.
int_sess = environments.interactive_sessions_using_environment(
project_uuid, environment_uuid)
for sess in int_sess:
StopInteractiveSession(self.tpe).transaction(sess.project_uuid,
sess.pipeline_uuid,
async_mode=True)
# Stop all interactive runs making use of the env.
int_runs = environments.interactive_runs_using_environment(
project_uuid, environment_uuid)
for run in int_runs:
AbortPipelineRun(self.tpe).transaction(run.uuid)
# Stop all jobs making use of the environment.
jobs = environments.jobs_using_environment(project_uuid,
environment_uuid)
for job in jobs:
AbortJob(self.tpe).transaction(job.uuid)
# Mark images to be removed from nodes and registry.
environments.mark_all_proj_env_images_to_be_removed_on_env_deletion(
project_uuid=project_uuid,
environment_uuid=environment_uuid,
)
# Cleanup references to the builds and dangling images of this
# environment.
DeleteProjectEnvironmentImageBuilds(self.tpe).transaction(
project_uuid, environment_uuid)
self.collateral_kwargs["project_uuid"] = project_uuid
self.collateral_kwargs["environment_uuid"] = environment_uuid
models.Environment.query.filter_by(project_uuid=project_uuid,
uuid=environment_uuid).delete()
def cleanup():
app = create_app(
config_class=CONFIG_CLASS, use_db=True, be_scheduler=False, register_api=False
)
with app.app_context():
app.logger.info("Starting app cleanup.")
try:
app.logger.info("Aborting interactive pipeline runs.")
runs = InteractivePipelineRun.query.filter(
InteractivePipelineRun.status.in_(["PENDING", "STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for run in runs:
AbortPipelineRun(tpe).transaction(run.uuid)
app.logger.info("Shutting down interactive sessions.")
int_sessions = InteractiveSession.query.all()
with TwoPhaseExecutor(db.session) as tpe:
for session in int_sessions:
StopInteractiveSession(tpe).transaction(
session.project_uuid, session.pipeline_uuid, async_mode=False
)
app.logger.info("Aborting environment builds.")
builds = EnvironmentImageBuild.query.filter(
EnvironmentImageBuild.status.in_(["PENDING", "STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for build in builds:
AbortEnvironmentImageBuild(tpe).transaction(
build.project_uuid,
build.environment_uuid,
build.image_tag,
)
app.logger.info("Aborting jupyter builds.")
builds = JupyterImageBuild.query.filter(
JupyterImageBuild.status.in_(["PENDING", "STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for build in builds:
AbortJupyterEnvironmentBuild(tpe).transaction(build.uuid)
app.logger.info("Aborting running one off jobs.")
jobs = Job.query.filter_by(schedule=None, status="STARTED").all()
with TwoPhaseExecutor(db.session) as tpe:
for job in jobs:
AbortJob(tpe).transaction(job.uuid)
app.logger.info("Aborting running pipeline runs of cron jobs.")
runs = NonInteractivePipelineRun.query.filter(
NonInteractivePipelineRun.status.in_(["STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for run in runs:
AbortPipelineRun(tpe).transaction(run.uuid)
# Delete old JupyterEnvironmentBuilds on to avoid
# accumulation in the DB. Leave the latest such that the
# user can see details about the last executed build after
# restarting Orchest.
jupyter_image_builds = (
JupyterImageBuild.query.order_by(
JupyterImageBuild.requested_time.desc()
)
.offset(1)
.all()
)
# Can't use offset and .delete in conjunction in sqlalchemy
# unfortunately.
for jupyter_image_build in jupyter_image_builds:
db.session.delete(jupyter_image_build)
db.session.commit()
except Exception as e:
app.logger.error("Cleanup failed.")
app.logger.error(e)
def create_app(config_class=None, use_db=True, be_scheduler=False):
"""Create the Flask app and return it.
Args:
config_class: Configuration class. See orchest-api/app/config.
use_db: If true, associate a database to the Flask app instance,
which implies connecting to a given database and possibly
creating such database and/or tables if they do not exist
already. The reason to differentiate instancing the app
through this argument is that the celery worker does not
need to connect to the db that "belongs" to the orchest-api.
be_scheduler: If true, a background thread will act as a job
scheduler, according to the logic in core/scheduler. While
Orchest runs, only a single process should be acting as
scheduler.
Returns:
Flask.app
"""
app = Flask(__name__)
app.config.from_object(config_class)
init_logging()
# Cross-origin resource sharing. Allow API to be requested from the
# different microservices such as the webserver.
CORS(app, resources={r"/*": {"origins": "*"}})
if os.getenv("FLASK_ENV") == "development":
app = register_teardown_request(app)
if use_db:
# Create the database if it does not exist yet. Roughly equal to
# a "CREATE DATABASE IF NOT EXISTS <db_name>" call.
if not database_exists(app.config["SQLALCHEMY_DATABASE_URI"]):
create_database(app.config["SQLALCHEMY_DATABASE_URI"])
db.init_app(app)
# necessary for migration
Migrate().init_app(app, db)
with app.app_context():
# Alembic does not support calling upgrade() concurrently
if not is_werkzeug_parent():
# Upgrade to the latest revision. This also takes
# care of bringing an "empty" db (no tables) on par.
try:
upgrade()
except Exception as e:
logging.error("Failed to run upgrade() %s [%s]" % (e, type(e)))
# In case of an ungraceful shutdown, these entities could be
# in an invalid state, so they are deleted, since for sure
# they are not running anymore.
# To avoid the issue of entities being deleted because of a
# flask app reload triggered by a --dev code change, we
# attempt to create a directory first. Since this is an
# atomic operation that will result in an error if the
# directory is already there, this cleanup operation will
# run only once per container.
try:
os.mkdir("/tmp/cleanup_done")
InteractiveSession.query.delete()
# Delete old JupyterBuilds on start to avoid
# accumulation in the DB. Leave the latest such that the
# user can see details about the last executed build
# after restarting Orchest.
jupyter_builds = (
JupyterBuild.query.order_by(JupyterBuild.requested_time.desc())
.offset(1)
.all()
)
# Can't use offset and .delete in conjunction in
# sqlalchemy unfortunately.
for jupyer_build in jupyter_builds:
db.session.delete(jupyer_build)
db.session.commit()
# Fix interactive runs.
runs = InteractivePipelineRun.query.filter(
InteractivePipelineRun.status.in_(["PENDING", "STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for run in runs:
AbortPipelineRun(tpe).transaction(run.uuid)
# Fix one off jobs (and their pipeline runs).
jobs = Job.query.filter_by(schedule=None, status="STARTED").all()
with TwoPhaseExecutor(db.session) as tpe:
for job in jobs:
AbortJob(tpe).transaction(job.uuid)
# This is to fix the state of cron jobs pipeline runs.
runs = NonInteractivePipelineRun.query.filter(
NonInteractivePipelineRun.status.in_(["STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for run in runs:
AbortPipelineRun(tpe).transaction(run.uuid)
# Fix env builds.
builds = EnvironmentBuild.query.filter(
EnvironmentBuild.status.in_(["PENDING", "STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for build in builds:
AbortEnvironmentBuild(tpe).transaction(build.uuid)
# Fix jupyter builds.
builds = JupyterBuild.query.filter(
JupyterBuild.status.in_(["PENDING", "STARTED"])
).all()
with TwoPhaseExecutor(db.session) as tpe:
for build in builds:
AbortJupyterBuild(tpe).transaction(build.uuid)
# Trigger a build of JupyterLab if no JupyterLab image
# is found for this version and JupyterLab setup_script
# is non-empty.
trigger_conditional_jupyter_build(app)
except FileExistsError:
app.logger.info("/tmp/cleanup_done exists. Skipping cleanup.")
except Exception as e:
app.logger.error("Cleanup failed")
app.logger.error(e)
if be_scheduler and not is_werkzeug_parent():
# Create a scheduler and have the scheduling logic running
# periodically.
scheduler = BackgroundScheduler(
job_defaults={
# Infinite amount of grace time, so that if a task
# cannot be instantly executed (e.g. if the webserver is
# busy) then it will eventually be.
"misfire_grace_time": 2 ** 31,
"coalesce": False,
# So that the same job can be in the queue an infinite
# amount of times, e.g. for concurrent requests issuing
# the same tasks.
"max_instances": 2 ** 31,
}
)
app.config["SCHEDULER"] = scheduler
scheduler.start()
scheduler.add_job(
Scheduler.check_for_jobs_to_be_scheduled,
"interval",
seconds=app.config["SCHEDULER_INTERVAL"],
args=[app],
)
# Register blueprints.
app.register_blueprint(api, url_prefix="/api")
return app
