Skip to content

Carry-over states between episodes, instead of resetting everything to 0. #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
enlorenz wants to merge 2 commits into from
Closed

Carry-over states between episodes, instead of resetting everything to 0. #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f16e150
Add carry-over state support and per-episode metrics. test_sanity_env…
Jan 15, 2026
1a15eea
Add: Carry over State. Some missing parts.
Jan 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions src/baseline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,12 @@ def baseline_step(baseline_state, baseline_cores_available, baseline_running_job
job_queue_2d, new_jobs_count, new_jobs_durations,
new_jobs_nodes, new_jobs_cores, baseline_next_empty_slot
)
metrics.baseline_jobs_submitted += len(new_baseline_jobs)
metrics.baseline_jobs_rejected_queue_full += (new_jobs_count - len(new_baseline_jobs))
metrics['baseline_jobs_submitted'] += len(new_baseline_jobs)
if 'episode_baseline_jobs_submitted' in metrics:
metrics['episode_baseline_jobs_submitted'] += len(new_baseline_jobs)
metrics['baseline_jobs_rejected_queue_full'] += (new_jobs_count - len(new_baseline_jobs))
if 'episode_baseline_jobs_rejected_queue_full' in metrics:
metrics['episode_baseline_jobs_rejected_queue_full'] += (new_jobs_count - len(new_baseline_jobs))

