From 1f880a87855fec93131cc704d02bdd65192c7f12 Mon Sep 17 00:00:00 2001 From: TohruSu <238259854+TohruSu@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:44:31 +0000 Subject: [PATCH] Add tabular Q-Learning meta-scheduler with shared simulation steps --- README.md | 20 ++- docs/simulation.md | 58 ++++++++ pyproject.toml | 2 +- src/adaptivequeue/__init__.py | 5 +- src/adaptivequeue/cli.py | 98 ++++++++++--- src/adaptivequeue/metrics.py | 5 +- src/adaptivequeue/models.py | 4 + src/adaptivequeue/qlearning.py | 235 ++++++++++++++++++++++++++++++++ src/adaptivequeue/simulation.py | 84 ++++++++---- tests/test_cli.py | 52 +++++++ tests/test_qlearning.py | 205 ++++++++++++++++++++++++++++ 11 files changed, 712 insertions(+), 56 deletions(-) create mode 100644 src/adaptivequeue/qlearning.py create mode 100644 tests/test_qlearning.py diff --git a/README.md b/README.md index 9438a31..2c0f095 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ # AdaptiveQueue A lightweight Python project for experimenting with adaptive background job -scheduling. The first milestone provides a discrete-time, single-worker simulator -and FIFO, SPT, EDF, Priority, and seeded Random baselines. +scheduling. Includes a discrete-time, single-worker simulator, FIFO, SPT, EDF, +Priority, and seeded Random baselines, and a tabular Q-Learning meta-scheduler. Requires Python 3.10 or newer. There are no runtime dependencies. @@ -23,6 +23,18 @@ rate, average and 95th-percentile waiting time, and average flow time. Workloads are `stable`, `bursty`, or `mixed`; policies are `fifo`, `spt`, `edf`, `priority`, or `random`. Use the same seed and workload settings to compare policies fairly. +Train a meta-scheduler that chooses FIFO, SPT, EDF, or Priority at each decision, +then evaluate its frozen Q-table on held-out workload seeds: + +```sh +adaptivequeue train --workload mixed --episodes 1000 --seed 42 --output policy.json +adaptivequeue evaluate --workload mixed --episodes 20 --model policy.json +``` + +Evaluation defaults to seeds following the saved training range, with exploration +disabled. Use `--seed` to choose a different disjoint range. Both commands accept +the workload options below; `train --help` lists learning and exploration settings. + ```sh python -m adaptivequeue --workload bursty --policy spt --seed 7 \ --jobs 60 --arrival-rate 0.3 --deadline-factor 3 @@ -34,5 +46,5 @@ mypy See [simulation rules and parameters](docs/simulation.md) or run `adaptivequeue --help`. The Python API also exposes per-job execution records. -Reinforcement learning is planned for a later milestone; the current scope is -simulation and conventional scheduling baselines. +The current scope is simulation, classical baselines, and CPU-only tabular +Q-Learning; no additional runtime dependencies are needed. diff --git a/docs/simulation.md b/docs/simulation.md index 9c424ff..8990022 100644 --- a/docs/simulation.md +++ b/docs/simulation.md @@ -86,3 +86,61 @@ episode. Cost is averaged over jobs, without normalizing by total priority weigh Compare policies on identical generated jobs and settings. Seeds control both workloads and Random decisions independently; several episode seeds give a more useful comparison than a single run. + +## Tabular Q-Learning meta-scheduler + +`Simulation(jobs)` exposes `waiting_jobs`, `current_time`, and `done` at each +decision. Its `step(policy)` executes exactly one job, returns its `JobResult`, +and advances to the next ready queue (or termination). `simulate()` uses this +same engine. Only the engine holds future arrivals. `qlearning.step_action` +accepts an action index and returns the completed job, reward, next encoded state, +and `done`. Terminal states are represented by `None`, including empty episodes. + +The state tuple has 96 possible values: three zero-based bucket indices and a +boolean, computed only from currently waiting jobs and the current time: + +| Feature | Buckets in order | +| --- | --- | +| Queue length | 1; 2–3; 4–7; 8+ | +| Minimum slack | overdue: <0; urgent: 0–2; moderate: 3–6; relaxed: 7+ | +| Short-job proportion | low: <1/3; medium: 1/3 to <2/3; high: >=2/3 | +| High-priority job waiting | true if any weight is >=4, otherwise false | + +A short job has duration <=2. Each job's slack is +`deadline - current_time - duration`: its remaining deadline budget after allowing +for processing. Waiting jobs have not started, so their remaining processing time +is their full duration. Negative slack means even immediate execution would miss +the deadline; it does not require the absolute deadline to have passed. + +Actions 0, 1, 2, and 3 invoke FIFO, SPT, EDF, and Priority respectively, preserving +their tie-breaking rules. The chosen baseline selects the next job. Reward is +`-priority_weight * (tardiness + 0.1 * flow_time)` for that completed job. The +episode return is `-n * J`; for a fixed job count, maximizing undiscounted return +minimizes the existing primary cost. + +The update is `Q(s,a) += alpha * (reward + gamma * max_a' Q(s',a') - Q(s,a))`, +omitting the bootstrap term for terminal transitions. Defaults are alpha=0.1, +gamma=1.0, and epsilon-greedy exploration with +`epsilon(episode) = max(0.05, 1.0 * 0.995**episode)` (episode numbering starts at +zero). Discounting, when configured, applies per scheduled job, not per time unit. +Q-values start at zero. Greedy ties choose the lowest action index, so unseen +states fall back to FIFO. Exploration chooses uniformly among all four actions. + +`qlearning.train()` generates workload seeds `seed` through `seed + episodes - 1`; +action selection uses a separate local RNG. RL seeds must be nonnegative. +`evaluate()` disables exploration and all Q-table updates, including insertion of +unseen states. It defaults to the next workload seed after training and rejects +overlapping training/evaluation seed ranges. All five metrics are computed over +pooled jobs from evaluation episodes; p95 is not an average of episode percentiles. +For classical comparisons, run baselines on those same evaluation seeds and settings. + +`QLearningAgent.save()` writes readable, versioned JSON containing action order, +state thresholds, reward definition, learning parameters, training workload and +seed range, and Q-values. `load()` preserves greedy predictions and rejects +incompatible formulations. The format supports evaluation, not exact continuation +of an interrupted training RNG stream. A loaded agent's `policy` method also works +with `simulate(jobs, agent.policy)`. + +The coarse state aliases different queues and omits future arrivals, so it is not +a full Markov description. Training does not guarantee improvement over a baseline; +evaluate on several held-out seeds and retrain when workload characteristics change. diff --git a/pyproject.toml b/pyproject.toml index 9d3b9a2..299c35b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "adaptivequeue" version = "0.1.0" -description = "A discrete-time job scheduling simulator and baseline policies" +description = "A job scheduling simulator with classical and tabular Q-Learning policies" readme = "README.md" requires-python = ">=3.10" license = {file = "LICENSE"} diff --git a/src/adaptivequeue/__init__.py b/src/adaptivequeue/__init__.py index 0b735f9..f2bd638 100644 --- a/src/adaptivequeue/__init__.py +++ b/src/adaptivequeue/__init__.py @@ -1,15 +1,16 @@ -"""Discrete-time simulation and conventional job scheduling baselines.""" +"""Discrete-time simulation, scheduling baselines, and a tabular meta-scheduler.""" from adaptivequeue.metrics import Metrics from adaptivequeue.models import Job, JobResult from adaptivequeue.policies import make_policy -from adaptivequeue.simulation import SimulationResult, simulate +from adaptivequeue.simulation import Simulation, SimulationResult, simulate from adaptivequeue.workloads import WorkloadConfig, bursty_arrivals, mixed_workload, stable_arrivals __all__ = [ "Job", "JobResult", "Metrics", + "Simulation", "SimulationResult", "WorkloadConfig", "bursty_arrivals", diff --git a/src/adaptivequeue/cli.py b/src/adaptivequeue/cli.py index d94a6f9..c4d050c 100644 --- a/src/adaptivequeue/cli.py +++ b/src/adaptivequeue/cli.py @@ -1,17 +1,20 @@ -"""Command-line entry point for one reproducible scheduling episode.""" +"""Baseline episodes and tabular Q-Learning training/evaluation.""" import argparse +from pathlib import Path +from adaptivequeue.metrics import Metrics from adaptivequeue.policies import POLICY_NAMES, make_policy +from adaptivequeue.qlearning import LearningConfig, QLearningAgent, evaluate, train from adaptivequeue.simulation import simulate from adaptivequeue.workloads import WORKLOADS, WorkloadConfig -def main() -> None: - parser = argparse.ArgumentParser(description="Run an AdaptiveQueue scheduling episode.") - parser.add_argument("--workload", choices=WORKLOADS, default="stable") - parser.add_argument("--policy", choices=POLICY_NAMES, default="fifo") - parser.add_argument("--seed", type=int, default=42) +def _add_workload_arguments( + parser: argparse.ArgumentParser, *, workload: str = "stable", seed: int | None = 42 +) -> None: + parser.add_argument("--workload", choices=WORKLOADS, default=workload) + parser.add_argument("--seed", type=int, default=seed) parser.add_argument("--jobs", type=int, default=60) parser.add_argument("--arrival-rate", type=float, default=0.28, help="Mean jobs per time unit") parser.add_argument( @@ -21,6 +24,36 @@ def main() -> None: parser.add_argument( "--burst-window", type=int, default=3, help="Burst arrival window in time units" ) + + +def _print_metrics(metrics: Metrics) -> None: + print(f"Primary cost J (lower is better): {metrics.primary_cost:.3f}") + print(f"Deadline miss rate: {metrics.deadline_miss_rate:.1%}") + print(f"Average waiting time: {metrics.average_waiting_time:.3f}") + print(f"95th-percentile waiting time: {metrics.p95_waiting_time:.3f}") + print(f"Average flow time: {metrics.average_flow_time:.3f}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run AdaptiveQueue scheduling experiments.") + _add_workload_arguments(parser) + parser.add_argument("--policy", choices=POLICY_NAMES, default="fifo") + commands = parser.add_subparsers(dest="command") + + training = commands.add_parser("train", help="Train and save a Q-Learning meta-scheduler") + _add_workload_arguments(training, workload="mixed") + training.add_argument("--episodes", type=int, default=1000) + training.add_argument("--learning-rate", type=float, default=0.1) + training.add_argument("--gamma", type=float, default=1.0) + training.add_argument("--epsilon-start", type=float, default=1.0) + training.add_argument("--epsilon-end", type=float, default=0.05) + training.add_argument("--epsilon-decay", type=float, default=0.995) + training.add_argument("--output", type=Path, required=True) + + evaluation = commands.add_parser("evaluate", help="Evaluate a frozen Q-table on held-out seeds") + _add_workload_arguments(evaluation, workload="mixed", seed=None) + evaluation.add_argument("--episodes", type=int, default=20) + evaluation.add_argument("--model", type=Path, required=True) args = parser.parse_args() try: @@ -31,17 +64,44 @@ def main() -> None: burst_size=args.burst_size, burst_window=args.burst_window, ) - except ValueError as error: + if args.command == "train": + learning = LearningConfig( + args.learning_rate, + args.gamma, + args.epsilon_start, + args.epsilon_end, + args.epsilon_decay, + ) + agent = train( + workload=args.workload, + episodes=args.episodes, + seed=args.seed, + workload_config=config, + learning_config=learning, + ) + agent.save(args.output) + print( + f"Trained {args.episodes} episodes; " + f"{len(agent.q_table)} states; saved {args.output}" + ) + elif args.command == "evaluate": + agent = QLearningAgent.load(args.model) + metrics = evaluate( + agent, + workload=args.workload, + episodes=args.episodes, + seed=args.seed, + workload_config=config, + ) + print(f"Workload: {args.workload} | Policy: qlearning | Episodes: {args.episodes}") + _print_metrics(metrics) + else: + jobs = WORKLOADS[args.workload](args.seed, config) + result = simulate(jobs, make_policy(args.policy, seed=args.seed)) + print( + f"Workload: {args.workload} | Policy: {args.policy} | " + f"Seed: {args.seed} | Jobs: {len(jobs)}" + ) + _print_metrics(result.metrics) + except (OSError, ValueError, KeyError, TypeError) as error: parser.error(str(error)) - - jobs = WORKLOADS[args.workload](args.seed, config) - result = simulate(jobs, make_policy(args.policy, seed=args.seed)) - metrics = result.metrics - print( - f"Workload: {args.workload} | Policy: {args.policy} | Seed: {args.seed} | Jobs: {len(jobs)}" - ) - print(f"Primary cost J (lower is better): {metrics.primary_cost:.3f}") - print(f"Deadline miss rate: {metrics.deadline_miss_rate:.1%}") - print(f"Average waiting time: {metrics.average_waiting_time:.3f}") - print(f"95th-percentile waiting time: {metrics.p95_waiting_time:.3f}") - print(f"Average flow time: {metrics.average_flow_time:.3f}") diff --git a/src/adaptivequeue/metrics.py b/src/adaptivequeue/metrics.py index f87e39b..70a7dca 100644 --- a/src/adaptivequeue/metrics.py +++ b/src/adaptivequeue/metrics.py @@ -26,10 +26,7 @@ def calculate_metrics(results: Sequence[JobResult]) -> Metrics: lower, upper = floor(position), ceil(position) p95 = waiting[lower] + (waiting[upper] - waiting[lower]) * (position - lower) return Metrics( - primary_cost=fmean( - result.job.priority_weight * (result.tardiness + 0.1 * result.flow_time) - for result in results - ), + primary_cost=fmean(result.cost for result in results), deadline_miss_rate=fmean(result.tardiness > 0 for result in results), average_waiting_time=fmean(waiting), p95_waiting_time=p95, diff --git a/src/adaptivequeue/models.py b/src/adaptivequeue/models.py index eb5b577..8888245 100644 --- a/src/adaptivequeue/models.py +++ b/src/adaptivequeue/models.py @@ -43,3 +43,7 @@ def flow_time(self) -> int: @property def tardiness(self) -> int: return max(0, self.completion_time - self.job.deadline) + + @property + def cost(self) -> float: + return self.job.priority_weight * (self.tardiness + 0.1 * self.flow_time) diff --git a/src/adaptivequeue/qlearning.py b/src/adaptivequeue/qlearning.py new file mode 100644 index 0000000..0df966c --- /dev/null +++ b/src/adaptivequeue/qlearning.py @@ -0,0 +1,235 @@ +"""Tabular Q-Learning over four conventional scheduling policies.""" + +import json +import random +from dataclasses import asdict, dataclass +from math import isfinite +from pathlib import Path +from typing import Sequence + +from adaptivequeue.metrics import Metrics, calculate_metrics +from adaptivequeue.models import Job, JobResult +from adaptivequeue.policies import edf, fifo, priority, spt +from adaptivequeue.simulation import Simulation, simulate +from adaptivequeue.workloads import WORKLOADS, WorkloadConfig + +State = tuple[int, int, int, bool] +ACTION_NAMES = ("fifo", "spt", "edf", "priority") +ACTION_POLICIES = (fifo, spt, edf, priority) +STATE_ENCODING = { + "fields": [ + "queue_length_bucket", + "min_slack_bucket", + "short_proportion_bucket", + "high_priority", + ], + "queue_length_upper_bounds": [1, 3, 7], + "slack": "deadline - current_time - duration", + "slack_upper_bounds": [-1, 2, 6], + "short_duration_max": 2, + "short_proportion_boundaries": ["1/3", "2/3"], + "high_priority_min": 4, +} +REWARD = "-priority_weight * (tardiness + 0.1 * flow_time)" + + +def encode_state(waiting_jobs: Sequence[Job], current_time: int) -> State: + if not waiting_jobs: + raise ValueError("state encoding requires a nonempty ready queue") + size = len(waiting_jobs) + queue_bucket = sum(size > bound for bound in (1, 3, 7)) + slack = min(job.deadline - current_time - job.duration for job in waiting_jobs) + slack_bucket = sum(slack > bound for bound in (-1, 2, 6)) + short_count = sum(job.duration <= 2 for job in waiting_jobs) + short_bucket = int(3 * short_count >= size) + int(3 * short_count >= 2 * size) + return ( + queue_bucket, + slack_bucket, + short_bucket, + any(job.priority_weight >= 4 for job in waiting_jobs), + ) + + +@dataclass(frozen=True) +class Transition: + job: JobResult + reward: float + next_state: State | None + done: bool + + +def step_action(simulation: Simulation, action: int) -> Transition: + if not 0 <= action < len(ACTION_POLICIES): + raise ValueError("action must be an index from 0 to 3") + record = simulation.step(ACTION_POLICIES[action]) + next_state = ( + None if simulation.done else encode_state(simulation.waiting_jobs, simulation.current_time) + ) + return Transition(record, -record.cost, next_state, simulation.done) + + +@dataclass(frozen=True) +class LearningConfig: + learning_rate: float = 0.1 + gamma: float = 1.0 + epsilon_start: float = 1.0 + epsilon_end: float = 0.05 + epsilon_decay: float = 0.995 + + def __post_init__(self) -> None: + if not 0 < self.learning_rate <= 1 or not 0 <= self.gamma <= 1: + raise ValueError("learning_rate must be in (0, 1] and gamma in [0, 1]") + if not 0 <= self.epsilon_end <= self.epsilon_start <= 1: + raise ValueError("epsilon must satisfy 0 <= end <= start <= 1") + if not 0 < self.epsilon_decay <= 1: + raise ValueError("epsilon_decay must be in (0, 1]") + + def epsilon(self, episode: int) -> float: + return max(self.epsilon_end, self.epsilon_start * self.epsilon_decay**episode) + + +@dataclass(frozen=True) +class TrainingInfo: + workload: str + workload_config: WorkloadConfig + episodes: int + seed: int + + +class QLearningAgent: + def __init__(self, *, seed: int, config: LearningConfig = LearningConfig()) -> None: + if seed < 0: + raise ValueError("seed must be nonnegative") + self.seed = seed + self.config = config + self.training: TrainingInfo | None = None + self.q_table: dict[State, list[float]] = {} + self._rng = random.Random(seed) + + def greedy_action(self, state: State) -> int: + values = self.q_table.get(state, [0.0] * len(ACTION_NAMES)) + return max(range(len(ACTION_NAMES)), key=lambda action: values[action]) + + def choose_action(self, state: State, *, epsilon: float) -> int: + if not 0 <= epsilon <= 1: + raise ValueError("epsilon must be in [0, 1]") + if epsilon > 0 and self._rng.random() < epsilon: + return self._rng.randrange(len(ACTION_NAMES)) + return self.greedy_action(state) + + def update(self, state: State, action: int, reward: float, next_state: State | None) -> None: + if not 0 <= action < len(ACTION_NAMES): + raise ValueError("action must be an index from 0 to 3") + target = reward + if next_state is not None: + target += self.config.gamma * max( + self.q_table.get(next_state, [0.0] * len(ACTION_NAMES)) + ) + values = self.q_table.setdefault(state, [0.0] * len(ACTION_NAMES)) + values[action] += self.config.learning_rate * (target - values[action]) + + def policy(self, waiting_jobs: Sequence[Job], current_time: int) -> Job: + """A frozen, greedy policy compatible with simulate().""" + action = self.greedy_action(encode_state(waiting_jobs, current_time)) + return ACTION_POLICIES[action](waiting_jobs, current_time) + + def save(self, path: str | Path) -> None: + data = { + "version": 1, + "actions": ACTION_NAMES, + "state_encoding": STATE_ENCODING, + "reward": REWARD, + "seed": self.seed, + "learning_config": asdict(self.config), + "training": asdict(self.training) if self.training else None, + "q_table": [ + {"state": state, "values": values} for state, values in sorted(self.q_table.items()) + ], + } + Path(path).write_text(json.dumps(data, indent=2, allow_nan=False) + "\n", encoding="utf-8") + + @classmethod + def load(cls, path: str | Path) -> "QLearningAgent": + data = json.loads(Path(path).read_text(encoding="utf-8")) + if ( + data["version"] != 1 + or data["actions"] != list(ACTION_NAMES) + or data["state_encoding"] != STATE_ENCODING + or data["reward"] != REWARD + ): + raise ValueError("incompatible Q-table format or scheduling formulation") + agent = cls(seed=data["seed"], config=LearningConfig(**data["learning_config"])) + if data["training"] is not None: + info = data["training"] + workload = info["workload_config"] + workload["duration_range"] = tuple(workload["duration_range"]) + agent.training = TrainingInfo( + info["workload"], WorkloadConfig(**workload), info["episodes"], info["seed"] + ) + for row in data["q_table"]: + raw = row["state"] + if len(raw) != 4 or len(row["values"]) != len(ACTION_NAMES): + raise ValueError("invalid Q-table row dimensions") + state: State = (raw[0], raw[1], raw[2], raw[3]) + if ( + any(type(value) is not int for value in state[:3]) + or not (0 <= state[0] < 4 and 0 <= state[1] < 4 and 0 <= state[2] < 3) + or not isinstance(state[3], bool) + ): + raise ValueError("invalid Q-table state") + values = [float(value) for value in row["values"]] + if not all(isfinite(value) for value in values): + raise ValueError("Q-values must be finite") + agent.q_table[state] = values + return agent + + +def train( + *, + workload: str = "mixed", + episodes: int = 1000, + seed: int = 42, + workload_config: WorkloadConfig = WorkloadConfig(), + learning_config: LearningConfig = LearningConfig(), +) -> QLearningAgent: + if episodes < 1: + raise ValueError("episodes must be positive") + agent = QLearningAgent(seed=seed, config=learning_config) + agent.training = TrainingInfo(workload, workload_config, episodes, seed) + for episode in range(episodes): + simulation = Simulation(WORKLOADS[workload](seed + episode, workload_config)) + state = ( + None + if simulation.done + else encode_state(simulation.waiting_jobs, simulation.current_time) + ) + epsilon = learning_config.epsilon(episode) + while state is not None: + action = agent.choose_action(state, epsilon=epsilon) + transition = step_action(simulation, action) + agent.update(state, action, transition.reward, transition.next_state) + state = transition.next_state + return agent + + +def evaluate( + agent: QLearningAgent, + *, + workload: str = "mixed", + episodes: int = 20, + seed: int | None = None, + workload_config: WorkloadConfig = WorkloadConfig(), +) -> Metrics: + """Pool job metrics across independently seeded, greedy evaluation episodes.""" + training = agent.training + if seed is None: + seed = training.seed + training.episodes if training else agent.seed + 1 + if episodes < 1 or seed < 0: + raise ValueError("episodes must be positive and seed nonnegative") + if training and seed < training.seed + training.episodes and seed + episodes > training.seed: + raise ValueError("evaluation workload seeds must not overlap the training seeds") + records: list[JobResult] = [] + for episode in range(episodes): + jobs = WORKLOADS[workload](seed + episode, workload_config) + records.extend(simulate(jobs, agent.policy).jobs) + return calculate_metrics(records) diff --git a/src/adaptivequeue/simulation.py b/src/adaptivequeue/simulation.py index f373489..def8456 100644 --- a/src/adaptivequeue/simulation.py +++ b/src/adaptivequeue/simulation.py @@ -14,31 +14,63 @@ class SimulationResult: metrics: Metrics -def simulate(jobs: Iterable[Job], policy: Policy) -> SimulationResult: - """Run until every job finishes, returning records in execution order.""" - arrivals = sorted(jobs, key=lambda job: (job.arrival_time, job.job_id)) - if len({job.job_id for job in arrivals}) != len(arrivals): - raise ValueError("job IDs must be unique within an episode") - - waiting: list[Job] = [] - completed: list[JobResult] = [] - time = 0 - next_arrival = 0 - - while waiting or next_arrival < len(arrivals): - if not waiting: - time = max(time, arrivals[next_arrival].arrival_time) - while next_arrival < len(arrivals) and arrivals[next_arrival].arrival_time <= time: - waiting.append(arrivals[next_arrival]) - next_arrival += 1 - - selected = policy(tuple(waiting), time) - if selected not in waiting: +class Simulation: + """Expose each ready queue and execute one job per scheduling decision.""" + + def __init__(self, jobs: Iterable[Job]) -> None: + self._arrivals = sorted(jobs, key=lambda job: (job.arrival_time, job.job_id)) + if len({job.job_id for job in self._arrivals}) != len(self._arrivals): + raise ValueError("job IDs must be unique within an episode") + self._waiting: list[Job] = [] + self._completed: list[JobResult] = [] + self._time = 0 + self._next_arrival = 0 + self._advance_to_decision() + + @property + def waiting_jobs(self) -> tuple[Job, ...]: + return tuple(self._waiting) + + @property + def current_time(self) -> int: + return self._time + + @property + def done(self) -> bool: + return not self._waiting + + def _advance_to_decision(self) -> None: + if not self._waiting and self._next_arrival < len(self._arrivals): + self._time = max(self._time, self._arrivals[self._next_arrival].arrival_time) + while ( + self._next_arrival < len(self._arrivals) + and self._arrivals[self._next_arrival].arrival_time <= self._time + ): + self._waiting.append(self._arrivals[self._next_arrival]) + self._next_arrival += 1 + + def step(self, policy: Policy) -> JobResult: + if self.done: + raise ValueError("cannot step a completed simulation") + selected = policy(self.waiting_jobs, self.current_time) + if selected not in self._waiting: raise ValueError("policy must select a job from the waiting queue") - waiting.remove(selected) - completion_time = time + selected.duration - completed.append(JobResult(selected, time, completion_time)) - time = completion_time + self._waiting.remove(selected) + record = JobResult(selected, self._time, self._time + selected.duration) + self._completed.append(record) + self._time = record.completion_time + self._advance_to_decision() + return record + + def result(self) -> SimulationResult: + """Return execution records and metrics for jobs completed so far.""" + records = tuple(self._completed) + return SimulationResult(records, calculate_metrics(records)) - records = tuple(completed) - return SimulationResult(records, calculate_metrics(records)) + +def simulate(jobs: Iterable[Job], policy: Policy) -> SimulationResult: + """Run until every job finishes, returning records in execution order.""" + simulation = Simulation(jobs) + while not simulation.done: + simulation.step(policy) + return simulation.result() diff --git a/tests/test_cli.py b/tests/test_cli.py index e3a19b5..8f60ad5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,5 +1,6 @@ import subprocess import sys +from pathlib import Path def test_module_cli_displays_metrics() -> None: @@ -41,3 +42,54 @@ def test_invalid_cli_configuration_has_a_useful_error() -> None: assert result.returncode == 2 assert "arrival_rate must be finite and positive" in result.stderr assert "Traceback" not in result.stderr + + +def test_cli_training_and_evaluation(tmp_path: Path) -> None: + model = tmp_path / "policy.json" + training = subprocess.run( + [ + sys.executable, + "-m", + "adaptivequeue", + "train", + "--workload", + "bursty", + "--episodes", + "4", + "--jobs", + "8", + "--seed", + "7", + "--learning-rate", + "0.3", + "--output", + str(model), + ], + check=True, + capture_output=True, + text=True, + ) + assert "Trained 4 episodes" in training.stdout + assert model.is_file() + evaluation = subprocess.run( + [ + sys.executable, + "-m", + "adaptivequeue", + "evaluate", + "--workload", + "bursty", + "--episodes", + "2", + "--jobs", + "8", + "--model", + str(model), + ], + check=True, + capture_output=True, + text=True, + ) + assert "Policy: qlearning | Episodes: 2" in evaluation.stdout + assert "Primary cost J" in evaluation.stdout + assert "95th-percentile waiting time" in evaluation.stdout diff --git a/tests/test_qlearning.py b/tests/test_qlearning.py new file mode 100644 index 0000000..39f9e14 --- /dev/null +++ b/tests/test_qlearning.py @@ -0,0 +1,205 @@ +import copy +import json +import random +from pathlib import Path + +import pytest + +from adaptivequeue import Job, Simulation, WorkloadConfig, simulate +from adaptivequeue.metrics import calculate_metrics +from adaptivequeue.qlearning import ( + ACTION_NAMES, + LearningConfig, + QLearningAgent, + encode_state, + evaluate, + step_action, + train, +) +from adaptivequeue.workloads import WORKLOADS + + +@pytest.mark.parametrize(("size", "bucket"), [(1, 0), (2, 1), (3, 1), (4, 2), (7, 2), (8, 3)]) +def test_queue_length_buckets(size: int, bucket: int) -> None: + jobs = [Job(str(i), 0, 4, 30, 1, "test") for i in range(size)] + assert encode_state(jobs, 10) == (bucket, 3, 0, False) + + +@pytest.mark.parametrize(("slack", "bucket"), [(-1, 0), (0, 1), (2, 1), (3, 2), (6, 2), (7, 3)]) +def test_slack_subtracts_current_time_and_processing_duration(slack: int, bucket: int) -> None: + job = Job("a", 0, 4, 10 + 4 + slack, 1, "test") + assert encode_state([job], 10)[1] == bucket + + +@pytest.mark.parametrize( + ("short", "total", "bucket"), [(0, 3, 0), (1, 4, 0), (1, 3, 1), (1, 2, 1), (2, 3, 2), (3, 3, 2)] +) +def test_short_job_proportion_buckets(short: int, total: int, bucket: int) -> None: + jobs = [Job(str(i), 0, 2 if i < short else 3, 30, 1, "test") for i in range(total)] + assert encode_state(jobs, 0)[2] == bucket + + +def test_state_uses_minimum_slack_and_detects_high_priority() -> None: + jobs = [Job("a", 0, 4, 14, 1, "test"), Job("b", 0, 6, 14, 4, "test")] + assert encode_state(jobs, 10) == (1, 0, 0, True) + with pytest.raises(ValueError, match="nonempty"): + encode_state([], 0) + + +def test_one_transition_executes_one_job_and_returns_its_cost() -> None: + jobs = [ + Job("first", 2, 4, 5, 4, "test"), + Job("arriving", 3, 1, 10, 1, "test"), + Job("future", 20, 1, 25, 1, "test"), + ] + simulation = Simulation(jobs) + assert simulation.current_time == 2 + assert simulation.waiting_jobs == (jobs[0],) + transition = step_action(simulation, 1) + assert (transition.job.start_time, transition.job.completion_time) == (2, 6) + assert transition.reward == pytest.approx(-4 * (1 + 0.1 * 4)) + assert not transition.done + assert transition.next_state == encode_state([jobs[1]], 6) + assert simulation.waiting_jobs == (jobs[1],) + assert len(simulation.result().jobs) == 1 + step_action(simulation, 0) + assert simulation.current_time == 20 + assert simulation.waiting_jobs == (jobs[2],) + terminal = step_action(simulation, 3) + assert terminal.done and terminal.next_state is None + assert terminal.job.completion_time == 21 + with pytest.raises(ValueError, match="completed"): + step_action(simulation, 0) + + +@pytest.mark.parametrize(("action", "expected"), [(0, "a"), (1, "b"), (2, "c"), (3, "d")]) +def test_action_indices_invoke_the_corresponding_baseline(action: int, expected: str) -> None: + jobs = [ + Job("a", 0, 5, 20, 1, "test"), + Job("b", 0, 1, 30, 1, "test"), + Job("c", 0, 3, 4, 2, "test"), + Job("d", 0, 4, 25, 4, "test"), + ] + transition = step_action(Simulation(jobs), action) + assert transition.job.job.job_id == expected + + +@pytest.mark.parametrize("action", [-1, 4]) +def test_invalid_action_does_not_execute_a_job(action: int) -> None: + simulation = Simulation([Job("a", 0, 2, 10, 1, "test")]) + with pytest.raises(ValueError, match="index"): + step_action(simulation, action) + assert simulation.current_time == 0 + assert simulation.result().jobs == () + + +def test_q_updates_use_next_max_and_no_bootstrap_at_terminal() -> None: + agent = QLearningAgent(seed=0, config=LearningConfig(learning_rate=0.5, gamma=0.9)) + state = (0, 1, 2, False) + next_state = (1, 1, 0, True) + agent.q_table[state] = [1.0, 0.0, 0.0, 0.0] + agent.q_table[next_state] = [2.0, 4.0, 3.0, 1.0] + agent.update(state, 0, -2.0, next_state) + assert agent.q_table[state][0] == pytest.approx(1 + 0.5 * (-2 + 0.9 * 4 - 1)) + agent.update(state, 0, -2.0, None) + assert agent.q_table[state] == pytest.approx([-0.35, 0.0, 0.0, 0.0]) + + +def test_training_updates_the_single_terminal_job() -> None: + config = WorkloadConfig(num_jobs=1, duration_range=(2, 2), deadline_factor=10) + agent = train( + episodes=1, + seed=7, + workload_config=config, + learning_config=LearningConfig(learning_rate=1, epsilon_start=0, epsilon_end=0), + ) + job = WORKLOADS["mixed"](7, config)[0] + state = encode_state([job], 0) + assert agent.q_table[state] == pytest.approx([-job.priority_weight * 0.2, 0, 0, 0]) + + +def test_decay_is_per_episode_and_stops_at_floor() -> None: + config = LearningConfig(epsilon_start=0.8, epsilon_end=0.1, epsilon_decay=0.5) + assert [config.epsilon(i) for i in range(5)] == pytest.approx([0.8, 0.4, 0.2, 0.1, 0.1]) + + +def test_training_is_reproducible_and_does_not_use_global_rng() -> None: + global_state = random.getstate() + config = WorkloadConfig(num_jobs=15) + first = train(episodes=12, seed=42, workload_config=config) + replay = train(episodes=12, seed=42, workload_config=config) + other = train(episodes=12, seed=43, workload_config=config) + assert first.q_table == replay.q_table + assert first.training == replay.training + assert first.q_table != other.q_table + assert random.getstate() == global_state + + +def test_save_load_preserves_predictions_metadata_and_metrics(tmp_path: Path) -> None: + agent = train(episodes=10, seed=42, workload_config=WorkloadConfig(num_jobs=12)) + path = tmp_path / "policy.json" + agent.save(path) + loaded = QLearningAgent.load(path) + assert loaded.q_table == agent.q_table + assert loaded.config == agent.config + assert loaded.training == agent.training + assert all(loaded.greedy_action(state) == agent.greedy_action(state) for state in agent.q_table) + assert loaded.greedy_action((3, 3, 2, False)) == agent.greedy_action((3, 3, 2, False)) + assert evaluate(loaded, episodes=3) == evaluate(agent, episodes=3) + data = json.loads(path.read_text()) + assert data["actions"] == list(ACTION_NAMES) + assert data["state_encoding"]["slack"] == "deadline - current_time - duration" + assert data["training"]["seed"] == 42 + + +def test_evaluation_is_greedy_frozen_and_uses_held_out_seeds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = train(episodes=3, seed=42, workload_config=WorkloadConfig(num_jobs=8)) + jobs = [Job("long", 0, 5, 30, 1, "test"), Job("short", 0, 1, 30, 1, "test")] + state = encode_state(jobs, 0) + agent.q_table[state] = [0.0, 2.0, 1.0, 0.0] + assert simulate(jobs, agent.policy).jobs[0].job.job_id == "short" + unseen = (3, 3, 2, False) + assert agent.greedy_action(unseen) == 0 + before = copy.deepcopy(agent.q_table) + rng_before = agent._rng.getstate() + + def fail(*args: object, **kwargs: object) -> None: + pytest.fail("evaluation must not explore or learn") + + monkeypatch.setattr(agent, "choose_action", fail) + monkeypatch.setattr(agent, "update", fail) + metrics = evaluate(agent, episodes=2, workload_config=WorkloadConfig(num_jobs=8)) + expected = [ + record + for seed in (45, 46) + for record in simulate( + WORKLOADS["mixed"](seed, WorkloadConfig(num_jobs=8)), agent.policy + ).jobs + ] + assert metrics == calculate_metrics(expected) + assert agent.q_table == before + assert agent._rng.getstate() == rng_before + assert unseen not in agent.q_table + with pytest.raises(ValueError, match="overlap"): + evaluate(agent, seed=43, episodes=1) + + +def test_episode_reward_is_negative_total_cost() -> None: + jobs = WORKLOADS["bursty"](5, WorkloadConfig(num_jobs=12)) + simulation = Simulation(jobs) + reward = 0.0 + while not simulation.done: + reward += step_action(simulation, 2).reward + assert reward == pytest.approx(-len(jobs) * simulation.result().metrics.primary_cost) + + +def test_incompatible_saved_formulation_is_rejected(tmp_path: Path) -> None: + path = tmp_path / "policy.json" + QLearningAgent(seed=0).save(path) + data = json.loads(path.read_text()) + data["actions"].reverse() + path.write_text(json.dumps(data)) + with pytest.raises(ValueError, match="incompatible"): + QLearningAgent.load(path)