class Queue(object): def __init__(self): """ Initializes the Queue class """ self.cm_config = Config() self.info = munch.munchify({ 'uid': None, "cloud": None, "kind": "batch-queue", "name": None, "cm": {}, "queue": { 'policy': None, 'status': None, 'active': False, 'charge': None, 'unit': None, "numJobs": 0, "numRunningJobs": 0, 'joblist': [] } }) # list of parameters that can be set self.settable_params = ['policy', 'charge', 'unit'] self.database = CmDatabase() @DatabaseUpdate() def create(self, queue_name, cloud_name, policy, charge=None, unit=None): """ This method is used to create a queue :param queue_name: name of the queue to create :param cloud_name: slurm cluster on which the job is gonna run :param policy: policy of the queue :param charge: charge of the queue :param unit: unit of the charge for the queue :return: """ name = Name(order=["cloud", "name"], cloud=cloud_name, name=queue_name) uid = name.id(cloud=cloud_name, name=queue_name) # print(uid) self.info = munch.munchify({ 'uid': uid, "cloud": cloud_name, "kind": "batch-queue", "name": queue_name, "cm": { "cloud": cloud_name, "kind": "batch-queue", "name": queue_name, "cluster": self.cm_config.get('cloudmesh').get('cluster')[cloud_name] }, "queue": { 'policy': policy, 'status': 'EMPTY', 'active': False, 'charge': charge, 'unit': unit, "numJobs": 0, "numRunningJobs": 0, } }) # Console.error(self.info) self.policyFunctionMap = munch.munchify({ 'FIFO': self.popFIFO, 'FILO': self.popFILO }) if self.database.exists(self.info)[0]: Console.error("Queue already exists") return return [self.info] def findQueue(self, cloud_name, queue_name): ''' finds a queue in the database based on the name :param name: name of the queue :return: ''' # if self.database.exists(self.info)[0]: # Console.error("Queue already exists") name = Name(order=["cloud", "name"], cloud=cloud_name, name=queue_name) uid = name.id(cloud=cloud_name, name=queue_name) queue = self.database.find_by_KeyValue( collection_name="{cloud}-{kind}".format(cloud=cloud_name, kind='batch-queue'), KeyValue={'uid': uid}) if type(queue) is cursor.Cursor: self.info = munch.munchify(queue[0]) return True # # queue found elif type(queue) is list and len(queue) == 0: return False # queue not found def findClouds(self): ''' finds all queues in the database based on the name :return: ''' for collection in self.database.collections(): if 'batch-queue' in collection: print(collection) # all_queues = self.database.db.find() # print(all_queues) def findQueues(self, cloud_name): ''' finds all queues in the database based on the name :return: ''' # TODO: find all queues info from the DB based on the ['cm'] all_queues = self.database.find_by_KeyValue(collection_name=cloud_name) all_queues = [munch.munchify(queue) for queue in all_queues] for queue in all_queues: print(queue.uid) def listJobs(self): ''' list the jobs in the current queue :return: ''' return def removeQueue(self): ''' remove the queue from the database :return: ''' # TODO: remove the queues info from the DB based on the ['cm'] return @DatabaseUpdate() # this should update the record not create a new one def push(self, job): ''' push job to stack :param job: :return: ''' self.info.queue.joblist.append(job) self.info.queue.numJobs += 1 self.updateStatus() return self.info @DatabaseUpdate() # this should update the record not create a new one def pop(self): ''' pop job from stack based on the policy :param job: :return: ''' self.info.queue.numJobs -= 1 self.updateStatus() policy = self.info.queue.policy return self.policyFunctionMap[policy]() def popFIFO(self): ''' pop job from stack based on FIFO policy :param job: :return: ''' return self.info['queue']['joblist'].pop(0) def popFILO(self): ''' pop job from stack based on FIFO policy :param job: :return: ''' return self.info['queue']['joblist'].pop() def isEmpty(self): ''' checks if the queue is empty :return: ''' if self.info.queue.numJobs > 0: return False return True @DatabaseUpdate() # this should update the record not create a new one def activate(self): ''' activates the queue :return: ''' # TODO: activating a queue should start submitting jobs self.info.queue.active = True return self.info @DatabaseUpdate() # this should update the record not create a new one def deactivate(self): ''' deactivates the queue :return: ''' # TODO: stop all jobs self.info.queue.active = False return self.info @DatabaseUpdate() # this should update the record not create a new one def updateStatus(self): ''' checks number of jobs and updates queue status :return: ''' if self.info.queue.numJobs > 0: self.info.queue.status = 'FULL' return self.info @DatabaseUpdate() def setParam(self, param, val): ''' set a particular parameter in the queue :param param: the parameter :param val: value of the parameter :return: ''' if param in self.settable_params: self.info.queue[param] = val else: Console.error("Only the following parameters could be set in a " "queue: \n" + ', '.join(self.settable_params)) return self.info
class WorkFlowDB(object): def __init__(self, name="workflow", active=False): self.database = CmDatabase() self.workflow_name = name self.collection = self.database.collection(f"{name}-flow") if active: self.switch_to_active_flow() def attributes(self, name): data = { "cm": { "kind": "flow", "cloud": self.workflow_name, "name": name }, "kind": "flow", "cloud": self.workflow_name, "name": name, "status": "defined" } return data @DatabaseUpdate() def add_node(self, node): name = node["name"] node.update(self.attributes(name)) VERBOSE(node) return node def add_edge(self, node, depends_on): self.collection.update_one({"name": node}, {"$push": { "dependencies": depends_on }}) def _node_from_db(self, db_obj): reconstructed = Node(db_obj["name"]) reconstructed.workflow = self.workflow_name reconstructed.dependencies = db_obj["dependencies"] reconstructed.status = db_obj.get("status", "pending") reconstructed.result = db_obj.get("result", {}) reconstructed.progress = "" reconstructed.modified = "" reconstructed.done = "" return reconstructed def remove_node(self, name): self.collection.delete_one({"name": name}) self.collection.update_many({}, {"$pull": {"dependencies": "name"}}) def remove_edge(self, node, depends_on): self.collection.update_one({"name": node}, {"$pull": { "dependencies": depends_on }}) def get_node(self, name=None): return self._node_from_db(self.collection.find_one({"name": name})) def list(self, node=None, edge=None): query = {} if node: query["name"] = node if edge: query["dependencies"] = edge return self.collection.find(query) def list_nodes(self): return [self._node_from_db(node) for node in self.list()] def list_edges(self): return self.collection.aggregate([{ "$unwind": "$dependencies" }, { "$project": { "to": "$name", "from": "$dependencies" } }]) def list_all_workflows(self): all_colls = self.database.collections() return [ name for name in all_colls if "flow" in name and "active" not in name ] def set_node_status(self, node, status): return self.collection.update_one({"name": node}, {"$set": { "status": status }}) def find_root_nodes(self): root_nodes = self.collection.find({ "dependencies.0": { "$exists": False }, "status": "pending" }) return [self._node_from_db(node) for node in root_nodes] def switch_to_active_flow(self): started_collection = f"{self.workflow_name}-flow-active" self.collection = self.database.collection(started_collection) def resolve_node_dependencies(self, name=None): return self.collection.update_many({"dependencies": name}, {"$pull": { "dependencies": name }}) def add_specification(self, spec): pass def start_flow(self): started_collection = f"{self.workflow_name}-flow-active" self.collection.aggregate([{ "$project": { "dependencies": 1, "cm": 1, "kind": 1, "cloud": 1, "name": 1, "status": "pending" } }, { "$out": started_collection }]) self.switch_to_active_flow() def add_node_result(self, nodename, result): return self.collection.update_one({"name": nodename}, {"$set": { "result": result }}) def add_graph(self, yamlfile): pass def last_update(self, workflow=None): """ This method returns the last modified value associated with a database update to a node. :param workflow: The name of the workflow :type workflow: string :return: The time of the last update :rtype: string """ raise NotImplementedError t = "the time string" return t
class WorkFlowDB(object): def __init__(self, name="workflow"): self.database = CmDatabase() self.workflow_name = name self.collection = self.database.collection(f"{name}-flow") def attributes(self, name): data = { "cm": { "kind": "flow", "cloud": self.workflow_name, "name": name }, "kind": "flow", "cloud": self.workflow_name, "name": name } return data @DatabaseUpdate() def add_node(self, node): name = node["name"] node.update(self.attributes(name)) return node def add_edge(self, node, depends_on): self.collection.update_one({"name": node}, {"$push": { "dependencies": depends_on }}) def _node_from_db(self, db_obj): reconstructed = Node(db_obj["name"]) reconstructed.workflow = self.workflow_name reconstructed.dependencies = db_obj["dependencies"] reconstructed.status = db_obj.get("status", "pending") return reconstructed def get_node(self, name=None): return self._node_from_db(self.collection.find_one({"name": name})) def list(self, node=None, edge=None): query = {} if node: query["name"] = node if edge: query["dependencies"] = edge return self.collection.find(query) def list_nodes(self): return [self._node_from_db(node) for node in self.list()] def list_edges(self): return self.collection.aggregate([{ "$unwind": "dependecies" }, { "$project": { "to": "$name", "from": "$dependecies" } }]) def list_all_workflows(self): all_colls = self.database.collections() return [ name for name in all_colls if "flow" in name and "active" not in name ] def set_node_status(self, node, status): return self.collection.update_one({"name": node}, {"$set": { "status": status }}) def find_root_nodes(self): root_nodes = self.collection.find({ "dependencies.0": { "$exists": False }, "status": "pending" }) return [self._node_from_db(node) for node in root_nodes] def switch_to_active_flow(self): started_collection = f"{self.workflow_name}-flow-active" self.collection = self.database.collection(started_collection) def resolve_node_dependency(self, name=None): return self.collection.update_many({"dependencies": name}, {"$pull": { "dependencies": name }}) def add_specification(self, spec): pass def start_flow(self): started_collection = f"{self.workflow_name}-flow-active" self.collection.aggregate([{ "$project": { "dependencies": 1, "cm": 1, "kind": 1, "cloud": 1, "name": 1, "status": "pending" } }, { "$out": started_collection }]) self.switch_to_active_flow() def add_graph(self, yamlfile): pass