def worker_cleanup(self, user_id, worker_id):
"""
Deletes the worker and all associated data from the database as well
as the socket directory.
"""
with self._engine.begin() as conn:
socket_rows = conn.execute(cl_worker_socket.select().where(
and_(
cl_worker_socket.c.user_id == user_id,
cl_worker_socket.c.worker_id == worker_id,
))).fetchall()
for socket_row in socket_rows:
self._cleanup_socket(socket_row.socket_id)
conn.execute(cl_worker_socket.delete().where(
and_(
cl_worker_socket.c.user_id == user_id,
cl_worker_socket.c.worker_id == worker_id,
)))
conn.execute(cl_worker_run.delete().where(
and_(cl_worker_run.c.user_id == user_id,
cl_worker_run.c.worker_id == worker_id)))
conn.execute(cl_worker_dependency.delete().where(
and_(
cl_worker_dependency.c.user_id == user_id,
cl_worker_dependency.c.worker_id == worker_id,
)))
conn.execute(cl_worker.delete().where(
and_(cl_worker.c.user_id == user_id,
cl_worker.c.worker_id == worker_id)))
def worker_cleanup(self, user_id, worker_id):
"""
Deletes the worker and all associated data from the database as well
as the socket directory.
"""
with self._engine.begin() as conn:
socket_rows = conn.execute(
cl_worker_socket.select()
.where(and_(cl_worker_socket.c.user_id == user_id,
cl_worker_socket.c.worker_id == worker_id))
).fetchall()
for socket_row in socket_rows:
self._cleanup_socket(socket_row.socket_id)
conn.execute(cl_worker_socket.delete()
.where(and_(cl_worker_socket.c.user_id == user_id,
cl_worker_socket.c.worker_id == worker_id)))
conn.execute(cl_worker_run.delete()
.where(and_(cl_worker_run.c.user_id == user_id,
cl_worker_run.c.worker_id == worker_id)))
conn.execute(cl_worker_dependency.delete()
.where(and_(cl_worker_dependency.c.user_id == user_id,
cl_worker_dependency.c.worker_id == worker_id)))
conn.execute(cl_worker.delete()
.where(and_(cl_worker.c.user_id == user_id,
cl_worker.c.worker_id == worker_id)))
def worker_checkin(self, user_id, worker_id, tag, slots, cpus, memory_bytes, dependencies):
"""
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,
'slots': slots,
'cpus': cpus,
'memory_bytes': memory_bytes,
'checkin_time': datetime.datetime.now(),
}
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))
conn.execute(
cl_worker_dependency.delete()
.where(and_(cl_worker_dependency.c.user_id == user_id,
cl_worker_dependency.c.worker_id == worker_id)))
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))
dependency_rows = [{
'user_id': user_id,
'worker_id': worker_id,
'dependency_uuid': uuid,
'dependency_path': path,
} for uuid, path in dependencies]
if dependency_rows:
conn.execute(cl_worker_dependency.insert(), dependency_rows)
return socket_id
def worker_checkin(self, user_id, worker_id, slots, dependency_uuids):
"""
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 = {
'slots': slots,
'checkin_time': datetime.datetime.now(),
}
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))
conn.execute(
cl_worker_dependency.delete()
.where(and_(cl_worker_dependency.c.user_id == user_id,
cl_worker_dependency.c.worker_id == worker_id)))
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))
dependency_rows = [{
'user_id': user_id,
'worker_id': worker_id,
'dependency_uuid': uuid,
} for uuid in dependency_uuids]
if dependency_rows:
conn.execute(cl_worker_dependency.insert(), dependency_rows)
return socket_id
