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
Exemple #2
0
    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
Exemple #3
0
    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
Exemple #4
0
    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
Exemple #5
0
    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