Metadata-Version: 2.4
Name: airflow-balancer
Version: 0.7.9
Summary: Utilities for tracking hosts and ports and load balancing DAGs
Project-URL: Repository, https://github.com/airflow-laminar/airflow-balancer
Project-URL: Homepage, https://github.com/airflow-laminar/airflow-balancer
Author-email: the airflow-balancer authors <t.paine154@gmail.com>
License: Apache-2.0
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: Implementation :: CPython
Classifier: Programming Language :: Python :: Implementation :: PyPy
Requires-Python: >=3.11
Requires-Dist: airflow-config<1.13,>=1.12
Requires-Dist: airflow-pydantic<1.7,>=1.6
Requires-Dist: pydantic>=2
Provides-Extra: airflow
Requires-Dist: apache-airflow-providers-ssh; extra == 'airflow'
Requires-Dist: apache-airflow-providers-standard; extra == 'airflow'
Requires-Dist: apache-airflow<3,>=2.8; extra == 'airflow'
Provides-Extra: airflow3
Requires-Dist: apache-airflow-providers-ssh; extra == 'airflow3'
Requires-Dist: apache-airflow-providers-standard; extra == 'airflow3'
Requires-Dist: apache-airflow<3.4,>=3; extra == 'airflow3'
Provides-Extra: develop
Requires-Dist: build; extra == 'develop'
Requires-Dist: bump-my-version; extra == 'develop'
Requires-Dist: check-dist; extra == 'develop'
Requires-Dist: codespell; extra == 'develop'
Requires-Dist: email-validator>=2; extra == 'develop'
Requires-Dist: fastapi; extra == 'develop'
Requires-Dist: flask; extra == 'develop'
Requires-Dist: flask-appbuilder>=4.5.4; extra == 'develop'
Requires-Dist: flask-limiter<3.13; extra == 'develop'
Requires-Dist: hatch-js; extra == 'develop'
Requires-Dist: hatchling; extra == 'develop'
Requires-Dist: mdformat; extra == 'develop'
Requires-Dist: mdformat-tables>=1; extra == 'develop'
Requires-Dist: pytest; extra == 'develop'
Requires-Dist: pytest-cov; extra == 'develop'
Requires-Dist: ruff; extra == 'develop'
Requires-Dist: twine; extra == 'develop'
Requires-Dist: ty; extra == 'develop'
Requires-Dist: uv; extra == 'develop'
Requires-Dist: uvicorn; extra == 'develop'
Requires-Dist: wheel; extra == 'develop'
Description-Content-Type: text/markdown

# airflow balancer

Utilities for tracking hosts and ports and load balancing DAGs

