Source code for airflow.models.pool

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

from typing import Iterable

from sqlalchemy import Column, Integer, String, Text, func
from sqlalchemy.orm.session import Session

from airflow.exceptions import AirflowException, PoolNotFound
from airflow.models.base import Base
from airflow.ti_deps.dependencies_states import EXECUTION_STATES
from airflow.typing_compat import TypedDict
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.sqlalchemy import nowait, with_row_locks
from airflow.utils.state import State


[docs]class PoolStats(TypedDict): """Dictionary containing Pool Stats"""
[docs] total: int
[docs] running: int
[docs] queued: int
[docs] open: int
[docs]class Pool(Base): """the class to get Pool info."""
[docs] __tablename__ = "slot_pool"
[docs] id = Column(Integer, primary_key=True)
[docs] pool = Column(String(256), unique=True)
# -1 for infinite
[docs] slots = Column(Integer, default=0)
[docs] description = Column(Text)
[docs] DEFAULT_POOL_NAME = "default_pool"
[docs] def __repr__(self): return str(self.pool)
@staticmethod @provide_session
[docs] def get_pools(session: Session = NEW_SESSION): """Get all pools.""" return session.query(Pool).all()
@staticmethod @provide_session
[docs] def get_pool(pool_name: str, session: Session = NEW_SESSION): """ Get the Pool with specific pool name from the Pools. :param pool_name: The pool name of the Pool to get. :param session: SQLAlchemy ORM Session :return: the pool object """ return session.query(Pool).filter(Pool.pool == pool_name).first()
@staticmethod @provide_session
[docs] def get_default_pool(session: Session = NEW_SESSION): """ Get the Pool of the default_pool from the Pools. :param session: SQLAlchemy ORM Session :return: the pool object """ return Pool.get_pool(Pool.DEFAULT_POOL_NAME, session=session)
@staticmethod @provide_session
[docs] def is_default_pool(id: int, session: Session = NEW_SESSION) -> bool: """ Check id if is the default_pool. :param id: pool id :param session: SQLAlchemy ORM Session :return: True if id is default_pool, otherwise False """ return ( session.query(func.count(Pool.id)) .filter(Pool.id == id, Pool.pool == Pool.DEFAULT_POOL_NAME) .scalar() > 0
) @staticmethod @provide_session
[docs] def create_or_update_pool(name: str, slots: int, description: str, session: Session = NEW_SESSION): """Create a pool with given parameters or update it if it already exists.""" if not name: return pool = session.query(Pool).filter_by(pool=name).first() if pool is None: pool = Pool(pool=name, slots=slots, description=description) session.add(pool) else: pool.slots = slots pool.description = description session.commit() return pool
@staticmethod @provide_session
[docs] def delete_pool(name: str, session: Session = NEW_SESSION): """Delete pool by a given name.""" if name == Pool.DEFAULT_POOL_NAME: raise AirflowException(f"{Pool.DEFAULT_POOL_NAME} cannot be deleted") pool = session.query(Pool).filter_by(pool=name).first() if pool is None: raise PoolNotFound(f"Pool '{name}' doesn't exist") session.delete(pool) session.commit() return pool
@staticmethod @provide_session
[docs] def slots_stats( *, lock_rows: bool = False, session: Session = NEW_SESSION, ) -> dict[str, PoolStats]: """ Get Pool stats (Number of Running, Queued, Open & Total tasks) If ``lock_rows`` is True, and the database engine in use supports the ``NOWAIT`` syntax, then a non-blocking lock will be attempted -- if the lock is not available then SQLAlchemy will throw an OperationalError. :param lock_rows: Should we attempt to obtain a row-level lock on all the Pool rows returns :param session: SQLAlchemy ORM Session """ from airflow.models.taskinstance import TaskInstance # Avoid circular import pools: dict[str, PoolStats] = {} query = session.query(Pool.pool, Pool.slots) if lock_rows: query = with_row_locks(query, session=session, **nowait(session)) pool_rows: Iterable[tuple[str, int]] = query.all() for (pool_name, total_slots) in pool_rows: if total_slots == -1: total_slots = float("inf") # type: ignore pools[pool_name] = PoolStats(total=total_slots, running=0, queued=0, open=0) state_count_by_pool = ( session.query(TaskInstance.pool, TaskInstance.state, func.sum(TaskInstance.pool_slots)) .filter(TaskInstance.state.in_(list(EXECUTION_STATES))) .group_by(TaskInstance.pool, TaskInstance.state) ).all() # calculate queued and running metrics for (pool_name, state, count) in state_count_by_pool: # Some databases return decimal.Decimal here. count = int(count) stats_dict: PoolStats | None = pools.get(pool_name) if not stats_dict: continue # TypedDict key must be a string literal, so we use if-statements to set value if state == "running": stats_dict["running"] = count elif state == "queued": stats_dict["queued"] = count else: raise AirflowException(f"Unexpected state. Expected values: {EXECUTION_STATES}.") # calculate open metric for pool_name, stats_dict in pools.items(): stats_dict["open"] = stats_dict["total"] - stats_dict["running"] - stats_dict["queued"] return pools
[docs] def to_json(self): """ Get the Pool in a json structure :return: the pool object in json format """ return { "id": self.id, "pool": self.pool, "slots": self.slots, "description": self.description,
} @provide_session
[docs] def occupied_slots(self, session: Session = NEW_SESSION): """ Get the number of slots used by running/queued tasks at the moment. :param session: SQLAlchemy ORM Session :return: the used number of slots """ from airflow.models.taskinstance import TaskInstance # Avoid circular import return int( session.query(func.sum(TaskInstance.pool_slots)) .filter(TaskInstance.pool == self.pool) .filter(TaskInstance.state.in_(list(EXECUTION_STATES))) .scalar() or 0
) @provide_session
[docs] def running_slots(self, session: Session = NEW_SESSION): """ Get the number of slots used by running tasks at the moment. :param session: SQLAlchemy ORM Session :return: the used number of slots """ from airflow.models.taskinstance import TaskInstance # Avoid circular import return int( session.query(func.sum(TaskInstance.pool_slots)) .filter(TaskInstance.pool == self.pool) .filter(TaskInstance.state == State.RUNNING) .scalar() or 0
) @provide_session
[docs] def queued_slots(self, session: Session = NEW_SESSION): """ Get the number of slots used by queued tasks at the moment. :param session: SQLAlchemy ORM Session :return: the used number of slots """ from airflow.models.taskinstance import TaskInstance # Avoid circular import return int( session.query(func.sum(TaskInstance.pool_slots)) .filter(TaskInstance.pool == self.pool) .filter(TaskInstance.state == State.QUEUED) .scalar() or 0
) @provide_session
[docs] def scheduled_slots(self, session: Session = NEW_SESSION): """ Get the number of slots scheduled at the moment. :param session: SQLAlchemy ORM Session :return: the number of scheduled slots """ from airflow.models.taskinstance import TaskInstance # Avoid circular import return int( session.query(func.sum(TaskInstance.pool_slots)) .filter(TaskInstance.pool == self.pool) .filter(TaskInstance.state == State.SCHEDULED) .scalar() or 0
) @provide_session
[docs] def open_slots(self, session: Session = NEW_SESSION) -> float: """ Get the number of slots open at the moment. :param session: SQLAlchemy ORM Session :return: the number of slots """ if self.slots == -1: return float("inf") else: return self.slots - self.occupied_slots(session)

Was this entry helpful?