_, baseline_next_empty_slot, _, next_job_id = assign_jobs_to_available_nodes(
job_queue_2d, baseline_state['nodes'], baseline_cores_available,
Expand All @@ -52,8 +56,11 @@ def baseline_step(baseline_state, baseline_cores_available, baseline_running_job
num_unprocessed_jobs = np.sum(job_queue_2d[:, 0] > 0)

# Track baseline max queue size
if num_unprocessed_jobs > metrics.baseline_max_queue_size_reached:
metrics.baseline_max_queue_size_reached = num_unprocessed_jobs
if num_unprocessed_jobs > metrics['baseline_max_queue_size_reached']:
metrics['baseline_max_queue_size_reached'] = num_unprocessed_jobs
if 'episode_baseline_max_queue_size_reached' in metrics:
if num_unprocessed_jobs > metrics['episode_baseline_max_queue_size_reached']:
metrics['episode_baseline_max_queue_size_reached'] = num_unprocessed_jobs

baseline_state['job_queue'] = job_queue_2d.flatten()

Expand Down
151 changes: 132 additions & 19 deletions src/environment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,8 @@ def __init__(self,
skip_plot_job_queue,
steps_per_iteration,
evaluation_mode=False,
workload_gen=None):
workload_gen=None,
carry_over_state=False):
super().__init__()

self.weights = weights
Expand Down Expand Up @@ -119,6 +120,7 @@ def __init__(self,
self.np_random = None
self._seed = None
self.workload_gen = workload_gen
self.carry_over_state = carry_over_state

if self.external_durations:
durations_sampler.init(self.external_durations)
Expand Down Expand Up @@ -152,6 +154,9 @@ def __init__(self,
# Initialize reward calculator
self.reward_calculator = RewardCalculator(self.prices)

self.reset_timeline_state()
self.metrics.reset_episode_metrics()

Comment on lines +157 to +159

@coderabbitai coderabbitai Bot Jan 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

# First, find the Prices class and check __init__
rg -n "class Prices" src/ -A 20 | head -50

Repository: FairRootGroup/powersched

Length of output: 2426

🏁 Script executed:

# Also check the environment.py file at the mentioned lines
sed -n '160,190p' src/environment.py

Repository: FairRootGroup/powersched

Length of output: 1410

🏁 Script executed:

# And check lines 268-277 as mentioned
sed -n '260,285p' src/environment.py

Repository: FairRootGroup/powersched

Length of output: 1050

🏁 Script executed:

# Now let's look at the reset_timeline_state implementation
rg -n "def reset_timeline_state" src/environment.py -A 10

Repository: FairRootGroup/powersched

Length of output: 577

🏁 Script executed:

# Search for reset method in Prices class
rg -n "def reset" src/prices_deterministic.py -A 15

Repository: FairRootGroup/powersched

Length of output: 796

🏁 Script executed:

# Check the full environment.__init__ to see the complete sequence
rg -n "def __init__" src/environment.py -A 80 | head -120

Repository: FairRootGroup/powersched

Length of output: 3707

🏁 Script executed:

# Also check if Prices.reset() is called anywhere in environment.__init__
rg -n "self.prices.reset" src/environment.py

Repository: FairRootGroup/powersched

Length of output: 183

🏁 Script executed:

# Get full context around these critical lines
sed -n '115,250p' src/environment.py | cat -n

Repository: FairRootGroup/powersched

Length of output: 6999

reset_timeline_state() attempts to copy self.prices.predicted_prices while it is still None.

In __init__ at line 175, reset_timeline_state() is called before Prices.reset() is invoked. However, Prices.__init__ initializes predicted_prices = None (line 19), and it only gets populated when Prices.reset() is called. This causes AttributeError: 'NoneType' object has no attribute 'copy' at line 277 when trying to execute self.prices.predicted_prices.copy().

Call self.prices.reset(start_index=0) before self.reset_timeline_state() in __init__, matching the pattern used in the reset() method (lines 120–123).

🤖 Prompt for AI Agents 
In `@src/environment.py` around lines 175 - 177, The constructor calls
reset_timeline_state() before Prices.reset(), so reset_timeline_state() tries to
copy self.prices.predicted_prices while it's still None; fix by invoking
self.prices.reset(start_index=0) before calling self.reset_timeline_state() in
__init__ (same order used in reset()), ensuring predicted_prices is initialized
before reset_timeline_state() uses self.prices.predicted_prices.copy().

# actions: - change number of available nodes:
# action_type: 0: decrease, 1: maintain, 2: increase
# action_magnitude: 0-MAX_CHANGE (+1ed in the action)
Expand Down Expand Up @@ -192,21 +197,59 @@ def reset(self, seed=None, options=None):
self.episode_idx += 1

# Reset metrics
self.metrics.reset_state_metrics()

# Choose starting index in the external price series
if self.prices is not None and getattr(self.prices, "external_prices", None) is not None:
n_prices = len(self.prices.external_prices)
episode_span = EPISODE_HOURS

# Episode k starts at hour k * episode_span (wrapping around the year)
start_index = (self.episode_idx * episode_span) % n_prices
if options and "price_start_index" in options: # For testing Purposes. Leave out 'options' to advance episode.
start_index = int(options["price_start_index"]) % n_prices
self.prices.reset(start_index=start_index)
if self.carry_over_state:
self.metrics.reset_episode_metrics()
else:
# Synthetic prices or no external prices
self.prices.reset(start_index=0)
self.metrics.reset_state_metrics()

self.metrics.current_hour = 0

if not self.carry_over_state:
# Choose starting index in the external price series
if self.prices is not None and getattr(self.prices, "external_prices", None) is not None:

@rbx rbx Jan 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do not use any getattr anywhere.
First of all we control the entire code, so we don't have to guess if something exists or not.
Second, in case for some reason the attribute does not exist, we want to have it fail so that we see the problem, rather than have it silently fill some default value and pretend that everything is okay.

n_prices = len(self.prices.external_prices)
episode_span = EPISODE_HOURS

# Episode k starts at hour k * episode_span (wrapping around the year)
start_index = (self.episode_idx * episode_span) % n_prices
if options and "price_start_index" in options: # For testing Purposes. Leave out 'options' to advance episode.
start_index = int(options["price_start_index"]) % n_prices
self.prices.reset(start_index=start_index)
else:
# Synthetic prices or no external prices
self.prices.reset(start_index=0)

self.state = {
# Initialize all nodes to be 'online but free' (0)
'nodes': np.zeros(MAX_NODES, dtype=np.int32),
# Initialize job queue to be empty
'job_queue': np.zeros((MAX_QUEUE_SIZE * 4), dtype=np.int32),
# Initialize predicted prices array
'predicted_prices': self.prices.predicted_prices.copy(),
}

self.baseline_state = {
'nodes': np.zeros(MAX_NODES, dtype=np.int32),
'job_queue': np.zeros((MAX_QUEUE_SIZE * 4), dtype=np.int32),
}

self.cores_available = np.full(MAX_NODES, CORES_PER_NODE, dtype=np.int32)
self.baseline_cores_available = np.full(MAX_NODES, CORES_PER_NODE, dtype=np.int32)

# Job tracking: { job_id: {'duration': remaining_hours, 'allocation': [(node_idx1, cores1), ...]}, ... }
self.running_jobs = {}
self.baseline_running_jobs = {}

self.next_job_id = 0 # shared between baseline and normal jobs

# Track next empty slot in job queue for O(1) insertion
self.next_empty_slot = 0
self.baseline_next_empty_slot = 0

return self.state, {}

def reset_timeline_state(self):
self.metrics.current_hour = 0

self.state = {
# Initialize all nodes to be 'online but free' (0)
Expand All @@ -224,7 +267,6 @@ def reset(self, seed=None, options=None):

self.cores_available = np.full(MAX_NODES, CORES_PER_NODE, dtype=np.int32)
self.baseline_cores_available = np.full(MAX_NODES, CORES_PER_NODE, dtype=np.int32)

# Job tracking: { job_id: {'duration': remaining_hours, 'allocation': [(node_idx1, cores1), ...]}, ... }
self.running_jobs = {}
self.baseline_running_jobs = {}
Expand All @@ -235,11 +277,10 @@ def reset(self, seed=None, options=None):
self.next_empty_slot = 0
self.baseline_next_empty_slot = 0

return self.state, {}

def step(self, action):
self.current_step += 1
self.metrics.current_hour += 1
self.metrics.total_time_hours += 1
if self.metrics.current_hour == 1:
self.current_episode += 1
self.env_print(Fore.GREEN + f"\n[[[ Starting episode: {self.current_episode}, step: {self.current_step}, hour: {self.metrics.current_hour}" + Fore.RESET)
Expand Down Expand Up @@ -272,7 +313,9 @@ def step(self, action):
new_jobs_nodes, new_jobs_cores, self.next_empty_slot
)
self.metrics.jobs_submitted += len(new_jobs)
self.metrics.episode_jobs_submitted += len(new_jobs)
self.metrics.jobs_rejected_queue_full += (new_jobs_count - len(new_jobs))
self.metrics.episode_jobs_rejected_queue_full += (new_jobs_count - len(new_jobs))

self.env_print("nodes: ", np.array2string(self.state['nodes'], separator=' ', max_line_width=np.inf))
self.env_print(f"cores_available: {np.array2string(self.cores_available, separator=' ', max_line_width=np.inf)} ({np.sum(self.cores_available)})")
Expand All @@ -288,11 +331,38 @@ def step(self, action):
# Assign jobs to available nodes
self.env_print(f"[4] Assigning jobs to available nodes...")

# Create metrics dict for job assignment
job_metrics = {
'jobs_completed': self.metrics.jobs_completed,
'total_job_wait_time': self.metrics.total_job_wait_time,
'jobs_dropped': self.metrics.jobs_dropped,
'dropped_this_episode': self.metrics.dropped_this_episode,
'baseline_jobs_completed': self.metrics.baseline_jobs_completed,
'baseline_total_job_wait_time': self.metrics.baseline_total_job_wait_time,
'baseline_jobs_dropped': self.metrics.baseline_jobs_dropped,
'baseline_dropped_this_episode': self.metrics.baseline_dropped_this_episode,
'episode_jobs_completed': self.metrics.episode_jobs_completed,
'episode_total_job_wait_time': self.metrics.episode_total_job_wait_time,
'episode_jobs_dropped': self.metrics.episode_jobs_dropped,
'episode_baseline_jobs_completed': self.metrics.episode_baseline_jobs_completed,
'episode_baseline_total_job_wait_time': self.metrics.episode_baseline_total_job_wait_time,
'episode_baseline_jobs_dropped': self.metrics.episode_baseline_jobs_dropped,
}

num_launched_jobs, self.next_empty_slot, num_dropped_this_step, self.next_job_id = assign_jobs_to_available_nodes(
job_queue_2d, self.state['nodes'], self.cores_available, self.running_jobs,
self.next_empty_slot, self.next_job_id, self.metrics, is_baseline=False
)

# Update metrics from job_metrics dict
self.metrics.jobs_completed = job_metrics['jobs_completed']
self.metrics.total_job_wait_time = job_metrics['total_job_wait_time']
self.metrics.jobs_dropped = job_metrics['jobs_dropped']
self.metrics.dropped_this_episode = job_metrics['dropped_this_episode']
self.metrics.episode_jobs_completed = job_metrics['episode_jobs_completed']
self.metrics.episode_total_job_wait_time = job_metrics['episode_total_job_wait_time']
self.metrics.episode_jobs_dropped = job_metrics['episode_jobs_dropped']

Comment on lines +334 to +365

@coderabbitai coderabbitai Bot Jan 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Pass the metrics dicts into job/baseline helpers (current code discards updates).
job_metrics / baseline_metrics are created but not passed, and later overwrite updated counters. This also risks TypeError if MetricsTracker isn’t subscriptable.

🐛 Proposed fix 
-        num_launched_jobs, self.next_empty_slot, num_dropped_this_step, self.next_job_id = assign_jobs_to_available_nodes(
-            job_queue_2d, self.state['nodes'], self.cores_available, self.running_jobs,
-            self.next_empty_slot, self.next_job_id, self.metrics, is_baseline=False
-        )
+        num_launched_jobs, self.next_empty_slot, num_dropped_this_step, self.next_job_id = assign_jobs_to_available_nodes(
+            job_queue_2d, self.state['nodes'], self.cores_available, self.running_jobs,
+            self.next_empty_slot, self.next_job_id, job_metrics, is_baseline=False
+        )
@@
-        baseline_cost, baseline_cost_off, self.baseline_next_empty_slot, self.next_job_id = baseline_step(
-            self.baseline_state, self.baseline_cores_available, self.baseline_running_jobs,
-            current_price, new_jobs_count, new_jobs_durations, new_jobs_nodes, new_jobs_cores,
-            self.baseline_next_empty_slot, self.next_job_id, self.metrics, self.env_print
-        )
+        baseline_cost, baseline_cost_off, self.baseline_next_empty_slot, self.next_job_id = baseline_step(
+            self.baseline_state, self.baseline_cores_available, self.baseline_running_jobs,
+            current_price, new_jobs_count, new_jobs_durations, new_jobs_nodes, new_jobs_cores,
+            self.baseline_next_empty_slot, self.next_job_id, baseline_metrics, self.env_print
+        )

Also applies to: 392-427

🤖 Prompt for AI Agents 
In `@src/environment.py` around lines 334 - 365, The job_metrics and
baseline_metrics dicts are created but not passed into the helper, causing
updated counters to be lost and risking TypeError; change the call to
assign_jobs_to_available_nodes to accept a metrics dict (e.g., pass job_metrics
for is_baseline=False and baseline_metrics for is_baseline=True) and update the
helper signature to either mutate that dict or return an updated metrics dict
(and assign it back to self.metrics or merge values), then remove the later
naive overwrite of self.metrics with the original job_metrics; update the
counterpart baseline call (the block referenced around 392-427) the same way so
metrics are preserved and no subscript errors occur when MetricsTracker is not
subscriptable.

self.env_print(f" {num_launched_jobs} jobs launched")

# Calculate node utilization stats
Expand All @@ -313,18 +383,53 @@ def step(self, action):
# Track max queue size
if num_unprocessed_jobs > self.metrics.max_queue_size_reached:
self.metrics.max_queue_size_reached = num_unprocessed_jobs
if num_unprocessed_jobs > self.metrics.episode_max_queue_size_reached:
self.metrics.episode_max_queue_size_reached = num_unprocessed_jobs

self.env_print(f"[5] Calculating reward...")

# Baseline step
baseline_metrics = {
'baseline_jobs_submitted': self.metrics.baseline_jobs_submitted,
'baseline_jobs_rejected_queue_full': self.metrics.baseline_jobs_rejected_queue_full,
'baseline_jobs_completed': self.metrics.baseline_jobs_completed,
'baseline_total_job_wait_time': self.metrics.baseline_total_job_wait_time,
'baseline_jobs_dropped': self.metrics.baseline_jobs_dropped,
'baseline_dropped_this_episode': self.metrics.baseline_dropped_this_episode,
'baseline_max_queue_size_reached': self.metrics.baseline_max_queue_size_reached,
'episode_baseline_jobs_submitted': self.metrics.episode_baseline_jobs_submitted,
'episode_baseline_jobs_rejected_queue_full': self.metrics.episode_baseline_jobs_rejected_queue_full,
'episode_baseline_jobs_completed': self.metrics.episode_baseline_jobs_completed,
'episode_baseline_total_job_wait_time': self.metrics.episode_baseline_total_job_wait_time,
'episode_baseline_jobs_dropped': self.metrics.episode_baseline_jobs_dropped,
'episode_baseline_max_queue_size_reached': self.metrics.episode_baseline_max_queue_size_reached,
}

baseline_cost, baseline_cost_off, self.baseline_next_empty_slot, self.next_job_id = baseline_step(
self.baseline_state, self.baseline_cores_available, self.baseline_running_jobs,
current_price, new_jobs_count, new_jobs_durations, new_jobs_nodes, new_jobs_cores,
self.baseline_next_empty_slot, self.next_job_id, self.metrics, self.env_print
)

# Update metrics from baseline_metrics dict
self.metrics.baseline_jobs_submitted = baseline_metrics['baseline_jobs_submitted']
self.metrics.baseline_jobs_rejected_queue_full = baseline_metrics['baseline_jobs_rejected_queue_full']
self.metrics.baseline_jobs_completed = baseline_metrics['baseline_jobs_completed']
self.metrics.baseline_total_job_wait_time = baseline_metrics['baseline_total_job_wait_time']
self.metrics.baseline_jobs_dropped = baseline_metrics['baseline_jobs_dropped']
self.metrics.baseline_dropped_this_episode = baseline_metrics['baseline_dropped_this_episode']
self.metrics.baseline_max_queue_size_reached = baseline_metrics['baseline_max_queue_size_reached']
self.metrics.episode_baseline_jobs_submitted = baseline_metrics['episode_baseline_jobs_submitted']
self.metrics.episode_baseline_jobs_rejected_queue_full = baseline_metrics['episode_baseline_jobs_rejected_queue_full']
self.metrics.episode_baseline_jobs_completed = baseline_metrics['episode_baseline_jobs_completed']
self.metrics.episode_baseline_total_job_wait_time = baseline_metrics['episode_baseline_total_job_wait_time']
self.metrics.episode_baseline_jobs_dropped = baseline_metrics['episode_baseline_jobs_dropped']
self.metrics.episode_baseline_max_queue_size_reached = baseline_metrics['episode_baseline_max_queue_size_reached']

self.metrics.baseline_cost += baseline_cost
self.metrics.baseline_cost_off += baseline_cost_off
self.metrics.episode_baseline_cost += baseline_cost
self.metrics.episode_baseline_cost_off += baseline_cost_off

step_reward, step_cost, eff_reward_norm, price_reward_norm, idle_penalty_norm, job_age_penalty_norm = self.reward_calculator.calculate(
num_used_nodes, num_idle_nodes, current_price, average_future_price,
Expand All @@ -334,6 +439,7 @@ def step(self, action):

self.metrics.episode_reward += step_reward
self.metrics.total_cost += step_cost
self.metrics.episode_total_cost += step_cost

# Store normalized reward components for plotting
self.metrics.eff_rewards.append(eff_reward_norm * 100)
Expand Down Expand Up @@ -385,4 +491,11 @@ def step(self, action):

self.env_print(Fore.GREEN + f"]]]" + Fore.RESET)

return self.state, step_reward, terminated, truncated, {}
info = {
"step_cost": float(step_cost),
"num_unprocessed_jobs": int(num_unprocessed_jobs),
"num_on_nodes": int(num_on_nodes),
"dropped_this_episode": int(getattr(self.metrics, "dropped_this_episode", 0)),
}

return self.state, step_reward, terminated, truncated, info
26 changes: 18 additions & 8 deletions src/job_management.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,11 +154,17 @@ def assign_jobs_to_available_nodes(job_queue_2d, nodes, cores_available, running

# Track job completion and wait time
if is_baseline:
metrics.baseline_jobs_completed += 1
metrics.baseline_total_job_wait_time += job_age
metrics['baseline_jobs_completed'] += 1
metrics['baseline_total_job_wait_time'] += job_age
if 'episode_baseline_jobs_completed' in metrics:
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['jobs_completed'] += 1
metrics['total_job_wait_time'] += job_age
if 'episode_jobs_completed' in metrics:
metrics['episode_jobs_completed'] += 1
metrics['episode_total_job_wait_time'] += job_age

num_processed_jobs += 1
continue
Expand All @@ -176,11 +182,15 @@ def assign_jobs_to_available_nodes(job_queue_2d, nodes, cores_available, running
num_dropped += 1

if is_baseline:
metrics.baseline_jobs_dropped += 1
metrics.baseline_dropped_this_episode += 1
metrics['baseline_jobs_dropped'] += 1
metrics['baseline_dropped_this_episode'] += 1
if 'episode_baseline_jobs_dropped' in metrics:
metrics['episode_baseline_jobs_dropped'] += 1
else:
metrics.jobs_dropped += 1
metrics.dropped_this_episode += 1
metrics['jobs_dropped'] += 1
metrics['dropped_this_episode'] += 1
if 'episode_jobs_dropped' in metrics:
metrics['episode_jobs_dropped'] += 1
else:
job_queue_2d[job_idx][1] = new_age

Expand Down
Loading