Class: NanoGPT::Web::TrainingWorker

Inherits:
Object
  • Object
show all
Defined in:
lib/nano_gpt/web/training_worker.rb

Overview

Processes all Torch operations on the main thread. The web server runs in a background thread while this worker owns the main thread and processes commands from a Queue.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(training_state:, metrics_store:, sse_notifier:) ⇒ TrainingWorker

Returns a new instance of TrainingWorker.



13
14
15
16
17
18
 
# File 'lib/nano_gpt/web/training_worker.rb', line 13

def initialize(training_state:, metrics_store:, sse_notifier:)
  @training_state = training_state
  @metrics_store = metrics_store
  @sse_notifier = sse_notifier
  @queue = Queue.new
end

Instance Attribute Details

#queue ⇒ Object (readonly)

Returns the value of attribute queue.



11
12
13
 
# File 'lib/nano_gpt/web/training_worker.rb', line 11

def queue
  @queue
end

Instance Method Details

#enqueue(command, **args) ⇒ Object

Enqueue a fire-and-forget training command



38
39
40
 
# File 'lib/nano_gpt/web/training_worker.rb', line 38

def enqueue(command, **args)
  @queue.push({ command: command, args: args })
end

#enqueue_sync(command, **args) ⇒ Object

Enqueue a command and wait for the result (used for generation)



43
44
45
46
47
 
# File 'lib/nano_gpt/web/training_worker.rb', line 43

def enqueue_sync(command, **args)
  result_queue = Queue.new
  @queue.push({ command: command, args: args, result: result_queue })
  result_queue.pop
end

#run ⇒ Object

Run the command loop on the MAIN thread (blocks forever). Call this AFTER starting the web server in a background thread.



22
23
24
25
26
27
28
29
30
31
32
33
34
35
 
# File 'lib/nano_gpt/web/training_worker.rb', line 22

def run
  loop do
    msg = @queue.pop
    case msg[:command]
    when :start then handle_start(**msg[:args])
    when :resume then handle_resume(**msg[:args])
    when :generate then handle_generate(msg)
    when :prepare_dataset then handle_prepare_dataset(msg)
    when :shutdown then break
    end
  end
rescue => e
  puts "Training worker crashed: #{e.message}\n#{e.backtrace.first(5).join("\n")}"
end