def list_groups(self, owner_id):
'''
Return a list of row dicts --one per group-- for the given owner.
'''
with self.engine.begin() as connection:
rows = connection.execute(cl_group.select().where(
cl_group.c.owner_id == owner_id
)).fetchall()
return [dict(row) for row in sorted(rows, key=lambda row: row.id)]
def batch_get_groups(self, **kwargs):
'''
Get a list of groups, all of which satisfy the clause given by kwargs.
'''
clause = self.make_kwargs_clause(cl_group, kwargs)
with self.engine.begin() as connection:
rows = connection.execute(
cl_group.select().where(clause)
).fetchall()
if not rows:
return []
values = {row.uuid: dict(row) for row in rows}
return [value for value in values.itervalues()]
def worker_checkin(
self,
user_id,
worker_id,
tag,
group_name,
cpus,
gpus,
memory_bytes,
free_disk_bytes,
dependencies,
shared_file_system,
tag_exclusive,
exit_after_num_runs,
is_terminating,
preemptible,
):
"""
Adds the worker to the database, if not yet there. Returns the socket ID
that the worker should listen for messages on.
"""
with self._engine.begin() as conn:
worker_row = {
'tag': tag,
'cpus': cpus,
'gpus': gpus,
'memory_bytes': memory_bytes,
'free_disk_bytes': free_disk_bytes,
'checkin_time': datetime.datetime.utcnow(),
'shared_file_system': shared_file_system,
'tag_exclusive': tag_exclusive,
'exit_after_num_runs': exit_after_num_runs,
'is_terminating': is_terminating,
'preemptible': preemptible,
}
# Populate the group for this worker, if group_name is valid
group_row = conn.execute(cl_group.select().where(
cl_group.c.name == group_name)).fetchone()
if group_row:
worker_row['group_uuid'] = group_row.uuid
existing_row = conn.execute(cl_worker.select().where(
and_(cl_worker.c.user_id == user_id,
cl_worker.c.worker_id == worker_id))).fetchone()
if existing_row:
socket_id = existing_row.socket_id
conn.execute(cl_worker.update().where(
and_(cl_worker.c.user_id == user_id, cl_worker.c.worker_id
== worker_id)).values(worker_row))
else:
socket_id = self.allocate_socket(user_id, worker_id, conn)
worker_row.update({
'user_id': user_id,
'worker_id': worker_id,
'socket_id': socket_id
})
conn.execute(cl_worker.insert().values(worker_row))
# Update dependencies
blob = self._serialize_dependencies(dependencies).encode()
if existing_row:
conn.execute(cl_worker_dependency.update().where(
and_(
cl_worker_dependency.c.user_id == user_id,
cl_worker_dependency.c.worker_id == worker_id,
)).values(dependencies=blob))
else:
conn.execute(cl_worker_dependency.insert().values(
user_id=user_id, worker_id=worker_id, dependencies=blob))
return socket_id
