<commit_before>require 'drb'
require 'thread'
require 'yaml'

SERVER_URL = 'druby://localhost:5555'
CONFIG_DIR = File.join(__dir__, './configs')
WORKER_DIR = File.join(__dir__, './workers')
TEMP_F_DIR = File.join(__dir__, './temp')
TEMP_W_DIR = File.join(__dir__, './temp_work')

# Require the worker helper.
require File.join(__dir__, './worker_helper.rb')

# Require all workers.
Dir["#{WORKER_DIR}/*.rb"].each { |file| require file }

################################################################################
#
# Worker classes must be created in WORKER_DIR. If a class needs an YAML config
# file it should be stored in CONFIG_DIR.
#
# Every worker class should the following methods:
#   * def initialize()
#   * def setup(id, config_loader, temp_saver, notifier, *args)
#   * def work()
#
################################################################################

# Worker global controller.
class MasterWorker

  def initialize
    @current_id = 0
    @process_mutex = Mutex.new
    @process_by_id = {}
  end

  def resume_workers
    Dir["#{TEMP_W_DIR}/*.yml"].each do |f|
      yaml = YAML.load(File.read(f))
      worker, args = yaml.first
      start_new_worker(worker, args)
      File.delete(f)
    end
  end

  # Launches new worker.
  def start_new_worker(worker, *args)
    wkr = _instantiate_worker(worker)
    if wkr
      id = _place_worker(wkr)
      _save_worker(id, args)
      wkr.setup(id, WorkerHelper.new(CONFIG_DIR, TEMP_F_DIR), *args)
      puts "Starting worker (#{worker})"
      _start_new_worker(wkr)
      return id
    else
      return nil
    end
  end

  # Frees worker resources.
  def notify_finish(id, msg = nil)
    @process_mutex.synchronize do
      @process_by_id.delete(id)
      File.delete(File.join(TEMP_W_DIR, "#{id}.yml"))
    end
    if msg
      puts "Worker #{id} finished: #{msg}"
    else
      puts "Worker #{id} finished"
    end
  end

  def working?
    true
  end

  private

  # Saves a copy of the worker arguments.
  def _save_worker(id, worker, args)
    File.open(File.join(TEMP_W_DIR, "#{id}.yaml"), 'w') do |f|
      f.write(YAML.dump({ worker => args }))
    end<commit_msg>Resumable workers added.<commit_after>require 'drb'
require 'thread'
require 'yaml'

SERVER_URL = 'druby://localhost:5555'
CONFIG_DIR = File.join(__dir__, './configs')
WORKER_DIR = File.join(__dir__, './workers')
TEMP_F_DIR = File.join(__dir__, './temp')
TEMP_W_DIR = File.join(__dir__, './temp_work')

# Require the worker helper.
require File.join(__dir__, './worker_helper.rb')

# Require all workers.
Dir["#{WORKER_DIR}/*.rb"].each { |file| require file }

################################################################################
#
# Worker classes must be created in WORKER_DIR. If a class needs an YAML config
# file it should be stored in CONFIG_DIR.
#
# Every worker class should the following methods:
#   * def initialize()
#   * def setup(id, config_loader, temp_saver, notifier, *args)
#   * def work()
#
################################################################################

# Worker global controller.
class MasterWorker

  def initialize
    @current_id = 0
    @process_mutex = Mutex.new
    @process_by_id = {}
  end

  def resume_workers
    Dir["#{TEMP_W_DIR}/*.yml"].each do |f|
      yaml = YAML.load(File.read(f))
      worker, args = yaml.first
      File.delete(f)
      start_new_worker(worker, *args)
    end
  end

  # Launches new worker.
  def start_new_worker(worker, *args)
    wkr = _instantiate_worker(worker)
    if wkr
      id = _place_worker(wkr)
      _save_worker(id, worker, args)
      wkr.setup(id, WorkerHelper.new(CONFIG_DIR, TEMP_F_DIR, SERVER_URL), *args)
      puts "Starting worker (#{worker})"
      _start_new_worker(wkr)
      return id
    else
      return nil
    end
  end

  # Frees worker resources.
  def notify_finish(id, msg = nil)
    @process_mutex.synchronize do
      @process_by_id.delete(id)
      File.delete(File.join(TEMP_W_DIR, "#{id}.yml"))
    end
    if msg
      puts "Worker #{id} finished: #{msg}"
    else
      puts "Worker #{id} finished"
    end
  end

  def working?
    true
  end

  private

  # Saves a copy of the worker arguments.
  def _save_worker(id, worker, args)
    File.open(File.join(TEMP_W_DIR, "#{id}.yml"), 'w') do |f|
      f.write(YAML.dump({ worker => args }))
    end
{'fraction_non_alphanumeric': 0.19157264563665669, 'fraction_numerical': 0.002302555836979047, 'mean_word_length': 3.4968944099378882, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4127774', 'n_tokens_mistral': 1468, 'n_tokens_neox': 1434, 'n_words': 367}