def get_workers(self): """ Returns information about all the workers in the database. The return value is a list of dicts with the structure shown in the code below. """ with self._engine.begin() as conn: worker_rows = conn.execute(cl_worker.select()).fetchall() worker_run_rows = ( conn.execute(cl_worker_run.select()).fetchall()) worker_dependency_rows = ( conn.execute(cl_worker_dependency.select()).fetchall()) worker_dict = {(row.user_id, row.worker_id): { 'user_id': row.user_id, 'worker_id': row.worker_id, 'slots': row.slots, 'checkin_time': row.checkin_time, 'socket_id': row.socket_id, 'run_uuids': [], 'dependency_uuids': [], } for row in worker_rows} for row in worker_run_rows: worker_dict[(row.user_id, row.worker_id)]['run_uuids'].append(row.run_uuid) for row in worker_dependency_rows: worker_dict[(row.user_id, row.worker_id)]['dependency_uuids'].append(row.dependency_uuid) return worker_dict.values()
def worker_checkin( self, user_id, worker_id, tag, cpus, gpus, memory_bytes, free_disk_bytes, dependencies, shared_file_system, tag_exclusive, ): """ 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, } 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
def worker_checkin(self, user_id, worker_id, tag, cpus, gpus, 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, 'cpus': cpus, 'gpus': gpus, '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) ) 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) 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
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 get_bundle_worker(self, uuid): """ Returns information about the worker that the given bundle is running on. This method should be called only for bundles that are running. """ with self._engine.begin() as connection: row = connection.execute(cl_worker_run.select() .where(cl_worker_run.c.run_uuid == uuid)).fetchone() precondition(row, 'Trying to find worker for bundle that is not running.') worker_row = connection.execute(cl_worker.select() .where(and_(cl_worker.c.user_id == row.user_id, cl_worker.c.worker_id == row.worker_id))).fetchone() return { 'user_id': worker_row.user_id, 'worker_id': worker_row.worker_id, 'socket_id': worker_row.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
def update_workers(self, user_id, worker_id, update): """ Update the designated worker with columns and values :param user_id: a user_id indicating whom a worker belongs to :param worker_id: a worker's identification number :param update: a dictionary of (key, value) pairs that specifies the columns and the values to update """ if not update: return with self._engine.begin() as conn: 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: conn.execute(cl_worker.update().where( and_(cl_worker.c.user_id == user_id, cl_worker.c.worker_id == worker_id)).values(update))
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