def _calculate_complexity(workflow): """Place workflow in queue and calculate and set its complexity.""" complexity = estimate_complexity(workflow.type_, workflow.reana_specification) workflow.complexity = complexity Session.commit() return complexity
def on_message(self, body, message): """On new message event handler.""" message.ack() body_dict = json.loads(body) workflow_uuid = body_dict.get("workflow_uuid") if workflow_uuid: workflow = ( Session.query(Workflow).filter_by(id_=workflow_uuid).one_or_none() ) next_status = body_dict.get("status") if next_status: next_status = RunStatus(next_status) print( " [x] Received workflow_uuid: {0} status: {1}".format( workflow_uuid, next_status ) ) logs = body_dict.get("logs") or "" if workflow.can_transition_to(next_status): _update_workflow_status(workflow, next_status, logs) if "message" in body_dict and body_dict.get("message"): msg = body_dict["message"] if "progress" in msg: _update_run_progress(workflow_uuid, msg) _update_job_progress(workflow_uuid, msg) # Caching: calculate input hash and store in JobCache if "caching_info" in msg: _update_job_cache(msg) Session.commit() else: logging.error( f"Cannot transition workflow {workflow.id_}" f" from status {workflow.status} to" f" {next_status}." )
def users_create_default(email, password, id_): """Create default user. This user has the administrator role and can retrieve other user information as well as create new users. """ reana_user_characteristics = { "id_": id_, "email": email, } try: user = User.query.filter_by(**reana_user_characteristics).first() if not user: reana_user_characteristics["access_token"] = secrets.token_urlsafe( 16) user = User(**reana_user_characteristics) create_user_workspace(user.get_user_workspace()) Session.add(user) Session.commit() # create invenio user, passing `confirmed_at` to mark it as confirmed register_user(email=email, password=password, confirmed_at=datetime.datetime.now()) click.echo(reana_user_characteristics["access_token"]) except Exception as e: click.echo("Something went wrong: {0}".format(e)) sys.exit(1)
def initialise_default_resources(): """Initialise default Resources.""" from reana_db.database import Session existing_resources = [r.name for r in Resource.query.all()] default_resources = [] resource_type_to_unit = { ResourceType.cpu: ResourceUnit.milliseconds, ResourceType.disk: ResourceUnit.bytes_, } for type_, name in DEFAULT_QUOTA_RESOURCES.items(): if name not in existing_resources: default_resources.append( Resource( name=name, type_=ResourceType[type_], unit=resource_type_to_unit[ResourceType[type_]], title="Default {} resource.".format(type_), )) if default_resources: Session.add_all(default_resources) Session.commit() return default_resources
def _create_and_associate_reana_user(sender, token=None, response=None, account_info=None): try: user_email = account_info['user']['email'] user_fullname = account_info['user']['profile']['full_name'] username = account_info['user']['profile']['username'] search_criteria = dict() search_criteria['email'] = user_email users = Session.query(User).filter_by(**search_criteria).all() if users: user = users[0] else: user_access_token = secrets.token_urlsafe(16) user_parameters = dict(access_token=user_access_token) user_parameters['email'] = user_email user_parameters['full_name'] = user_fullname user_parameters['username'] = username user = User(**user_parameters) Session.add(user) Session.commit() except (InvalidRequestError, IntegrityError): Session.rollback() raise ValueError('Could not create user, ' 'possible constraint violation') except Exception: raise ValueError('Could not create user') return user
def _create_and_associate_reana_user(sender, token=None, response=None, account_info=None): try: user_email = account_info["user"]["email"] user_fullname = account_info["user"]["profile"]["full_name"] username = account_info["user"]["profile"]["username"] search_criteria = dict() search_criteria["email"] = user_email users = Session.query(User).filter_by(**search_criteria).all() if users: user = users[0] else: user_parameters = dict(email=user_email, full_name=user_fullname, username=username) user = User(**user_parameters) Session.add(user) Session.commit() except (InvalidRequestError, IntegrityError): Session.rollback() raise ValueError("Could not create user, " "possible constraint violation") except Exception: raise ValueError("Could not create user") return user
def store_logs(logs, job_id): """Write logs to DB.""" try: logging.info("Storing job logs: {}".format(job_id)) Session.query(Job).filter_by(id_=job_id).update(dict(logs=logs)) Session.commit() except Exception as e: logging.error("Exception while saving logs: {}".format(str(e)), exc_info=True)
def token_grant(admin_access_token, id_, email): """Grant a token to the selected user.""" try: admin = User.query.filter_by(id_=ADMIN_USER_ID).one_or_none() if admin_access_token != admin.access_token: raise ValueError("Admin access token invalid.") user = _get_user_by_criteria(id_, email) error_msg = None if not user: error_msg = f"User {id_ or email} does not exist." elif user.access_token: error_msg = ( f"User {user.id_} ({user.email}) has already an active access token." ) if error_msg: click.secho(f"ERROR: {error_msg}", fg="red") sys.exit(1) if user.access_token_status in [UserTokenStatus.revoked.name, None]: click.confirm( f"User {user.id_} ({user.email}) access token status" f" is {user.access_token_status}, do you want to" " proceed?", abort=True, ) user_granted_token = secrets.token_urlsafe(16) user.access_token = user_granted_token Session.commit() log_msg = (f"Token for user {user.id_} ({user.email}) granted.\n" f"\nToken: {user_granted_token}") click.secho(log_msg, fg="green") admin.log_action(AuditLogAction.grant_token, {"reana_admin": log_msg}) # send notification to user by email email_subject = "REANA access token granted" email_body = JinjaEnv.render_template( "emails/token_granted.txt", user_full_name=user.full_name, reana_hostname=REANA_HOSTNAME, ui_config=REANAConfig.load("ui"), sender_email=ADMIN_EMAIL, ) send_email(user.email, email_subject, email_body) except click.exceptions.Abort: click.echo("Grant token aborted.") except REANAEmailNotificationError as e: click.secho( "Something went wrong while sending email:\n{}".format(e), fg="red", err=True, ) except Exception as e: click.secho( "Something went wrong while granting token:\n{}".format(e), fg="red", err=True, )
def cache_job(self): """Cache a job.""" workflow = Session.query(Workflow).filter_by( id_=self.workflow_uuid).one_or_none() access_times = calculate_file_access_time(workflow.workspace_path) prepared_job_cache = JobCache() prepared_job_cache.job_id = self.job_id prepared_job_cache.access_times = access_times Session.add(prepared_job_cache) Session.commit()
def _load_yadage_spec(workflow, operational_options): """Load and save in DB the Yadage workflow specification.""" operational_options.update({"accept_metadir": True}) toplevel = operational_options.get("toplevel", "") workflow.reana_specification = yadage_load_from_workspace( workflow.workspace_path, workflow.reana_specification, toplevel, ) Session.commit()
def on_message(self, body, message): """Process messages on ``jobs-status`` queue for alive workflows. This function will ignore events about workflows that have been already terminated since a graceful finalisation of the workflow cannot be guaranteed if the workflow engine (orchestrator) is not alive. """ try: message.ack() body_dict = json.loads(body) workflow_uuid = body_dict.get("workflow_uuid") workflow = (Session.query(Workflow).filter( Workflow.id_ == workflow_uuid, Workflow.status.in_(ALIVE_STATUSES), ).one_or_none()) if workflow: next_status = body_dict.get("status") if next_status: next_status = RunStatus(next_status) logging.info( " [x] Received workflow_uuid: {0} status: {1}".format( workflow_uuid, next_status)) logs = body_dict.get("logs") or "" if workflow.can_transition_to(next_status): _update_workflow_status(workflow, next_status, logs) if "message" in body_dict and body_dict.get("message"): msg = body_dict["message"] if "progress" in msg: _update_run_progress(workflow_uuid, msg) _update_job_progress(workflow_uuid, msg) # Caching: calculate input hash and store in JobCache if "caching_info" in msg: _update_job_cache(msg) Session.commit() else: logging.error(f"Cannot transition workflow {workflow.id_}" f" from status {workflow.status} to" f" {next_status}.") elif workflow_uuid: logging.warning( "Event for not alive workflow {workflow_uuid} received:\n" "{body}\n" "Ignoring ...".format(workflow_uuid=workflow_uuid, body=body)) except REANAWorkflowControllerError as rwce: logging.error(rwce, exc_info=True) except SQLAlchemyError as sae: logging.error( f"Something went wrong while querying the database for workflow: {workflow.id_}" ) logging.error(sae, exc_info=True) except Exception as e: logging.error(f"Unexpected error while processing workflow: {e}", exc_info=True)
def stop_workflow(workflow): """Stop a given workflow.""" if workflow.status == RunStatus.running: kwrm = KubernetesWorkflowRunManager(workflow) kwrm.stop_batch_workflow_run() workflow.status = RunStatus.stopped Session.add(workflow) Session.commit() else: message = ("Workflow {id_} is not running.").format(id_=workflow.id_) raise REANAWorkflowControllerError(message)
def update_workflow_logs(workflow_uuid, log_message): """Update workflow logs.""" try: logging.info('Storing workflow logs: {}'.format(workflow_uuid)) workflow = Session.query(Workflow).filter_by(id_=workflow_uuid).\ one_or_none() workflow.logs += '\n' + log_message Session.commit() except Exception as e: logging.error('Exception while saving logs: {}'.format(str(e)), exc_info=True)
def remove_workflow_jobs_from_cache(workflow): """Remove any cached jobs from given workflow. :param workflow: The workflow object that spawned the jobs. :return: None. """ jobs = Session.query(Job).filter_by(workflow_uuid=workflow.id_).all() for job in jobs: job_path = os.path.join(workflow.workspace_path, "..", "archive", str(job.id_)) Session.query(JobCache).filter_by(job_id=job.id_).delete() remove_workflow_workspace(job_path) Session.commit()
def remove_workflow_jobs_from_cache(workflow): """Remove any cached jobs from given workflow. :param workflow: The workflow object that spawned the jobs. :return: None. """ jobs = Session.query(Job).filter_by(workflow_uuid=workflow.id_).all() for job in jobs: job_path = remove_upper_level_references( os.path.join(workflow.get_workspace(), '..', 'archive', str(job.id_))) Session.query(JobCache).filter_by(job_id=job.id_).delete() remove_workflow_workspace(job_path) Session.commit()
def set_quota_limit(ctx, admin_access_token, emails, resource_name, limit): """Set quota limits to the given users per resource.""" try: for email in emails: error_msg = None user = _get_user_by_criteria(None, email) resource = Resource.query.filter_by( name=resource_name).one_or_none() if not user: error_msg = f"ERROR: Provided user {email} does not exist." elif not resource: error_msg = ( "ERROR: Provided resource name does not exist. Available " f"resources are {[resource.name for resource in Resource.query]}." ) if error_msg: click.secho( error_msg, fg="red", err=True, ) sys.exit(1) user_resource = UserResource.query.filter_by( user=user, resource=resource).one_or_none() if user_resource: user_resource.quota_limit = limit Session.add(user_resource) else: # Create user resource in case there isn't one. Useful for old users. user.resources.append( UserResource( user_id=user.id_, resource_id=resource.id_, quota_limit=limit, quota_used=0, )) Session.commit() click.secho( f"Quota limit {limit} for '{resource.name}' successfully set to users {emails}.", fg="green", ) except Exception as e: logging.debug(traceback.format_exc()) logging.debug(str(e)) click.echo( click.style("Quota could not be set: \n{}".format(str(e)), fg="red"), err=True, )
def token_revoke(admin_access_token, id_, email): """Revoke selected user's token.""" try: admin = User.query.filter_by(id_=ADMIN_USER_ID).one_or_none() if admin_access_token != admin.access_token: raise ValueError("Admin access token invalid.") user = _get_user_by_criteria(id_, email) error_msg = None if not user: error_msg = f"User {id_ or email} does not exist." elif not user.access_token: error_msg = (f"User {user.id_} ({user.email}) does not have an" " active access token.") if error_msg: click.secho(f"ERROR: {error_msg}", fg="red") sys.exit(1) revoked_token = user.access_token user.active_token.status = UserTokenStatus.revoked Session.commit() log_msg = (f"User token {revoked_token} ({user.email}) was" " successfully revoked.") click.secho(log_msg, fg="green") admin.log_action(AuditLogAction.revoke_token, {"reana_admin": log_msg}) # send notification to user by email email_subject = "REANA access token revoked" email_body = JinjaEnv.render_template( "emails/token_revoked.txt", user_full_name=user.full_name, reana_hostname=REANA_HOSTNAME, ui_config=REANAConfig.load("ui"), sender_email=ADMIN_EMAIL, ) send_email(user.email, email_subject, email_body) except REANAEmailNotificationError as e: click.secho( "Something went wrong while sending email:\n{}".format(e), fg="red", err=True, ) except Exception as e: click.secho( "Something went wrong while revoking token:\n{}".format(e), fg="red", err=True, )
def store_workflow_disk_quota(workflow, bytes_to_sum=None): """ Update or create disk workflow resource. :param workflow: Workflow whose disk resource usage must be calculated. :param bytes_to_sum: Amount of bytes to sum to workflow disk quota, if None, `du` will be used to recalculate it. :type workflow: reana_db.models.Workflow :type bytes_to_sum: int """ from reana_commons.errors import REANAMissingWorkspaceError from reana_commons.utils import get_disk_usage from reana_db.database import Session from reana_db.models import ResourceType, WorkflowResource def _get_disk_usage_or_zero(workflow): """Get disk usage for a workflow if the workspace exists, zero if not.""" try: disk_bytes = get_disk_usage(workflow.workspace_path, summarize=True) return int(disk_bytes[0]["size"]["raw"]) except REANAMissingWorkspaceError: return 0 disk_resource = get_default_quota_resource(ResourceType.disk.name) workflow_resource = ( Session.query(WorkflowResource) .filter_by(workflow_id=workflow.id_, resource_id=disk_resource.id_) .one_or_none() ) if workflow_resource: if bytes_to_sum: workflow_resource.quota_used += bytes_to_sum else: workflow_resource.quota_used = _get_disk_usage_or_zero(workflow) Session.commit() elif inspect(workflow).persistent: workflow_resource = WorkflowResource( workflow_id=workflow.id_, resource_id=disk_resource.id_, quota_used=_get_disk_usage_or_zero(workflow), ) Session.add(workflow_resource) Session.commit() return workflow_resource
def token_grant(admin_access_token, id_, email): """Grant a token to the selected user.""" try: admin = User.query.filter_by(id_=ADMIN_USER_ID).one_or_none() if admin_access_token != admin.access_token: raise ValueError("Admin access token invalid.") user = _get_user_by_criteria(id_, email) error_msg = None if not user: error_msg = f"User {id_ or email} does not exist." elif user.access_token: error_msg = (f"User {user.id_} ({user.email}) has already an" " active access token.") if error_msg: click.secho(f"ERROR: {error_msg}", fg="red") sys.exit(1) if user.access_token_status in [UserTokenStatus.revoked.name, None]: click.confirm( f"User {user.id_} ({user.email}) access token status" f" is {user.access_token_status}, do you want to" " proceed?", abort=True, ) user_granted_token = secrets.token_urlsafe(16) user.access_token = user_granted_token Session.commit() log_msg = (f"Token for user {user.id_} ({user.email}) granted.\n" f"\nToken: {user_granted_token}") click.secho(log_msg, fg="green") admin.log_action(AuditLogAction.grant_token, {"reana_admin": log_msg}) # send notification to user by email email_subject = "REANA access token granted" email_body = ( f"Dear {user.full_name},\n\nYour REANA access token has" f" been granted, please find it on https://{REANA_URL}/profile" "\n\nThe REANA support team") send_email(user.email, email_subject, email_body) except click.exceptions.Abort as e: click.echo("Grant token aborted.") except Exception as e: click.secho( "Something went wrong while granting token:\n{}".format(e), fg="red", err=True, )
def _import_users(admin_access_token, users_csv_file): """Import list of users to database. :param admin_access_token: Admin access token. :type admin_access_token: str :param users_csv_file: CSV file object containing a list of users. :type users_csv_file: _io.TextIOWrapper """ admin = User.query.filter_by(id_=ADMIN_USER_ID).one_or_none() if admin_access_token != admin.access_token: raise ValueError('Admin access token invalid.') csv_reader = csv.reader(users_csv_file) for row in csv_reader: user = User(id_=row[0], email=row[1], access_token=row[2]) Session.add(user) Session.commit() Session.remove()
def create_job_in_db(self, backend_job_id): """Create job in db.""" job_db_entry = JobTable(backend_job_id=backend_job_id, workflow_uuid=self.workflow_uuid, status=JobStatus.created.name, compute_backend=self.compute_backend, cvmfs_mounts=self.cvmfs_mounts or '', shared_file_system=self.shared_file_system or False, docker_img=self.docker_img, cmd=json.dumps(self.cmd), env_vars=json.dumps(self.env_vars), deleted=False, job_name=self.job_name, prettified_cmd=self.prettified_cmd) Session.add(job_db_entry) Session.commit() self.job_id = str(job_db_entry.id_)
def _create_user(email, user_access_token, admin_access_token): """Create user with provided credentials.""" try: admin = Session.query(User).filter_by(id_=ADMIN_USER_ID).one_or_none() if admin_access_token != admin.access_token: raise ValueError('Admin access token invalid.') if not user_access_token: user_access_token = secrets.token_urlsafe(16) user_parameters = dict(access_token=user_access_token) user_parameters['email'] = email user = User(**user_parameters) Session.add(user) Session.commit() except (InvalidRequestError, IntegrityError) as e: Session.rollback() raise ValueError('Could not create user, ' 'possible constraint violation') return user
def doesnt_exceed_max_reana_workflow_count(): """Check upper limit on running REANA batch workflows.""" doesnt_exceed = True try: running_workflows = (Session.query(func.count()).filter( or_( Workflow.status == RunStatus.pending, Workflow.status == RunStatus.running, )).scalar()) if running_workflows >= REANA_MAX_CONCURRENT_BATCH_WORKFLOWS: doesnt_exceed = False except SQLAlchemyError as e: logging.error( "Something went wrong while querying for number of running workflows." ) logging.error(e) doesnt_exceed = False Session.commit() return doesnt_exceed
def _create_and_associate_reana_user(email, fullname, username): try: search_criteria = dict() search_criteria["email"] = email users = Session.query(User).filter_by(**search_criteria).all() if users: user = users[0] else: user_parameters = dict(email=email, full_name=fullname, username=username) user = User(**user_parameters) Session.add(user) Session.commit() except (InvalidRequestError, IntegrityError): Session.rollback() raise ValueError( "Could not create user, possible constraint violation") except Exception: raise ValueError("Could not create user") return user
def update_workflow_cpu_quota(workflow) -> int: """Update workflow CPU quota based on started and finished/stopped times. :return: Workflow running time in milliseconds if workflow has terminated, else 0. """ from reana_db.database import Session from reana_db.models import ( ResourceType, UserResource, WorkflowResource, ) terminated_at = workflow.run_finished_at or workflow.run_stopped_at if workflow.run_started_at and terminated_at: cpu_time = terminated_at - workflow.run_started_at cpu_milliseconds = int(cpu_time.total_seconds() * 1000) cpu_resource = get_default_quota_resource(ResourceType.cpu.name) # WorkflowResource might exist already if the cluster # follows a combined termination + periodic policy (eg. created # by the status listener, revisited by the cronjob) workflow_resource = WorkflowResource.query.filter_by( workflow_id=workflow.id_, resource_id=cpu_resource.id_).one_or_none() if workflow_resource: workflow_resource.quota_used = cpu_milliseconds else: workflow_resource = WorkflowResource( workflow_id=workflow.id_, resource_id=cpu_resource.id_, quota_used=cpu_milliseconds, ) user_resource_quota = UserResource.query.filter_by( user_id=workflow.owner_id, resource_id=cpu_resource.id_).first() user_resource_quota.quota_used += cpu_milliseconds Session.add(workflow_resource) Session.commit() return cpu_milliseconds return 0
def store_workflow_disk_quota(workflow, bytes_to_sum: Optional[int] = None): """ Update or create disk workflow resource. :param workflow: Workflow whose disk resource usage must be calculated. :param bytes_to_sum: Amount of bytes to sum to workflow disk quota, if None, `du` will be used to recalculate it. :type workflow: reana_db.models.Workflow :type bytes_to_sum: int """ from reana_db.database import Session from reana_db.models import ResourceType, WorkflowResource if (ResourceType.disk.name not in WORKFLOW_TERMINATION_QUOTA_UPDATE_POLICY and not PERIODIC_RESOURCE_QUOTA_UPDATE_POLICY): return disk_resource = get_default_quota_resource(ResourceType.disk.name) workflow_resource = (Session.query(WorkflowResource).filter_by( workflow_id=workflow.id_, resource_id=disk_resource.id_).one_or_none()) if workflow_resource: if bytes_to_sum: workflow_resource.quota_used += bytes_to_sum else: workflow_resource.quota_used = get_disk_usage_or_zero( workflow.workspace_path) Session.commit() elif inspect(workflow).persistent: workflow_resource = WorkflowResource( workflow_id=workflow.id_, resource_id=disk_resource.id_, quota_used=get_disk_usage_or_zero(workflow.workspace_path), ) Session.add(workflow_resource) Session.commit() return workflow_resource
def on_message(self, body, message): """On new message event handler.""" message.ack() body_dict = json.loads(body) workflow_uuid = body_dict.get('workflow_uuid') if workflow_uuid: status = body_dict.get('status') if status: status = WorkflowStatus(status) print(" [x] Received workflow_uuid: {0} status: {1}".format( workflow_uuid, status)) logs = body_dict.get('logs') or '' _update_workflow_status(workflow_uuid, status, logs) if 'message' in body_dict and body_dict.get('message'): msg = body_dict['message'] if 'progress' in msg: _update_run_progress(workflow_uuid, msg) _update_job_progress(workflow_uuid, msg) # Caching: calculate input hash and store in JobCache if 'caching_info' in msg: _update_job_cache(msg) Session.commit()
def users_create_default(email, id_): """Create default user. This user has the administrator role and can retrieve other user information as well as create new users. """ user_characteristics = {"id_": id_, "email": email, } try: user = User.query.filter_by(**user_characteristics).first() if not user: user_characteristics['access_token'] = secrets.token_urlsafe() user = User(**user_characteristics) create_user_workspace(user.get_user_workspace()) Session.add(user) Session.commit() click.echo('Created 1st user with access_token: {}'. format(user_characteristics['access_token'])) except Exception as e: click.echo('Something went wrong: {0}'.format(e)) sys.exit(1)
def token_revoke(admin_access_token, id_, email): """Revoke selected user's token.""" try: admin = User.query.filter_by(id_=ADMIN_USER_ID).one_or_none() if admin_access_token != admin.access_token: raise ValueError("Admin access token invalid.") user = _get_user_by_criteria(id_, email) error_msg = None if not user: error_msg = f"User {id_ or email} does not exist." elif not user.access_token: error_msg = (f"User {user.id_} ({user.email}) does not have an" " active access token.") if error_msg: click.secho(f"ERROR: {error_msg}", fg="red") sys.exit(1) revoked_token = user.access_token user.active_token.status = UserTokenStatus.revoked Session.commit() log_msg = (f"User token {revoked_token} ({user.email}) was" " successfully revoked.") click.secho(log_msg, fg="green") admin.log_action(AuditLogAction.revoke_token, {"reana_admin": log_msg}) # send notification to user by email email_subject = "REANA access token revoked" email_body = (f"Dear {user.full_name},\n\nYour REANA access token has" " been revoked.\n\nThe REANA support team") send_email(user.email, email_subject, email_body) except Exception as e: click.secho( "Something went wrong while revoking token:\n{}".format(e), fg="red", err=True, )
def update_users_cpu_quota(user=None) -> None: """Update users CPU quota usage. :param user: User whose CPU quota will be updated. If None, applies to all users. :type user: reana_db.models.User """ from reana_db.database import Session from reana_db.models import ( ResourceType, User, UserResource, UserToken, UserTokenStatus, ) if (ResourceType.cpu.name not in WORKFLOW_TERMINATION_QUOTA_UPDATE_POLICY and not PERIODIC_RESOURCE_QUOTA_UPDATE_POLICY): return if user: users = [user] else: users = User.query.join(UserToken).filter_by( status=UserTokenStatus.active # skip users with no active token ) for user in users: cpu_milliseconds = 0 for workflow in user.workflows: cpu_milliseconds += update_workflow_cpu_quota(workflow=workflow) cpu_resource = get_default_quota_resource(ResourceType.cpu.name) user_resource_quota = UserResource.query.filter_by( user_id=user.id_, resource_id=cpu_resource.id_).first() user_resource_quota.quota_used = cpu_milliseconds Session.commit()