diff --git a/AGENTS.md b/AGENTS.md index 2ba064c..62d663c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -68,6 +68,7 @@ src/devx/ │ ├── push_badges.py # Generate and push quality badges (--retries for retry on git push failures) │ ├── notify_failure.py # Create Gitea issues on CI failures (--auto-login) │ ├── distribute_files.py # Distribute files across parallel runners (LPT scheduling) +│ ├── distribute_items.py # Distribute generic items (VMs, hosts) across parallel runners (LPT) │ ├── integration_guard.py # Run pytest with cross-runner fail-fast │ ├── check_translations.py # Translation completeness check │ └── doc_coverage.py # Documentation coverage check diff --git a/src/devx/ci/distribute_items.py b/src/devx/ci/distribute_items.py new file mode 100644 index 0000000..ffb8230 --- /dev/null +++ b/src/devx/ci/distribute_items.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +"""Distribute a list of items across N parallel runners using LPT scheduling. + +Generic item distribution for CI matrix jobs. Items are read from a JSON +array on stdin (or from a file via --items-file), sorted for deterministic +ordering, then assigned to *max_runners* groups using LPT (Longest +Processing Time first) scheduling. + +Each item is a string (e.g. an Ansible ``--limit`` pattern like +``observability`` or ``infra-314-vm``). Optionally, items can be objects +with ``{"id": "...", "weight": N}`` to provide explicit weights. + +The assigned group for *runner_index* is written to ``$GITHUB_ENV`` as +``ASSIGNED_ITEMS`` (space-delimited) for use by subsequent steps. + +Usage:: + + echo '["observability", "infra-314-vm"]' | \\ + python3 -m devx.ci.distribute_items \\ + --runner-index 1 --max-runners 3 \\ + --github-env --skip-if-excess + + # With weights: + echo '[{"id": "observability", "weight": 5}, {"id": "customer-1", "weight": 3}]' | \\ + python3 -m devx.ci.distribute_items \\ + --runner-index 1 --max-runners 3 --github-env +""" + +from __future__ import annotations + +import json +import os +import sys + +import click + +from devx.i18n import _ + +DEFAULT_MAX_RUNNERS = 3 +DEFAULT_WEIGHT = 1 + + +def parse_items(raw: str) -> list[str]: + """Parse a JSON array into a list of item identifier strings. + + Accepts both plain string arrays (``["a", "b"]``) and object arrays + (``[{"id": "a", "weight": 2}]``). Returns just the identifier strings. + """ + data = json.loads(raw) + if not isinstance(data, list): + raise click.ClickException(_("Items input must be a JSON array, got {type}", type=type(data).__name__)) + items: list[str] = [] + for entry in data: + if isinstance(entry, str): + items.append(entry) + elif isinstance(entry, dict) and "id" in entry: + items.append(str(entry["id"])) + else: + raise click.ClickException( + _("Each item must be a string or an object with 'id', got {type}", type=type(entry).__name__) + ) + return items + + +def parse_weighted_items(raw: str) -> tuple[list[str], list[int]]: + """Parse a JSON array into (items, weights) lists. + + For plain string arrays, all items get ``DEFAULT_WEIGHT``. + For object arrays, the ``weight`` field is used (default: ``DEFAULT_WEIGHT``). + """ + data = json.loads(raw) + if not isinstance(data, list): + raise click.ClickException(_("Items input must be a JSON array, got {type}", type=type(data).__name__)) + items: list[str] = [] + weights: list[int] = [] + for entry in data: + if isinstance(entry, str): + items.append(entry) + weights.append(DEFAULT_WEIGHT) + elif isinstance(entry, dict) and "id" in entry: + items.append(str(entry["id"])) + weights.append(int(entry.get("weight", DEFAULT_WEIGHT))) + else: + raise click.ClickException( + _("Each item must be a string or an object with 'id', got {type}", type=type(entry).__name__) + ) + return items, weights + + +def distribute(items: list[str], weights: list[int], max_runners: int) -> list[list[str]]: + """Split *items* into *max_runners* balanced groups using LPT scheduling. + + Items are sorted by weight (descending), then assigned to the runner + with the least total weight. + """ + groups: list[list[str]] = [[] for _ in range(max_runners)] + loads = [0] * max_runners + indexed = sorted(enumerate(items), key=lambda x: (-weights[x[0]], x[0])) + for orig_idx, item in indexed: + min_runner = min(range(max_runners), key=lambda r: loads[r]) + groups[min_runner].append(item) + loads[min_runner] += weights[orig_idx] + return groups + + +def items_for_runner(items: list[str], weights: list[int], runner_index: int, max_runners: int) -> list[str]: + """Return the subset of items assigned to *runner_index* (0-based).""" + groups = distribute(items, weights, 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: + gh_env = os.environ.get("GITHUB_ENV") + if not gh_env: + raise click.ClickException("GITHUB_ENV environment variable is not set") + with open(gh_env, "a") as f: # noqa: PTH123 + if "\n" in value: + delimiter = "EOF" + f.write(f"{key}<<{delimiter}\n{value}\n{delimiter}\n") + else: + f.write(f"{key}={value}\n") + + +@click.command() +@click.option( + "--items-file", + type=click.Path(exists=True, file_okay=True, path_type=None), + default=None, + help="Read items from a JSON file instead of stdin.", +) +@click.option( + "--runner-index", + type=int, + default=None, + help="One-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( + "--github-env", + is_flag=True, + default=False, + help="Write ASSIGNED_ITEMS and SKIP to $GITHUB_ENV.", +) +@click.option( + "--skip-if-excess", + is_flag=True, + default=False, + help="With --github-env: write SKIP=true when runner-index exceeds max-runners.", +) +def main( + items_file: str | None, + runner_index: int | None, + max_runners: int, + github_env: bool, + skip_if_excess: bool, +) -> None: + # Read items from file or stdin + if items_file is not None: + with open(items_file) as f: # noqa: PTH123 + raw = f.read() + else: + raw = sys.stdin.read() + + raw = raw.strip() + if not raw: + raw = "[]" + + items, weights = parse_weighted_items(raw) + + if runner_index is None: + groups = distribute(items, weights, max_runners) + for i, group in enumerate(groups): + labels = " ".join(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("ASSIGNED_ITEMS", "") + _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 = items_for_runner(items, weights, zero_based, max_runners) + encoded = " ".join(assigned) + + if github_env: + _write_github_env("ASSIGNED_ITEMS", encoded) + _write_github_env("SKIP", "false") + click.echo(f"Assigned {len(assigned)} items to runner {runner_index}: {encoded}") + return + + click.echo(encoded) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/src/devx/translations.json b/src/devx/translations.json index 11543f2..daa4802 100644 --- a/src/devx/translations.json +++ b/src/devx/translations.json @@ -639,6 +639,14 @@ "ru": "Dockerfile not found: {path}", "zh": "Dockerfile not found: {path}" }, + "Each item must be a string or an object with 'id', got {type}": { + "bg": "Всеки елемент трябва да е низ или обект с 'id', получено {type}", + "de": "Jedes Element muss ein String oder ein Objekt mit 'id' sein, erhalten {type}", + "en": "Each item must be a string or an object with 'id', got {type}", + "pl": "Każdy element musi być ciągiem lub obiektem z 'id', otrzymano {type}", + "ru": "Каждый элемент должен быть строкой или объектом с 'id', получено {type}", + "zh": "每个元素必须是字符串或带有 'id' 的对象,得到 {type}" + }, "Dry-run mode: on branch '{branch}' (not master). Some checks may behave differently.": { "bg": "Dry-run mode: on branch '{branch}' (not master). Some checks may behave differently.", "de": "Dry-run mode: on branch '{branch}' (not master). Some checks may behave differently.", @@ -2247,6 +2255,14 @@ "ru": "Invalid checklist category: {cat}. Must be numbers.", "zh": "Invalid checklist category: {cat}. Must be numbers." }, + "Items input must be a JSON array, got {type}": { + "bg": "Входните данни трябва да са JSON масив, получено {type}", + "de": "Eingabe muss ein JSON-Array sein, erhalten {type}", + "en": "Items input must be a JSON array, got {type}", + "pl": "Dane wejściowe muszą być tablicą JSON, otrzymano {type}", + "ru": "Входные данные должны быть JSON-массивом, получено {type}", + "zh": "输入必须是 JSON 数组,得到 {type}" + }, "Review body must be at least 50 characters.": { "en": "Review body must be at least 50 characters.", "bg": "Review body must be at least 50 characters.", diff --git a/tests/unit/test_distribute_items.py b/tests/unit/test_distribute_items.py new file mode 100644 index 0000000..95d1540 --- /dev/null +++ b/tests/unit/test_distribute_items.py @@ -0,0 +1,250 @@ +"""Unit tests for devx.ci.distribute_items.""" + +import pytest +from click.testing import CliRunner + +from devx.ci.distribute_items import ( + DEFAULT_WEIGHT, + distribute, + items_for_runner, + main, + parse_items, + parse_weighted_items, +) + + +class TestParseItems: + def test_string_array(self) -> None: + assert parse_items('["a", "b", "c"]') == ["a", "b", "c"] + + def test_object_array(self) -> None: + raw = '[{"id": "a", "weight": 2}, {"id": "b"}]' + assert parse_items(raw) == ["a", "b"] + + def test_empty_array(self) -> None: + assert parse_items("[]") == [] + + def test_not_an_array(self) -> None: + with pytest.raises(Exception, match="must be a JSON array"): + parse_items('{"key": "value"}') + + def test_invalid_entry_type(self) -> None: + with pytest.raises(Exception, match="must be a string or an object"): + parse_items("[42]") + + def test_object_without_id(self) -> None: + with pytest.raises(Exception, match="must be a string or an object"): + parse_items('[{"weight": 2}]') + + +class TestParseWeightedItems: + def test_string_array_default_weights(self) -> None: + items, weights = parse_weighted_items('["a", "b"]') + assert items == ["a", "b"] + assert weights == [DEFAULT_WEIGHT, DEFAULT_WEIGHT] + + def test_object_array_with_weights(self) -> None: + items, weights = parse_weighted_items('[{"id": "a", "weight": 5}, {"id": "b", "weight": 1}]') + assert items == ["a", "b"] + assert weights == [5, 1] + + def test_object_array_missing_weight(self) -> None: + items, weights = parse_weighted_items('[{"id": "a"}]') + assert items == ["a"] + assert weights == [DEFAULT_WEIGHT] + + def test_not_an_array(self) -> None: + with pytest.raises(Exception, match="must be a JSON array"): + parse_weighted_items('"hello"') + + def test_invalid_entry(self) -> None: + with pytest.raises(Exception, match="must be a string or an object"): + parse_weighted_items("[true]") + + +class TestDistribute: + def test_even_split(self) -> None: + items = [f"vm-{i}" for i in range(6)] + weights = [1] * 6 + groups = distribute(items, weights, 3) + assert len(groups) == 3 + assert all(len(g) == 2 for g in groups) + + def test_uneven_split(self) -> None: + items = [f"vm-{i}" for i in range(5)] + weights = [1] * 5 + groups = distribute(items, weights, 3) + assert len(groups[0]) == 2 + assert len(groups[1]) == 2 + assert len(groups[2]) == 1 + + def test_more_runners_than_items(self) -> None: + items = ["vm-a"] + weights = [1] + groups = distribute(items, weights, 5) + assert len(groups) == 5 + assert len(groups[0]) == 1 + assert all(len(g) == 0 for g in groups[1:]) + + def test_lpt_heavy_item_on_least_loaded(self) -> None: + items = ["heavy", "light1", "light2", "light3"] + weights = [10, 1, 1, 1] + groups = distribute(items, weights, 2) + # Heavy item goes to runner 0, lights go to runner 1 (least loaded) + assert "heavy" in groups[0] + # Runner 1 should have more items but less total weight + assert len(groups[1]) >= 2 + + def test_empty_items(self) -> None: + groups = distribute([], [], 3) + assert len(groups) == 3 + assert all(len(g) == 0 for g in groups) + + def test_single_runner(self) -> None: + items = ["a", "b", "c"] + weights = [1, 2, 3] + groups = distribute(items, weights, 1) + assert len(groups) == 1 + assert len(groups[0]) == 3 + + +class TestItemsForRunner: + def test_returns_assigned_subset(self) -> None: + items = ["a", "b", "c", "d", "e", "f"] + weights = [1] * 6 + result = items_for_runner(items, weights, 0, 3) + assert len(result) == 2 + assert all(item in items for item in result) + + def test_out_of_range(self) -> None: + with pytest.raises(Exception, match="out of range"): + items_for_runner(["a"], [1], 5, 3) + + def test_negative_index(self) -> None: + with pytest.raises(Exception, match="out of range"): + items_for_runner(["a"], [1], -1, 3) + + +class TestMain: + def test_stdin_string_array(self) -> None: + runner = CliRunner() + result = runner.invoke(main, ["--runner-index", "1", "--max-runners", "2"], input='["a", "b", "c"]') + assert result.exit_code == 0 + # LPT: heaviest first, so "a" goes to runner 0, "b" to runner 1, "c" to runner 0 + # All weights equal, so round-robin-ish: runner 0 gets "a","c"; runner 1 gets "b" + assert "a" in result.output + + def test_stdin_object_array(self) -> None: + runner = CliRunner() + result = runner.invoke( + main, + ["--runner-index", "1", "--max-runners", "2"], + input='[{"id": "a", "weight": 5}, {"id": "b", "weight": 1}]', + ) + assert result.exit_code == 0 + assert "a" in result.output + + def test_items_file(self, tmp_path: object) -> None: + import pathlib + + items_file = pathlib.Path(str(tmp_path)) / "items.json" + items_file.write_text('["x", "y", "z"]') + runner = CliRunner() + result = runner.invoke(main, ["--items-file", str(items_file), "--runner-index", "1", "--max-runners", "3"]) + assert result.exit_code == 0 + assert "x" in result.output + + def test_print_all_groups_no_runner_index(self) -> None: + runner = CliRunner() + result = runner.invoke(main, ["--max-runners", "2"], input='["a", "b"]') + assert result.exit_code == 0 + assert "Runner 0:" in result.output + assert "Runner 1:" in result.output + + def test_empty_stdin(self) -> None: + runner = CliRunner() + result = runner.invoke(main, ["--runner-index", "1", "--max-runners", "3"], input="") + assert result.exit_code == 0 + # Empty input → empty assigned items + assert result.output.strip() == "" + + def test_github_env(self, tmp_path: object, monkeypatch: pytest.MonkeyPatch) -> None: + import pathlib + + gh_env = pathlib.Path(str(tmp_path)) / "gh_env" + gh_env.write_text("") + monkeypatch.setenv("GITHUB_ENV", str(gh_env)) + runner = CliRunner() + result = runner.invoke( + main, + ["--runner-index", "1", "--max-runners", "2", "--github-env"], + input='["a", "b"]', + ) + assert result.exit_code == 0 + content = gh_env.read_text() + assert "ASSIGNED_ITEMS=" in content + assert "SKIP=false" in content + + def test_skip_if_excess(self, tmp_path: object, monkeypatch: pytest.MonkeyPatch) -> None: + import pathlib + + gh_env = pathlib.Path(str(tmp_path)) / "gh_env" + gh_env.write_text("") + monkeypatch.setenv("GITHUB_ENV", str(gh_env)) + runner = CliRunner() + result = runner.invoke( + main, + ["--runner-index", "5", "--max-runners", "3", "--github-env", "--skip-if-excess"], + input='["a"]', + ) + assert result.exit_code == 0 + content = gh_env.read_text() + assert "ASSIGNED_ITEMS=" in content + assert "SKIP=true" in content + + def test_runner_index_zero(self) -> None: + runner = CliRunner() + result = runner.invoke(main, ["--runner-index", "0"], input='["a"]') + assert result.exit_code != 0 + assert "out of range" in result.output + + def test_default_max_runners(self) -> None: + runner = CliRunner() + result = runner.invoke(main, ["--runner-index", "1"], input='["a"]') + assert result.exit_code == 0 + assert "a" in result.output + + def test_github_env_not_set(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("GITHUB_ENV", raising=False) + runner = CliRunner() + result = runner.invoke( + main, + ["--runner-index", "1", "--github-env"], + input='["a"]', + ) + assert result.exit_code != 0 + assert "GITHUB_ENV" in result.output + + def test_invalid_json(self) -> None: + runner = CliRunner() + result = runner.invoke(main, ["--runner-index", "1"], input="not json") + assert result.exit_code != 0 + + def test_multiline_github_env(self, tmp_path: object, monkeypatch: pytest.MonkeyPatch) -> None: + import pathlib + + gh_env = pathlib.Path(str(tmp_path)) / "gh_env" + gh_env.write_text("") + monkeypatch.setenv("GITHUB_ENV", str(gh_env)) + runner = CliRunner() + # Items with newlines in their IDs would trigger multiline syntax + # Normal items don't have newlines, but test the path anyway + result = runner.invoke( + main, + ["--runner-index", "1", "--max-runners", "1", "--github-env"], + input='["a\\nb"]', + ) + assert result.exit_code == 0 + content = gh_env.read_text() + # Item "a\nb" contains a newline → heredoc syntax + assert "ASSIGNED_ITEMS<<" in content