scripts.schedule_tests

Reference › scripts › schedule_tests

scripts/schedule_tests.py is the command-line front end for the CLP(FD) test-module scheduler (section 9, row B28), whose solver lives in orchestrata_ultimate.core.threading.test_scheduler. A human or a runner runs it to get a schedule of test modules across workers. Its duration input is the cache that test_durations maintains.

There is one subcommand, plan. It requires --durations (a JSON file mapping module to seconds) and --workers. Optional inputs are --infra-map (a JSON file with a modules map of infrastructure needs and a workers list of what each worker offers, read by _load_infra_map, _parse_infra_modules and _parse_infra_workers), --conflicts (a JSON list of module pairs that must not share a worker), --budget-ms (a wall-clock budget for the solver) and --json. run computes the plan and _print_plan prints a manifest, or JSON with --json. main exits with the code run returns.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/scripts/schedule_tests.py

#!/usr/bin/env python3
"""CLI for the §9 row B28 CLP(FD) test-module scheduler.

    python scripts/schedule_tests.py plan --durations F --workers N \
        [--infra-map F --conflicts F --budget-ms B --json]

Prints a per-worker manifest (module list) + predicted makespan.

Exit codes:
    0 -- a budget-satisfying (or, with no ``--budget-ms``, unconstrained-
         optimal) plan was found.
    2 -- usage/input error (bad JSON, malformed infra-map, malformed
         conflicts file, etc).
    3 -- ``--budget-ms`` was given and is infeasible; the printed plan is
         still the best achievable (unconstrained-optimum) fallback --
         callers can match on this distinct code and fall back to naive
         batching.

``--infra-map F`` (optional) is a JSON object:
    {"modules": {"path/to/mod_test.py": ["neo4j"]},
     "workers": [["redis"], ["redis", "neo4j"]]}
Both keys are optional. Omitted ``modules`` -> no module has any infra
requirement. Omitted ``workers`` -> every one of the ``--workers N`` workers
gets an EMPTY capability set (matches today's real fleet: every worker has
equal access to the same shared Redis/graph -- see docs/LEVEL0.md -- so
infra differentiation is a hypothetical the CLP(FD) model supports, not a
real constraint of the current fleet). Omitting ``--infra-map`` entirely is
the same as an infra-map with neither key.

B28 ADOPTION SEAM: this CLI is a standalone planning tool -- it is NOT
wired into ``scripts/run_tests.py`` or ``scripts/run_tests_threadpool.py``
(both load-bearing, out of scope for this task). Wiring it in would mean:
(1) ``run_tests_threadpool.py``'s batch-assignment step (``_lpt_pack_modules``
et al., ``core/threading/test_executor.py``) calling ``plan_schedule()``
instead of (or beside) its own LPT packer, fed by the SAME duration cache
this CLI reads (see ``load_durations``'s docstring); (2) a real
``--conflicts`` producer -- ``scripts/gate_parallel_safety.py`` is the
intended one but does not emit this pair format yet (see
``load_conflicts``'s docstring); (3) an infra-map reflecting a genuinely
heterogeneous worker fleet, which does not exist today.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

# Ensure repo src is on path when run directly (mirrors scripts/context_pack.py).
_repo_src = Path(__file__).resolve().parent.parent / "src"
if str(_repo_src) not in sys.path:
    sys.path.insert(0, str(_repo_src))

from orchestrata_ultimate.core.threading.test_scheduler import (  # noqa: E402
    ScheduledPlan,
    load_conflicts,
    load_durations,
    plan_schedule,
)

EXIT_OK = 0
EXIT_USAGE_ERROR = 2
EXIT_INFEASIBLE = 3


def _parse_infra_modules(path: str, raw_modules: object) -> dict[str, frozenset[str]]:
    """Parse the optional 'modules' key of an ``--infra-map`` file."""
    if not isinstance(raw_modules, dict):
        raise ValueError(
            f"{path}: 'modules' must be a JSON object of {{module: [infra,...]}}"
        )
    module_infra: dict[str, frozenset[str]] = {}
    for module, reqs in raw_modules.items():
        if not isinstance(reqs, list) or not all(isinstance(r, str) for r in reqs):
            raise ValueError(
                f"{path}: infra requirements for {module!r} must be a list of strings"
            )
        module_infra[str(module)] = frozenset(reqs)
    return module_infra


def _parse_infra_workers(
    path: str, raw_workers: object, num_workers: int
) -> list[frozenset[str]]:
    """Parse the optional 'workers' key of an ``--infra-map`` file; defaults
    to ``num_workers`` empty-capability workers when the key is omitted."""
    if raw_workers is None:
        return [frozenset() for _ in range(num_workers)]
    if not isinstance(raw_workers, list) or len(raw_workers) != num_workers:
        raise ValueError(
            f"{path}: 'workers' must be a list of exactly {num_workers} "
            f"infra-capability lists (--workers {num_workers})"
        )
    workers: list[frozenset[str]] = []
    for idx, caps in enumerate(raw_workers):
        if not isinstance(caps, list) or not all(isinstance(c, str) for c in caps):
            raise ValueError(
                f"{path}: worker {idx}'s capabilities must be a list of strings"
            )
        workers.append(frozenset(caps))
    return workers


def _load_infra_map(
    path: str | None, num_workers: int
) -> tuple[dict[str, frozenset[str]], list[frozenset[str]]]:
    """Read the optional ``--infra-map`` file; see this module's docstring
    for its schema and defaults."""
    if path is None:
        return {}, [frozenset() for _ in range(num_workers)]

    with open(path, encoding="utf-8") as fh:
        data = json.load(fh)
    if not isinstance(data, dict):
        raise ValueError(
            f"{path}: expected a JSON object with 'modules'/'workers' keys"
        )

    module_infra = _parse_infra_modules(path, data.get("modules", {}))
    workers = _parse_infra_workers(path, data.get("workers"), num_workers)
    return module_infra, workers


def _print_plan(plan: ScheduledPlan, num_workers: int, as_json: bool) -> None:
    if as_json:
        print(
            json.dumps(
                {
                    "feasible": plan.feasible,
                    "predicted_makespan_ms": plan.predicted_makespan_ms,
                    "assignments": plan.assignments,
                },
                indent=2,
                sort_keys=True,
            )
        )
        return

    per_worker: dict[int, list[str]] = {idx: [] for idx in range(num_workers)}
    for module, worker_idx in sorted(plan.assignments.items()):
        per_worker[worker_idx].append(module)
    for idx in range(num_workers):
        print(f"worker {idx}:")
        for module in per_worker[idx]:
            print(f"  {module}")
    print(f"predicted makespan: {plan.predicted_makespan_ms}ms")
    print(f"feasible: {plan.feasible}")


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="schedule_tests",
        description="CLP(FD) test-module scheduler (§9 row B28).",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    plan_parser = subparsers.add_parser("plan", help="Compute a schedule.")
    plan_parser.add_argument(
        "--durations",
        required=True,
        help="Path to a duration-cache JSON file (module -> seconds).",
    )
    plan_parser.add_argument(
        "--workers", type=int, required=True, help="Number of workers."
    )
    plan_parser.add_argument(
        "--infra-map",
        default=None,
        help="Path to a JSON {'modules': {...}, 'workers': [...]} infra map.",
    )
    plan_parser.add_argument(
        "--conflicts",
        default=None,
        help="Path to a JSON list of [modA, modB] conflict pairs.",
    )
    plan_parser.add_argument(
        "--budget-ms", type=int, default=None, help="Wall-clock budget in ms."
    )
    plan_parser.add_argument(
        "--json",
        action="store_true",
        help="Emit machine-readable JSON instead of a manifest.",
    )
    return parser


def run(argv: list[str] | None = None) -> int:
    """Parse args, load inputs, solve, print the plan, and return the
    process exit code (see this module's docstring for the code table)."""
    parser = _build_parser()
    args = parser.parse_args(argv)

    if args.workers < 1:
        print("ERROR: --workers must be >= 1", file=sys.stderr)
        return EXIT_USAGE_ERROR

    try:
        durations = load_durations(args.durations)
        module_infra, workers = _load_infra_map(args.infra_map, args.workers)
        conflicts = load_conflicts(args.conflicts) if args.conflicts else set()
    except (OSError, ValueError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return EXIT_USAGE_ERROR

    try:
        plan = plan_schedule(
            durations, workers, module_infra, conflicts, args.budget_ms
        )
    except ValueError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return EXIT_USAGE_ERROR

    _print_plan(plan, args.workers, args.json)
    return EXIT_OK if plan.feasible else EXIT_INFEASIBLE


def main(argv: list[str] | None = None) -> None:
    """Process entry point: ``sys.exit`` with ``run``'s return code."""
    sys.exit(run(argv))


if __name__ == "__main__":
    main()