Public Access
484 lines
18 KiB
Python
484 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""Distribute molecule (scenario, platform) pairs across N parallel runners.
|
|
|
|
Discovers all molecule scenarios under ansible/roles/*/molecule/ and
|
|
crosses them with the supported OS platform matrix, then splits the
|
|
resulting test pairs evenly across the requested number of runners.
|
|
|
|
Each pair is printed as ``scenario|platform_name|platform_image|platform_command``
|
|
so the CI workflow can set the appropriate environment variables.
|
|
|
|
Usage:
|
|
python3 -m devx.molecule.distribute_molecule --runner-index 1 --max-runners 3
|
|
# prints: default|ubuntu-2204|ubuntu:22.04| lifecycle|ubuntu-2204|ubuntu:22.04| ...
|
|
python3 -m devx.molecule.distribute_molecule --list
|
|
# prints all scenarios, one per line
|
|
python3 -m devx.molecule.distribute_molecule --list-platforms
|
|
# prints all platforms, one per line
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import tomllib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import click
|
|
|
|
from devx.ci._shared import lpt_distribute, write_github_env
|
|
from devx.i18n import _
|
|
from devx.molecule.platforms import PLATFORMS, load_platforms
|
|
|
|
DEFAULT_MAX_RUNNERS = 3
|
|
DEFAULT_ROLES_ROOT = Path("ansible/roles")
|
|
|
|
|
|
def _default_molecule_root() -> Path:
|
|
"""Auto-discover the single molecule directory under ansible/roles/.
|
|
|
|
If exactly one role has a molecule/ subdirectory, return it.
|
|
Otherwise, fall back to the first role with a molecule/ directory.
|
|
"""
|
|
roles_root = DEFAULT_ROLES_ROOT
|
|
if not roles_root.is_dir():
|
|
return roles_root / "gitea_runner" / "molecule" # sensible default for error message
|
|
mol_dirs = sorted(d / "molecule" for d in roles_root.iterdir() if (d / "molecule").is_dir())
|
|
if mol_dirs:
|
|
return mol_dirs[0]
|
|
return roles_root / "molecule" # will produce a clear "not found" error
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TestPair:
|
|
"""A (scenario, platform) combination to test."""
|
|
|
|
scenario: str
|
|
platform: dict[str, str]
|
|
|
|
def encode(self) -> str:
|
|
"""Serialize to a pipe-delimited string for CI consumption."""
|
|
cmd = self.platform["command"].replace(" ", "__SPACE__")
|
|
return f"{self.scenario}|{self.platform['name']}|{self.platform['image']}|{cmd}"
|
|
|
|
@staticmethod
|
|
def decode(encoded: str) -> TestPair:
|
|
"""Deserialize from a pipe-delimited string."""
|
|
parts = encoded.split("|")
|
|
return TestPair(
|
|
scenario=parts[0],
|
|
platform={"name": parts[1], "image": parts[2], "command": parts[3].replace("__SPACE__", " ")},
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MultiRoleTestPair:
|
|
"""A (role, scenario, platform) combination for multi-role projects."""
|
|
|
|
role: str
|
|
scenario: str
|
|
platform: dict[str, str]
|
|
|
|
def encode(self) -> str:
|
|
"""Serialize to a pipe-delimited string: ``role|scenario|platform_name|image|command``."""
|
|
cmd = self.platform["command"].replace(" ", "__SPACE__")
|
|
return f"{self.role}|{self.scenario}|{self.platform['name']}|{self.platform['image']}|{cmd}"
|
|
|
|
@staticmethod
|
|
def decode(encoded: str) -> MultiRoleTestPair:
|
|
"""Deserialize from a pipe-delimited string."""
|
|
parts = encoded.split("|")
|
|
return MultiRoleTestPair(
|
|
role=parts[0],
|
|
scenario=parts[1],
|
|
platform={"name": parts[2], "image": parts[3], "command": parts[4].replace("__SPACE__", " ")},
|
|
)
|
|
|
|
|
|
def discover_scenarios(root: Path | None = None) -> list[str]:
|
|
"""Return sorted list of molecule scenario directory names."""
|
|
if root is None:
|
|
root = _default_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 discover_multi_role_scenarios(
|
|
roles_root: Path | None = None,
|
|
include_roles: list[str] | None = None,
|
|
exclude_roles: list[str] | None = None,
|
|
) -> list[tuple[str, str]]:
|
|
"""Discover (role, scenario) pairs across all roles under *roles_root*.
|
|
|
|
Scans ``roles_root/*/molecule/*/`` for scenario directories, skipping
|
|
``common`` and directories starting with ``_``. Returns a sorted list of
|
|
``(role_name, scenario_name)`` tuples.
|
|
|
|
If *include_roles* is given, only roles whose name is in the list are
|
|
returned. If *exclude_roles* is given, roles whose name is in the list
|
|
are skipped. Both filters are case-insensitive.
|
|
"""
|
|
if roles_root is None:
|
|
roles_root = DEFAULT_ROLES_ROOT
|
|
if not roles_root.is_dir():
|
|
raise click.ClickException(_("Roles directory not found: {path}", path=str(roles_root)))
|
|
include_set = {r.lower() for r in include_roles} if include_roles else None
|
|
exclude_set = {r.lower() for r in exclude_roles} if exclude_roles else None
|
|
pairs: list[tuple[str, str]] = []
|
|
for role_dir in sorted(roles_root.iterdir()):
|
|
if not role_dir.is_dir():
|
|
continue
|
|
role_name = role_dir.name
|
|
if include_set is not None and role_name.lower() not in include_set:
|
|
continue
|
|
if exclude_set is not None and role_name.lower() in exclude_set:
|
|
continue
|
|
mol_dir = role_dir / "molecule"
|
|
if not mol_dir.is_dir():
|
|
continue
|
|
for scenario_dir in mol_dir.iterdir():
|
|
if not scenario_dir.is_dir():
|
|
continue
|
|
if scenario_dir.name.startswith("_") or scenario_dir.name == "common":
|
|
continue
|
|
pairs.append((role_dir.name, scenario_dir.name))
|
|
return pairs
|
|
|
|
|
|
def build_pairs(scenarios: list[str], platforms: list[dict[str, str]] | None = None) -> list[TestPair]:
|
|
"""Build the full cross-product of scenarios and platforms."""
|
|
if platforms is None:
|
|
platforms = PLATFORMS
|
|
return [TestPair(s, p) for s in scenarios for p in platforms]
|
|
|
|
|
|
def build_multi_role_pairs(
|
|
role_scenarios: list[tuple[str, str]],
|
|
platforms: list[dict[str, str]] | None = None,
|
|
) -> list[MultiRoleTestPair]:
|
|
"""Build the full cross-product of (role, scenario) pairs and platforms."""
|
|
if platforms is None:
|
|
platforms = PLATFORMS
|
|
return [MultiRoleTestPair(r, s, p) for r, s in role_scenarios for p in platforms]
|
|
|
|
|
|
# --- Molecule weight configuration ---
|
|
#
|
|
# Weights are loaded from ``[tool.devx.molecule.weights]`` in
|
|
# ``pyproject.toml``. Each project contributes its own
|
|
# weights calibrated from actual CI execution times.
|
|
#
|
|
# Two key formats are supported:
|
|
# - ``"scenario" = weight`` — applies to any role with that scenario name
|
|
# - ``"role/scenario" = weight`` — role-specific (takes priority)
|
|
#
|
|
# Example pyproject.toml::
|
|
#
|
|
# [tool.devx.molecule.weights]
|
|
# "nextcloud" = 15
|
|
# "app_container/customer-apps" = 11
|
|
# "restore/default" = 11
|
|
# "default" = 3
|
|
#
|
|
# If no configuration is found, a generic default weight is used for all
|
|
# scenarios (producing a round-robin distribution).
|
|
|
|
_DEFAULT_SCENARIO_WEIGHT = 3
|
|
|
|
|
|
def _load_molecule_weights(pyproject_path: str = "pyproject.toml") -> tuple[dict[str, int], dict[tuple[str, str], int]]:
|
|
"""Load molecule weights from ``[tool.devx.molecule.weights]`` in pyproject.toml.
|
|
|
|
Returns a tuple of ``(scenario_weights, role_scenario_weights)``:
|
|
- ``scenario_weights``: maps scenario name → weight (applies to any role)
|
|
- ``role_scenario_weights``: maps (role, scenario) → weight (role-specific)
|
|
"""
|
|
path = Path(pyproject_path)
|
|
if not path.exists():
|
|
return {}, {}
|
|
try:
|
|
with open(path, "rb") as f: # noqa: PTH123
|
|
data = tomllib.load(f)
|
|
except (tomllib.TOMLDecodeError, OSError):
|
|
return {}, {}
|
|
|
|
weights_raw = data.get("tool", {}).get("devx", {}).get("molecule", {}).get("weights", {})
|
|
if not isinstance(weights_raw, dict):
|
|
return {}, {}
|
|
|
|
scenario_weights: dict[str, int] = {}
|
|
role_scenario_weights: dict[tuple[str, str], int] = {}
|
|
|
|
for key, value in weights_raw.items():
|
|
if not isinstance(value, int):
|
|
continue
|
|
if "/" in key:
|
|
role, scenario = key.split("/", 1)
|
|
role_scenario_weights[(role.lower(), scenario.lower())] = value
|
|
else:
|
|
scenario_weights[key.lower()] = value
|
|
|
|
return scenario_weights, role_scenario_weights
|
|
|
|
|
|
# Load weights once at import time (like devx.config and classify_changes)
|
|
_SCENARIO_WEIGHTS, _ROLE_SCENARIO_WEIGHTS = _load_molecule_weights()
|
|
|
|
|
|
def _scenario_weight(scenario: str, role: str | None = None) -> int:
|
|
"""Estimate a weight for a scenario based on its name and optionally its role.
|
|
|
|
Role-specific weights (``"role/scenario"``) take priority over
|
|
scenario-name-only weights (``"scenario"``). Falls back to the
|
|
default weight if no configuration matches.
|
|
"""
|
|
s = scenario.lower()
|
|
if role is not None:
|
|
r = role.lower()
|
|
key = (r, s)
|
|
if key in _ROLE_SCENARIO_WEIGHTS:
|
|
return _ROLE_SCENARIO_WEIGHTS[key]
|
|
for key, weight in _SCENARIO_WEIGHTS.items():
|
|
if key in s:
|
|
return weight
|
|
return _DEFAULT_SCENARIO_WEIGHT
|
|
|
|
|
|
def _lpt_distribute[T](items: list[T], weights: list[int], max_runners: int) -> list[list[T]]:
|
|
"""Distribute *items* across *max_runners* using LPT (delegates to shared utility)."""
|
|
return lpt_distribute(items, weights, max_runners)
|
|
|
|
|
|
def distribute_multi_role(pairs: list[MultiRoleTestPair], max_runners: int) -> list[list[MultiRoleTestPair]]:
|
|
"""Split *pairs* into *max_runners* balanced groups using LPT scheduling.
|
|
|
|
Each pair is weighted by role+scenario heuristics (e.g. ``nextcloud`` is
|
|
heavier than ``simple-app``). Pairs are sorted by weight descending and
|
|
assigned to the runner with the least total weight.
|
|
"""
|
|
weights = [_scenario_weight(p.scenario, p.role) for p in pairs]
|
|
return _lpt_distribute(pairs, weights, max_runners)
|
|
|
|
|
|
def multi_role_pairs_for_runner(
|
|
pairs: list[MultiRoleTestPair], runner_index: int, max_runners: int
|
|
) -> list[MultiRoleTestPair]:
|
|
"""Return the subset of multi-role pairs assigned to *runner_index* (0-based)."""
|
|
groups = distribute_multi_role(pairs, 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]
|
|
|
|
|
|
def distribute(pairs: list[TestPair], max_runners: int) -> list[list[TestPair]]:
|
|
"""Split *pairs* into *max_runners* balanced groups using LPT scheduling.
|
|
|
|
Each pair is weighted by scenario name heuristics (e.g. ``nextcloud`` is
|
|
heavier than ``binary``). Pairs are sorted by weight descending and
|
|
assigned to the runner with the least total weight.
|
|
"""
|
|
weights = [_scenario_weight(p.scenario) for p in pairs]
|
|
return _lpt_distribute(pairs, weights, max_runners)
|
|
|
|
|
|
def pairs_for_runner(pairs: list[TestPair], runner_index: int, max_runners: int) -> list[TestPair]:
|
|
"""Return the subset of pairs assigned to *runner_index*."""
|
|
groups = distribute(pairs, 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]
|
|
|
|
|
|
def _write_github_env(key: str, value: str) -> None:
|
|
"""Append a key=value line to the $GITHUB_ENV file (delegates to shared utility)."""
|
|
write_github_env(key, value)
|
|
|
|
|
|
@click.command()
|
|
@click.option(
|
|
"--runner-index",
|
|
type=int,
|
|
default=None,
|
|
help="One-based runner index (Gitea Actions renders 0 as empty). "
|
|
"Converted to zero-based internally. 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.",
|
|
)
|
|
@click.option(
|
|
"--list-platforms",
|
|
"list_platforms",
|
|
is_flag=True,
|
|
help="List all supported platforms, one per line.",
|
|
)
|
|
@click.option(
|
|
"--github-env",
|
|
"github_env",
|
|
is_flag=True,
|
|
default=False,
|
|
help="Write TEST_PAIRS and SKIP to $GITHUB_ENV (for CI workflow steps).",
|
|
)
|
|
@click.option(
|
|
"--skip-if-excess",
|
|
is_flag=True,
|
|
default=False,
|
|
help="With --github-env: write SKIP=true when runner-index exceeds max-runners.",
|
|
)
|
|
@click.option(
|
|
"--molecule-root",
|
|
type=click.Path(exists=True, file_okay=False, path_type=Path),
|
|
default=None,
|
|
help="Custom molecule directory (single-role mode). Default: auto-discovered under ansible/roles/*/molecule.",
|
|
)
|
|
@click.option(
|
|
"--roles-root",
|
|
type=click.Path(exists=True, file_okay=False, path_type=Path),
|
|
default=None,
|
|
help="Roles directory for multi-role discovery (scans */molecule/*/). "
|
|
"Use this for projects with multiple Ansible roles. Default: disabled (single-role mode).",
|
|
)
|
|
@click.option(
|
|
"--platforms-file",
|
|
type=click.Path(exists=True, file_okay=True, path_type=Path),
|
|
default=None,
|
|
help="JSON file with custom platform list (each entry: name, image, command). "
|
|
"Overrides the default platform matrix. Useful for projects with custom test images.",
|
|
)
|
|
@click.option(
|
|
"--include-roles",
|
|
"include_roles",
|
|
type=str,
|
|
default=None,
|
|
help="Comma-separated list of role names to include (multi-role mode only). "
|
|
"Only scenarios from these roles are distributed. Case-insensitive. "
|
|
"Example: --include-roles docker_base,crowdsec,disk_cleanup,app_hardening",
|
|
)
|
|
@click.option(
|
|
"--exclude-roles",
|
|
"exclude_roles",
|
|
type=str,
|
|
default=None,
|
|
help="Comma-separated list of role names to exclude (multi-role mode only). "
|
|
"Scenarios from these roles are skipped. Case-insensitive. "
|
|
"Example: --exclude-roles docker_base,crowdsec,disk_cleanup,app_hardening",
|
|
)
|
|
def cli(
|
|
runner_index: int | None,
|
|
max_runners: int,
|
|
list_all: bool,
|
|
list_platforms: bool,
|
|
github_env: bool,
|
|
skip_if_excess: bool,
|
|
molecule_root: Path | None,
|
|
roles_root: Path | None,
|
|
platforms_file: Path | None,
|
|
include_roles: str | None,
|
|
exclude_roles: str | None,
|
|
) -> None:
|
|
platforms = load_platforms(platforms_file)
|
|
# Parse role filters
|
|
include_list = [r.strip() for r in include_roles.split(",")] if include_roles else None
|
|
exclude_list = [r.strip() for r in exclude_roles.split(",")] if exclude_roles else None
|
|
# Multi-role mode: discover (role, scenario) pairs across all roles
|
|
if roles_root is not None:
|
|
role_scenarios = discover_multi_role_scenarios(
|
|
roles_root, include_roles=include_list, exclude_roles=exclude_list
|
|
)
|
|
if list_all:
|
|
for role, scenario in role_scenarios:
|
|
click.echo(f"{role}|{scenario}")
|
|
return
|
|
if list_platforms:
|
|
for p in platforms:
|
|
click.echo(f"{p['name']}|{p['image']}|{p['command']}")
|
|
return
|
|
pairs_mr = build_multi_role_pairs(role_scenarios, platforms)
|
|
if runner_index is None:
|
|
groups = distribute_multi_role(pairs_mr, max_runners)
|
|
for i, group in enumerate(groups):
|
|
labels = " ".join(p.encode() for p in group) if group else "(none)"
|
|
click.echo(f"Runner {i}: {labels}")
|
|
return
|
|
if skip_if_excess and github_env and runner_index > max_runners:
|
|
click.echo(f"Skipping — runner index {runner_index} > max runners {max_runners}")
|
|
_write_github_env("TEST_PAIRS", "")
|
|
_write_github_env("SKIP", "true")
|
|
return
|
|
if runner_index < 1:
|
|
raise click.ClickException(f"Runner index {runner_index} is out of range (must be >= 1)")
|
|
zero_based = runner_index - 1
|
|
assigned = multi_role_pairs_for_runner(pairs_mr, zero_based, max_runners)
|
|
encoded = " ".join(p.encode() for p in assigned)
|
|
if github_env:
|
|
_write_github_env("TEST_PAIRS", encoded)
|
|
_write_github_env("SKIP", "false")
|
|
click.echo(f"Assigned pairs: {encoded}")
|
|
return
|
|
click.echo(encoded)
|
|
return
|
|
|
|
# Single-role mode (default or --molecule-root)
|
|
scenarios = discover_scenarios(molecule_root)
|
|
if list_all:
|
|
for s in scenarios:
|
|
click.echo(s)
|
|
return
|
|
if list_platforms:
|
|
for p in platforms:
|
|
click.echo(f"{p['name']}|{p['image']}|{p['command']}")
|
|
return
|
|
pairs = build_pairs(scenarios, platforms)
|
|
if runner_index is None:
|
|
groups = distribute(pairs, max_runners)
|
|
for i, group in enumerate(groups):
|
|
labels = " ".join(p.encode() for p in group) if group else "(none)"
|
|
click.echo(f"Runner {i}: {labels}")
|
|
return
|
|
|
|
# Skip if runner index exceeds available runners (CI static matrix has 3 slots)
|
|
if skip_if_excess and github_env and runner_index > max_runners:
|
|
click.echo(f"Skipping — runner index {runner_index} > max runners {max_runners}")
|
|
_write_github_env("TEST_PAIRS", "")
|
|
_write_github_env("SKIP", "true")
|
|
return
|
|
|
|
# Validate runner index is in range
|
|
if runner_index < 1:
|
|
raise click.ClickException(f"Runner index {runner_index} is out of range (must be >= 1)")
|
|
|
|
# Convert 1-based CLI index to 0-based internal index
|
|
zero_based = runner_index - 1
|
|
assigned = pairs_for_runner(pairs, zero_based, max_runners)
|
|
encoded = " ".join(p.encode() for p in assigned)
|
|
|
|
if github_env:
|
|
_write_github_env("TEST_PAIRS", encoded)
|
|
_write_github_env("SKIP", "false")
|
|
click.echo(f"Assigned pairs: {encoded}")
|
|
return
|
|
|
|
click.echo(encoded)
|
|
|
|
|
|
if __name__ == "__main__": # pragma: no cover
|
|
cli()
|