From b2afdf6a37f634f56446fd0d990a99a17feda4ba Mon Sep 17 00:00:00 2001 From: Alexey Rybalchenko Date: Fri, 20 Mar 2026 13:27:32 +0100 Subject: [PATCH 1/3] Fix: Increment completed jobs metric when a job actually finishes, not when it is launched --- src/baseline.py | 6 ++---- src/environment.py | 2 +- src/job_management.py | 41 ++++++++++++++++++++++++++--------------- 3 files changed, 29 insertions(+), 20 deletions(-) diff --git a/src/baseline.py b/src/baseline.py index 366a6dc..b692822 100644 --- a/src/baseline.py +++ b/src/baseline.py @@ -61,13 +61,11 @@ def baseline_step( """ job_queue_2d = baseline_state['job_queue'].reshape(-1, 4) - process_ongoing_jobs(baseline_state['nodes'], baseline_cores_available, baseline_running_jobs) + process_ongoing_jobs(baseline_state['nodes'], baseline_cores_available, baseline_running_jobs, metrics, is_baseline=True) # Age helper queue and fill real queue before new arrivals age_backlog_queue(baseline_backlog_queue, metrics, _is_baseline=True) - baseline_next_empty_slot, _ = fill_queue_from_backlog( - job_queue_2d, baseline_backlog_queue, baseline_next_empty_slot - ) + baseline_next_empty_slot, _ = fill_queue_from_backlog(job_queue_2d, baseline_backlog_queue, baseline_next_empty_slot) _new_baseline_jobs, baseline_next_empty_slot, baseline_backlog_dropped = add_new_jobs( job_queue_2d, new_jobs_count, new_jobs_durations, diff --git a/src/environment.py b/src/environment.py index b218363..ec136da 100644 --- a/src/environment.py +++ b/src/environment.py @@ -322,7 +322,7 @@ def step(self, action: np.ndarray) -> tuple[dict[str, np.ndarray], float, bool, # Decrement booked time for nodes and complete running jobs self.env_print("[1] Processing ongoing jobs...") - completed_jobs = process_ongoing_jobs(self.state['nodes'], self.cores_available, self.running_jobs) + completed_jobs = process_ongoing_jobs(self.state['nodes'], self.cores_available, self.running_jobs, self.metrics, is_baseline=False) self.env_print(f"{len(completed_jobs)} jobs completed: [{' '.join(['#' + str(job_id) for job_id in completed_jobs]) if len(completed_jobs) > 0 else ''}]") # Age helper queues (jobs waiting outside the fixed queue) diff --git a/src/job_management.py b/src/job_management.py index d693dc6..ae9ec02 100644 --- a/src/job_management.py +++ b/src/job_management.py @@ -73,20 +73,28 @@ def validate_next_empty(job_queue_2d: np.ndarray, next_empty: int) -> None: assert np.all(job_queue_2d[:next_empty, 0] != 0), "hole before next_empty_slot" -def process_ongoing_jobs(nodes: np.ndarray, cores_available: np.ndarray, running_jobs: dict[int, dict[str, Any]]) -> list[int]: +def process_ongoing_jobs(nodes: np.ndarray, cores_available: np.ndarray, running_jobs: dict[int, dict[str, Any]], metrics: MetricsTracker, is_baseline: bool) -> list[int]: """ Process ongoing jobs: decrement their duration, complete finished jobs, - and release resources. + release resources, and record completion metrics. + + Completion is counted here (when duration hits zero), not at launch time. + 'wait_time' on each job is the time spent in the queue before being launched, + which is the standard HPC metric for scheduler responsiveness. Args: nodes: Array of node states cores_available: Array of available cores per node running_jobs: Dictionary of currently running jobs + metrics: Optional metrics tracker; when provided, completed job counts + and queue wait times are recorded here + is_baseline: Whether finished jobs belong to the baseline simulation Returns: List of completed job IDs """ completed_jobs = [] + completed_wait_time = 0 for job_id, job_data in running_jobs.items(): job_data['duration'] -= 1 @@ -94,6 +102,7 @@ def process_ongoing_jobs(nodes: np.ndarray, cores_available: np.ndarray, running # Check if job is completed if job_data['duration'] <= 0: completed_jobs.append(job_id) + completed_wait_time += int(job_data.get('wait_time', 0)) # Release resources for node_idx, cores_used in job_data['allocation']: cores_available[node_idx] += cores_used @@ -102,6 +111,19 @@ def process_ongoing_jobs(nodes: np.ndarray, cores_available: np.ndarray, running for job_id in completed_jobs: del running_jobs[job_id] + if completed_jobs: + completed_count = len(completed_jobs) + if is_baseline: + metrics.baseline_jobs_completed += completed_count + metrics.baseline_total_job_wait_time += completed_wait_time + metrics.episode_baseline_jobs_completed += completed_count + metrics.episode_baseline_total_job_wait_time += completed_wait_time + else: + metrics.jobs_completed += completed_count + metrics.total_job_wait_time += completed_wait_time + metrics.episode_jobs_completed += completed_count + metrics.episode_total_job_wait_time += completed_wait_time + # Update node times based on remaining jobs # Reset all nodes first for i in range(MAX_NODES): @@ -198,7 +220,7 @@ def assign_jobs_to_available_nodes( running_jobs: Dictionary of currently running jobs next_empty_slot: Index of next empty slot in queue next_job_id: Next available job ID - metrics: MetricsTracker object to update with job completion metrics + metrics: MetricsTracker object to update with drop/rejection counts is_baseline: Whether this is baseline simulation Returns: @@ -229,6 +251,7 @@ def assign_jobs_to_available_nodes( running_jobs[next_job_id] = { "duration": job_duration, "allocation": job_allocation, + "wait_time": int(job_age), # hours spent in queue; recorded at completion time } next_job_id += 1 @@ -239,18 +262,6 @@ def assign_jobs_to_available_nodes( if job_idx < next_empty_slot: next_empty_slot = job_idx - # Track job completion and wait time - if is_baseline: - metrics.baseline_jobs_completed += 1 - metrics.baseline_total_job_wait_time += job_age - metrics.episode_baseline_jobs_completed += 1 - metrics.episode_baseline_total_job_wait_time += job_age - else: - metrics.jobs_completed += 1 - metrics.total_job_wait_time += job_age - metrics.episode_jobs_completed += 1 - metrics.episode_total_job_wait_time += job_age - num_processed_jobs += 1 continue From d4f02e567e7e5b1dac5f172258a6cd1604cf9d8d Mon Sep 17 00:00:00 2001 From: Alexey Rybalchenko Date: Fri, 20 Mar 2026 17:24:24 +0100 Subject: [PATCH 2/3] tests: adjust no longer relevant params --- test/run_all.py | 1 - test/test_sanity_env.py | 13 ++++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/test/run_all.py b/test/run_all.py index f88ecbd..57ac853 100644 --- a/test/run_all.py +++ b/test/run_all.py @@ -15,7 +15,6 @@ ["python", "-m", "test.test_sanity_env", "--steps", "200"], ["python", "-m", "test.test_sanity_env", "--check-gym", "--check-determinism", "--steps", "300"], ["python", "-m", "test.test_sanity_env", "--prices", "data/prices_2023.csv", "--hourly-jobs", "data/allusers-gpu-30.log", "--steps", "300"], - ["python", "-m", "test.test_sanity_env", "--prices", "data/prices_2023.csv", "--hourly-jobs", "data/allusers-gpu-30.log", "--steps", "300", "--carry-over-state"], ["python", "-m", "test.test_sampler_duration", "--print-stats", "--test-samples", "10"], ["python", "-m", "test.test_sampler_hourly", "--file-path", "data/allusers-gpu-30.log", "--test-day"], ["python", "-m", "test.test_sampler_hourly_aggregated", "--file-path", "data/allusers-gpu-30.log"], diff --git a/test/test_sanity_env.py b/test/test_sanity_env.py index a3691a5..ae60adf 100644 --- a/test/test_sanity_env.py +++ b/test/test_sanity_env.py @@ -206,6 +206,7 @@ def parse_args(): p.add_argument("--steps", type=int, default=500) p.add_argument("--episodes", type=int, default=1) p.add_argument("--check-determinism", action="store_true") + p.add_argument("--check-gym", action="store_true", help="Run stable-baselines3 check_env() on the environment.") # mirror train.py-ish knobs (mostly optional) p.add_argument("--session", default="sanity") p.add_argument("--render", type=str, default="none", choices=["human", "none"]) @@ -380,7 +381,17 @@ def cmp(name, a, b): determinism_test(lambda: make_env_with_carry(), seed=args.seed, n_steps=min(args.steps, 500)) print("[OK] determinism test passed") - # 4) Carry-over continuity + # 4) Gym interface check (optional) + if args.check_gym: + from stable_baselines3.common.env_checker import check_env + gym_env = make_env_with_carry() + try: + check_env(gym_env, warn=True) + finally: + gym_env.close() + print("[OK] gym check passed") + + # 5) Carry-over continuity carry_over_test(lambda: make_env_with_carry(), seed=args.seed, n_steps=min(args.steps, 10)) print("[OK] carry-over continuity test passed") From 8454317b5009f289d923c5235d4cb84020f8e3dd Mon Sep 17 00:00:00 2001 From: Alexey Rybalchenko Date: Fri, 20 Mar 2026 18:01:27 +0100 Subject: [PATCH 3/3] Switch power model to proportional per-core Replace binary node-level power (idle or fully-used per node) with a proportional model: all on-nodes draw idle baseline, compute delta scales linearly with cores_used/CORES_PER_NODE. Updates power_cost(), power_consumption_mwh(), energy efficiency reward, and all call sites in environment.py and baseline.py. --- src/baseline.py | 17 +++++------ src/environment.py | 4 +-- src/job_management.py | 3 ++ src/reward_calculation.py | 62 +++++++++++++++++++++------------------ 4 files changed, 46 insertions(+), 40 deletions(-) diff --git a/src/baseline.py b/src/baseline.py index b692822..88e7638 100644 --- a/src/baseline.py +++ b/src/baseline.py @@ -92,10 +92,10 @@ def baseline_step( baseline_running_jobs, baseline_next_empty_slot, next_job_id, metrics, is_baseline=True ) - num_used_nodes = np.sum(baseline_state['nodes'] > 0) - num_on_nodes = np.sum(baseline_state['nodes'] > -1) + num_used_nodes = int(np.sum(baseline_state['nodes'] > 0)) + num_on_nodes = int(np.sum(baseline_state['nodes'] > -1)) num_idle_nodes = num_on_nodes - num_used_nodes - num_unprocessed_jobs = np.sum(job_queue_2d[:, 0] > 0) + num_unprocessed_jobs = int(np.sum(job_queue_2d[:, 0] > 0)) # Track baseline max queue size (queue only, without backlog) if num_unprocessed_jobs > metrics.baseline_max_queue_size_reached: @@ -112,14 +112,13 @@ def baseline_step( baseline_state['job_queue'] = job_queue_2d.flatten() - baseline_cost = power_cost(num_used_nodes, num_idle_nodes, current_price) + num_used_cores = num_on_nodes * CORES_PER_NODE - np.sum(baseline_cores_available) + baseline_cost = power_cost(num_on_nodes, num_used_cores, current_price) env_print(f" > baseline_cost: €{baseline_cost:.4f} | used nodes: {num_used_nodes}, idle nodes: {num_idle_nodes}") - baseline_cost_off = power_cost(num_used_nodes, 0, current_price) + baseline_cost_off = power_cost(num_used_nodes, num_used_cores, current_price) env_print(f" > baseline_cost_off: €{baseline_cost_off:.4f} | used nodes: {num_used_nodes}, idle nodes: 0") - baseline_power_mwh = power_consumption_mwh(num_used_nodes, num_idle_nodes) - baseline_power_off_mwh = power_consumption_mwh(num_used_nodes, 0) - - num_used_cores = num_on_nodes * CORES_PER_NODE - np.sum(baseline_cores_available) + baseline_power_mwh = power_consumption_mwh(num_on_nodes, num_used_cores) + baseline_power_off_mwh = power_consumption_mwh(num_used_nodes, num_used_cores) return ( baseline_cost, baseline_cost_off, diff --git a/src/environment.py b/src/environment.py index ec136da..e3770b0 100644 --- a/src/environment.py +++ b/src/environment.py @@ -472,7 +472,7 @@ def step(self, action: np.ndarray) -> tuple[dict[str, np.ndarray], float, bool, self.metrics.episode_baseline_used_cores.append(baseline_num_used_cores) step_reward, step_cost, eff_reward_norm, price_reward, idle_penalty_norm, job_age_penalty_norm = self.reward_calculator.calculate( - num_used_nodes, num_idle_nodes, current_price, average_future_price, + num_used_nodes, num_idle_nodes, num_used_cores, current_price, average_future_price, num_off_nodes, num_launched_jobs, num_node_changes, job_queue_2d, num_unprocessed_jobs, self.weights, num_dropped_this_step, self.env_print ) @@ -480,7 +480,7 @@ def step(self, action: np.ndarray) -> tuple[dict[str, np.ndarray], float, bool, self.metrics.episode_reward += step_reward self.metrics.total_cost += step_cost self.metrics.episode_total_cost += step_cost - step_power_mwh = power_consumption_mwh(num_used_nodes, num_idle_nodes) + step_power_mwh = power_consumption_mwh(num_on_nodes, num_used_cores) self.metrics.total_power_consumption_mwh += step_power_mwh self.metrics.episode_total_power_consumption_mwh += step_power_mwh diff --git a/src/job_management.py b/src/job_management.py index ae9ec02..05fafc0 100644 --- a/src/job_management.py +++ b/src/job_management.py @@ -113,6 +113,9 @@ def process_ongoing_jobs(nodes: np.ndarray, cores_available: np.ndarray, running if completed_jobs: completed_count = len(completed_jobs) + # episode_jobs_completed can exceed episode_jobs_submitted: jobs submitted + # in the previous episode and carried over in running_jobs complete here. + # This is expected — it reflects the continuous simulation and is informative. if is_baseline: metrics.baseline_jobs_completed += completed_count metrics.baseline_total_job_wait_time += completed_wait_time diff --git a/src/reward_calculation.py b/src/reward_calculation.py index d067953..0347a7e 100644 --- a/src/reward_calculation.py +++ b/src/reward_calculation.py @@ -5,49 +5,50 @@ import numpy as np from src.config import ( - COST_IDLE_MW, COST_USED_MW, PENALTY_IDLE_NODE, + COST_IDLE_MW, COST_USED_MW, CORES_PER_NODE, PENALTY_IDLE_NODE, PENALTY_DROPPED_JOB, MAX_NODES, MAX_NEW_JOBS_PER_HOUR, WEEK_HOURS ) from src.prices import Prices from src.weights import Weights -def power_cost(num_used_nodes: int, num_idle_nodes: int, current_price: float) -> float: +def power_cost(num_on_nodes: int, cores_used: int, current_price: float) -> float: """ Calculate power cost based on node usage and current electricity price. + Proportional model: all on-nodes draw idle power, plus additional compute + power scaled linearly by actual core utilization. + Formula: (COST_IDLE_MW * num_on + (COST_USED_MW - COST_IDLE_MW) * cores_used/CORES_PER_NODE) * price + Args: - num_used_nodes: Number of nodes with jobs running - num_idle_nodes: Number of idle (on but unused) nodes + num_on_nodes: Number of nodes that are on (used + idle) + cores_used: Total number of cores actively running jobs current_price: Current electricity price Returns: Total power cost """ - idle_cost = COST_IDLE_MW * current_price * num_idle_nodes - usage_cost = COST_USED_MW * current_price * num_used_nodes - total_cost = idle_cost + usage_cost - return total_cost + return (COST_IDLE_MW * num_on_nodes + (COST_USED_MW - COST_IDLE_MW) * (cores_used / CORES_PER_NODE)) * current_price -def power_consumption_mwh(num_used_nodes: int, num_idle_nodes: int) -> float: +def power_consumption_mwh(num_on_nodes: int, cores_used: int) -> float: """ Calculate energy consumption for one environment step. One environment step equals one hour, so this is both average MW and MWh/step. + Uses the same proportional model as power_cost. Args: - num_used_nodes: Number of nodes with jobs running - num_idle_nodes: Number of idle (on but unused) nodes + num_on_nodes: Number of nodes that are on (used + idle) + cores_used: Total number of cores actively running jobs Returns: Energy consumption in MWh for this step """ - return COST_IDLE_MW * num_idle_nodes + COST_USED_MW * num_used_nodes + return COST_IDLE_MW * num_on_nodes + (COST_USED_MW - COST_IDLE_MW) * (cores_used / CORES_PER_NODE) class RewardCalculator: - """Calculates rewards with pre-computed normalization bounds.""" # Faster response so price signal reacts on the same horizon as node-efficiency actions. # Price scaling uses active used nodes as work proxy, matching efficiency semantics. @@ -80,9 +81,9 @@ def __init__(self, prices: Prices) -> None: def _compute_bounds(self) -> None: """Compute min/max bounds for reward normalization.""" - # Efficiency bounds - cost_for_min_efficiency = power_cost(0, MAX_NODES, self.prices.MAX_PRICE) - cost_for_max_efficiency = power_cost(MAX_NODES, 0, self.prices.MIN_PRICE) + # Efficiency bounds (for legacy _reward_efficiency_normalized only) + cost_for_min_efficiency = power_cost(MAX_NODES, 0, self.prices.MAX_PRICE) + cost_for_max_efficiency = power_cost(MAX_NODES, MAX_NODES * CORES_PER_NODE, self.prices.MIN_PRICE) self._min_efficiency_reward = self._reward_efficiency(0, cost_for_min_efficiency) self._max_efficiency_reward = max(1.0, self._reward_efficiency(MAX_NODES, cost_for_max_efficiency)) @@ -153,7 +154,7 @@ def _reward_price(self, current_price: float, average_future_price: float, num_u - Saturates quickly with better-than-context prices and used nodes. - Always applies overdrive when current price is negative. """ - + if num_used_nodes <= 0: return 0.0 @@ -230,16 +231,18 @@ def _penalty_job_age_normalized(self, num_off_nodes: int, job_queue_2d: np.ndarr normalized_penalty = -current_penalty return float(np.clip(normalized_penalty, -1, 0)) - def _reward_energy_efficiency_normalized(self, num_used_nodes: int, num_idle_nodes: int) -> float: - '''Redefine meaning of "efficiency". Use purely as "energy efficiency", aka: How much of the energy (in MW) which is currently needed, gets used for work. - NOTE: Original efficiency function was doing 3 things at once. 1. Handled Blackout logic, with (2.) penalty-ish reward delay for unprocessed jobs, while blackout. - But this log1p function would start to become "harsh" only for a very high number of unprocessed. This rewarded shutting everything off. - 3. rewarded used/cost, but cost was defined in units of price. Price reward should handle this solely, otherwise double counting. - Hence, here new efficiency definition.''' - total_work = num_used_nodes * COST_USED_MW + num_idle_nodes * COST_IDLE_MW - if total_work <= 0.0: + def _reward_energy_efficiency_normalized(self, num_on_nodes: int, cores_used: int) -> float: + '''Energy efficiency: fraction of total power draw that goes to actual computation. + + Proportional model: total power = COST_IDLE_MW * num_on + compute_delta * cores/CORES_PER_NODE. + Compute power = compute_delta * cores/CORES_PER_NODE (the portion above idle baseline). + Efficiency = compute_power / total_power, scaled to [-1, 1]. + ''' + compute_power = (COST_USED_MW - COST_IDLE_MW) * (cores_used / CORES_PER_NODE) + total_power = COST_IDLE_MW * num_on_nodes + compute_power + if total_power <= 0.0: return 0.0 # nothing on => no "efficiency" signal - return 2*(float(np.clip((num_used_nodes * COST_USED_MW) / total_work, 0.0, 1.0))) - 1.0 # scale to [-1, 1] so that it can be weighted in either direction without exceeding bounds. + return 2 * float(np.clip(compute_power / total_power, 0.0, 1.0)) - 1.0 # scale to [-1, 1] def _blackout_term(self, num_used_nodes: int, num_idle_nodes: int, num_unprocessed_jobs: int) -> float: """ @@ -260,7 +263,7 @@ def _blackout_term(self, num_used_nodes: int, num_idle_nodes: int, num_unprocess penalty = np.exp(-ratio * SATURATION_FACTOR) - 1.0 return float(np.clip(penalty, -1.0, 0.0)) - def calculate(self, num_used_nodes: int, num_idle_nodes: int, current_price: float, average_future_price: float, + def calculate(self, num_used_nodes: int, num_idle_nodes: int, num_used_cores: int, current_price: float, average_future_price: float, num_off_nodes: int, _num_processed_jobs: int, num_node_changes: int, job_queue_2d: np.ndarray, # noqa: ARG002 - _num_processed_jobs legacy; num_node_changes reserved for future node-change penalty num_unprocessed_jobs: int, weights: Weights, num_dropped_this_step: int, env_print: Callable[..., None]) -> tuple[float, float, float, float, float, float]: @@ -286,8 +289,9 @@ def calculate(self, num_used_nodes: int, num_idle_nodes: int, current_price: flo idle_penalty_norm, job_age_penalty_norm) """ # 0. Energy efficiency. Reward calculation based on Workload (used nodes) (W) / Cost (C) - total_cost = power_cost(num_used_nodes, num_idle_nodes, current_price) - efficiency_reward_norm = self._reward_energy_efficiency_normalized(num_used_nodes, num_idle_nodes) + self._blackout_term(num_used_nodes, num_idle_nodes, num_unprocessed_jobs) + num_on_nodes = num_used_nodes + num_idle_nodes + total_cost = power_cost(num_on_nodes, num_used_cores, current_price) + efficiency_reward_norm = self._reward_energy_efficiency_normalized(num_on_nodes, num_used_cores) + self._blackout_term(num_used_nodes, num_idle_nodes, num_unprocessed_jobs) efficiency_reward_weighted = weights.efficiency_weight * efficiency_reward_norm # 2. Increase reward if current price is favorable and currently used nodes are high.