Class: PromptObjects::Execution::Scheduler

Inherits:
Object
  • Object
show all
Defined in:
lib/prompt_objects/execution/scheduler.rb

Overview

Thread-aware in-process scheduler. It owns admission, per-thread writer leases, worker slots, and cross-thread priority selection; Runner owns the execution lifecycle after a lease is granted.

Defined Under Namespace

Classes: ClosedError, Entry, RejectedError, Ticket

Instance Method Summary collapse

Constructor Details

#initialize(worker_limit:, max_active_runs:, max_active_per_po:, max_queued_per_thread:, execute:, ready: nil, cancel_queued: nil) ⇒ Scheduler

Returns a new instance of Scheduler.

Raises:

  • (ArgumentError) 


73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 73

def initialize(worker_limit:, max_active_runs:, max_active_per_po:, max_queued_per_thread:,
               execute:, ready: nil, cancel_queued: nil)
  raise ArgumentError, "max_active_runs must be positive" unless max_active_runs.positive?
  raise ArgumentError, "max_active_per_po must be positive" unless max_active_per_po.positive?
  raise ArgumentError, "max_queued_per_thread must be positive" unless max_queued_per_thread.positive?

  @execute = execute
  @ready = ready || ->(_run) { true }
  @cancel_queued = cancel_queued
  @max_active_runs = max_active_runs
  @max_active_per_po = max_active_per_po
  @max_queued_per_thread = max_queued_per_thread
  @queues = Hash.new { |hash, key| hash[key] = [] }
  @entries_by_run_id = {}
  @running_tokens = {}
  @active_po_counts = Hash.new(0)
  @active_threads = {}
  @mutex = Mutex.new
  @condition = ConditionVariable.new
  @closed = false
  @enqueue_sequence = 0
  @workers = Array.new(worker_limit) { start_worker }
  @deadline_monitor = start_deadline_monitor
end

Instance Method Details

#cancel(run_id, reason: "cancelled") ⇒ Object 



126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 126

def cancel(run_id, reason: "cancelled")
  queued = nil
  token = nil
  @mutex.synchronize do
    entry = @entries_by_run_id[run_id]
    return false unless entry

    token = @running_tokens[run_id]
    unless token
      @queues[entry.run.thread_id].delete(entry)
      cleanup_queue(entry.run.thread_id)
      @entries_by_run_id.delete(run_id)
      queued = entry
    end
    @condition.broadcast
  end

  if token
    token.cancel(reason)
  else
    @cancel_queued&.call(queued.run, reason)
    queued.ticket.resolve(:cancelled)
  end
  true
end

#help(run_id) ⇒ Object

Execute a queued run on the caller's current capacity. This is used by a parent joining its child so a one-worker scheduler cannot deadlock with the parent waiting and the child unable to start.



167
168
169
170
171
172
173
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 167

def help(run_id)
  entry, token = claim_queued(run_id)
  return ticket(run_id) unless entry

  execute(entry, token)
  entry.ticket
end

#queued_count ⇒ Object 



152
153
154
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 152

def queued_count
  @mutex.synchronize { @queues.values.sum(&:size) }
end

#running_count ⇒ Object 



156
157
158
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 156

def running_count
  @mutex.synchronize { @running_tokens.size }
end

#shutdown(wait: true) ⇒ Object 



179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 179

def shutdown(wait: true)
  queued = []
  @mutex.synchronize do
    return if @closed

    @closed = true
    queued = @queues.values.flatten
    @queues.clear
    queued.each { |entry| @entries_by_run_id.delete(entry.run.id) }
    @running_tokens.each_value { |token| token.cancel("scheduler shutdown") }
    @condition.broadcast
  end
  queued.each { |entry| entry.ticket.reject(ClosedError.new("scheduler shut down")) }
  @workers.each(&:join) if wait
  @deadline_monitor.join if wait
end

#submit(run, deadline: nil, **options) ⇒ Object 



98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 98

def submit(run, deadline: nil, **options)
  @mutex.synchronize do
    raise ClosedError, "scheduler is closed" if @closed
    return @entries_by_run_id.fetch(run.id).ticket if @entries_by_run_id.key?(run.id)
    raise RejectedError, "workspace active-run limit reached" if @entries_by_run_id.size >= @max_active_runs

    queue = @queues[run.thread_id]
    if queue.size >= @max_queued_per_thread
      raise RejectedError, "thread queue limit reached for #{run.thread_id}"
    end

    @enqueue_sequence += 1
    ticket = Ticket.new(run.id)
    entry = Entry.new(
      run: run,
      ticket: ticket,
      deadline: deadline,
      enqueue_sequence: @enqueue_sequence,
      options: options.freeze
    )
    queue << entry
    queue.sort_by! { |queued| [queued.run.turn_ordinal, queued.enqueue_sequence] }
    @entries_by_run_id[run.id] = entry
    @condition.broadcast
    ticket
  end
end

#ticket(run_id) ⇒ Object 



160
161
162
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 160

def ticket(run_id)
  @mutex.synchronize { @entries_by_run_id[run_id]&.ticket }
end

#wake ⇒ Object 



175
176
177
 
# File 'lib/prompt_objects/execution/scheduler.rb', line 175

def wake
  @mutex.synchronize { @condition.broadcast }
end