104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Distribute molecule scenarios across N parallel runners.
|
|
|
|
Discovers all molecule scenarios under ansible/roles/*/molecule/ and
|
|
splits them evenly across the requested number of runners.
|
|
|
|
Usage:
|
|
python3 scripts/distribute_molecule.py --runner-index 0 --max-runners 3
|
|
# prints: default deregister
|
|
python3 scripts/distribute_molecule.py --list
|
|
# prints all scenarios, one per line
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import click
|
|
|
|
from gitea_runner_manager.i18n import _
|
|
|
|
DEFAULT_MAX_RUNNERS = 3
|
|
MOLECULE_ROOT = Path("ansible/roles/gitea-runner/molecule")
|
|
|
|
|
|
def discover_scenarios(root: Path | None = None) -> list[str]:
|
|
"""Return sorted list of molecule scenario directory names."""
|
|
if root is None:
|
|
root = MOLECULE_ROOT
|
|
if not root.is_dir():
|
|
raise click.ClickException(
|
|
_("Molecule directory not found: {path}", path=str(root))
|
|
)
|
|
scenarios = [
|
|
d.name
|
|
for d in root.iterdir()
|
|
if d.is_dir() and not d.name.startswith("_") and d.name != "common"
|
|
]
|
|
return sorted(scenarios)
|
|
|
|
|
|
def distribute(scenarios: list[str], max_runners: int) -> list[list[str]]:
|
|
"""Split *scenarios* into *max_runners* balanced groups (round-robin)."""
|
|
groups: list[list[str]] = [[] for _ in range(max_runners)]
|
|
for i, scenario in enumerate(scenarios):
|
|
groups[i % max_runners].append(scenario)
|
|
return groups
|
|
|
|
|
|
def scenarios_for_runner(
|
|
scenarios: list[str], runner_index: int, max_runners: int
|
|
) -> list[str]:
|
|
"""Return the subset of scenarios assigned to *runner_index*."""
|
|
groups = distribute(scenarios, max_runners)
|
|
if runner_index < 0 or runner_index >= len(groups):
|
|
raise click.ClickException(
|
|
_(
|
|
"Runner index {index} out of range (0..{max})",
|
|
index=runner_index,
|
|
max=max_runners - 1,
|
|
)
|
|
)
|
|
return groups[runner_index]
|
|
|
|
|
|
@click.command()
|
|
@click.option(
|
|
"--runner-index",
|
|
type=int,
|
|
default=None,
|
|
help="Zero-based runner index. If omitted, prints all groups.",
|
|
)
|
|
@click.option(
|
|
"--max-runners",
|
|
type=int,
|
|
default=DEFAULT_MAX_RUNNERS,
|
|
show_default=True,
|
|
help="Total number of parallel runners.",
|
|
)
|
|
@click.option(
|
|
"--list",
|
|
"list_all",
|
|
is_flag=True,
|
|
help="List all discovered scenarios, one per line.",
|
|
)
|
|
def cli(runner_index: int | None, max_runners: int, list_all: bool) -> None:
|
|
scenarios = discover_scenarios()
|
|
if list_all:
|
|
for s in scenarios:
|
|
click.echo(s)
|
|
return
|
|
if runner_index is None:
|
|
groups = distribute(scenarios, max_runners)
|
|
for i, group in enumerate(groups):
|
|
click.echo(f"Runner {i}: {' '.join(group) if group else '(none)'}")
|
|
return
|
|
assigned = scenarios_for_runner(scenarios, runner_index, max_runners)
|
|
click.echo(" ".join(assigned))
|
|
|
|
|
|
if __name__ == "__main__": # pragma: no cover
|
|
cli() # pragma: no cover
|