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