[![Build Status](https://github.com/airflow-laminar/airflow-balancer/actions/workflows/build.yaml/badge.svg?branch=main&event=push)](https://github.com/airflow-laminar/airflow-balancer/actions/workflows/build.yaml)
[![codecov](https://codecov.io/gh/airflow-laminar/airflow-balancer/branch/main/graph/badge.svg)](https://codecov.io/gh/airflow-laminar/airflow-balancer)
[![License](https://img.shields.io/github/license/airflow-laminar/airflow-balancer)](https://github.com/airflow-laminar/airflow-balancer)
[![PyPI](https://img.shields.io/pypi/v/airflow-balancer.svg)](https://pypi.python.org/pypi/airflow-balancer)

## Overview

`airflow-balancer` is a utility library for Apache Airflow to track host and port usage via yaml files. It enables you to:

- **Track Hosts**: Define and manage a pool of worker hosts with different capabilities (OS, queues, tags)
- **Manage Ports**: Track port usage across your host infrastructure to avoid conflicts
- **Load Balance**: Intelligently select hosts based on queues, operating systems, tags, or custom criteria
- **Integrate with Airflow**: Automatically create Airflow pools for each host and port

### Integration with airflow-laminar Stack

`airflow-balancer` is tightly integrated with the [airflow-laminar](https://github.com/airflow-laminar) ecosystem:

| Library                                                                 | Integration                                                                                                                                                                          |
| ----------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
| [airflow-pydantic](https://github.com/airflow-laminar/airflow-pydantic) | Core data models (`Host`, `Port`, `BalancerConfiguration`) are defined in `airflow-pydantic`, providing full Pydantic validation, type checking, and JSON/YAML serialization support |
| [airflow-config](https://github.com/airflow-laminar/airflow-config)     | Configuration loading via [Hydra](https://hydra.cc) for hierarchical configs with defaults, overrides, and environment-specific settings                                             |

With `airflow-balancer`, you can register host and port usage in configuration:

```yaml
_target_: airflow_balancer.BalancerConfiguration
default_username: timkpaine
hosts:
  - name: host1
    size: 16
    os: ubuntu
    queues: [primary]

  - name: host2
    os: ubuntu
    size: 16
    queues: [workers]

  - name: host3
    os: macos
    size: 8
    queues: [workers]

ports:
  - host: host1
    port: 8080

  - host_name: host2
    port: 8793
```

Either via `airflow-config` or directly, you can then select amongst available hosts for use in your DAGs.

```python
from airflow_balaner import BalancerConfiguration, load

balancer_config: BalancerConfiguration = load("balancer.yaml")

host = balancer_config.select_host(queue="workers")
port = balancer_config.free_port(host=host)

...

operator = SSHOperator(ssh_hook=host.hook(), ...)

```

### Visualization

Configuration, and Host and Port listing is built into the extension, available either from the topbar in Airflow or as a standalone viewer (via the `airflow-balancer-viewer` CLI).

<img src="https://raw.githubusercontent.com/airflow-laminar/airflow-balancer/refs/heads/main/docs/img/toolbar.png" width=400>

<img src="https://raw.githubusercontent.com/airflow-laminar/airflow-balancer/refs/heads/main/docs/img/home.png" width=400>

<img src="https://raw.githubusercontent.com/airflow-laminar/airflow-balancer/refs/heads/main/docs/img/hosts.png" width=800>

## Installation

You can install from pip:

```bash
pip install airflow-balancer
```

For use with Apache Airflow 2.x:

```bash
pip install airflow-balancer[airflow]
```

For use with Apache Airflow 3.x:

```bash
pip install airflow-balancer[airflow3]
```

Or via conda:

```bash
conda install airflow-balancer -c conda-forge
```

## Using with airflow-config

The recommended approach is to use `airflow-balancer` as an extension within your `airflow-config` configuration:

```yaml
# config/config.yaml
# @package _global_
_target_: airflow_config.Configuration
defaults:
  - extensions/balancer@extensions.balancer
```

```yaml
# config/extensions/balancer.yaml
# @package extensions.balancer
_target_: airflow_balancer.BalancerConfiguration

default_username: airflow
default_key_file: /home/airflow/.ssh/id_rsa
hosts:
  - name: worker1
    size: 16
    os: ubuntu
    queues: [workers]
```

```python
from airflow_config import load_config

config = load_config("config", "config")
balancer = config.extensions["balancer"]

# Select a host and use its SSH hook
host = balancer.select_host(queue="workers")
operator = SSHOperator(ssh_hook=host.hook(), ...)
```

## Using with airflow-pydantic

Since the core models are defined in `airflow-pydantic`, you can leverage its testing utilities:

```python
from airflow_balancer import BalancerConfiguration, Host
from airflow_balancer.testing import pools, variables
from airflow_pydantic import Variable

# Testing with mocked pools
with pools():
    config = BalancerConfiguration(
        hosts=[Host(name="test-host", size=8, queues=["test"])]
    )
    assert config.select_host(queue="test").name == "test-host"

# Using Airflow Variables for credentials
host = Host(
    name="secure-host",
    username="admin",
    password=Variable(key="host_password"),
)
```

## License

This software is licensed under the Apache 2.0 license. See the [LICENSE](LICENSE) file for details.

> [!NOTE]
> This library was generated using [copier](https://copier.readthedocs.io/en/stable/) from the [Base Python Project Template repository](https://github.com/python-project-templates/base).
