<commit_before>require 'yaml'

require 'hive'
require 'hive/file_system'
require 'hive/execution_script'
require 'hive/diagnostic_runner'
require 'hive/messages'
require 'hive/port_allocator'
require 'code_cache'
require 'res'

module Hive
  # The generic worker class
  class Worker
    class InvalidJobReservationError < StandardError
    end

    class DeviceNotReady < StandardError
    end

    class NoPortsAvailable < StandardError
    end

    # Device API Object for device associated with this worker
    attr_accessor :device_api, :queues 

    # The main worker process loop
    def initialize(options)
      @options = options
      @parent_pid = @options['parent_pid']
      @device_id = @options['id']
      @hive_id = @options['hive_id']
      @default_component ||= self.class.to_s
      @current_job_start_time = nil
      @hive_mind ||= mind_meld_klass.new(
        url: Chamber.env.network.hive_mind? ? Chamber.env.network.hive_mind : nil,
        pem: Chamber.env.network.cert ? Chamber.env.network.cert : nil,
        ca_file: Chamber.env.network.cafile ? Chamber.env.network.cafile : nil,
        verify_mode: Chamber.env.network.verify_mode ? Chamber.env.network.verify_mode : nil,
        device: hive_mind_device_identifiers
      )
      @device_identity = @options['device_identity'] || 'unknown-device'
      pid = Process.pid
      $PROGRAM_NAME = "#{@options['name_stub'] || 'WORKER'}.#{pid}"
      @log = Hive::Log.new
      @log.add_logger(
        "#{LOG_DIRECTORY}/#{pid}.#{@device_identity}.log",
        Hive.config.logging.worker_level || 'INFO'
      )
      @log.hive_mind = @hive_mind
      @log.default_progname = @default_component

      self.update_queues

      @port_allocator = (@options.has_key?('port_allocator') ? @options['port_allocator'] : Hive::PortAllocator.new(ports: []))
      
      platform = self.class.to_s.scan(/[^:][^:]*/)[2].downcase
      @diagnostic_runner = Hive::DiagnosticRunner.new(@options, Hive.config.diagnostics, platform, @hive_mind) if Hive.config.diagnostics? && Hive.config.diagnostics[platform]

      Hive::Messages.configure do |config|
        config.base_path = Hive.config.network.scheduler
        config.pem_file = Hive.config.network.cert
        config.ssl_verify_mode = OpenSSL::SSL::VERIFY_NONE
      end

      Signal.trap('TERM') do
        @log.info("Worker terminated")
        exit
      end

      @log.info('Starting worker')
      while keep_worker_running?
        begin
          @log.clear({component: @default_component, level: Hive.config.logging.hm_logs_to_delete})
          update_queues
          poll_queue if diagnostics
        rescue DeviceNotReady => e
          @log.warn("#{e.message}\n");
        rescue StandardError => e
          @log.warn("Worker loop aborted: #{e.message}\n  : #{e.backtrace.join("\n  : ")}")
        end
        sleep Hive.config.timings.worker_loop_interval
      end
      @log.info('Exiting worker')
    end

    # Check the queues for work
    def poll_queue
      @job = reserve_job
      if @job.nil?
        @log.info('No job found')
      else
        @log.info('Job starting')
        begin
          @current_job_start_time = Time.now
          execute_job
        rescue => e
          @log.info("Error running test: #{e.message}\n : #{e.backtrace.join("\n :")}")
        end
        cleanup
      end
    end

    # Try to find and reserve a job
    def reserve_job
      @log.info "Trying to reserve job for queues: #{@queues.join(', ')}"
      job = job_message_klass.reserve(@queues, reservation_details)
      @log.debug "Job: #{job.inspect}"
      raise InvalidJobReservationError.new("Invalid Job Reserved") if ! (job.nil? || job.valid?)
      job
    end

    # Get the correct job class
    # This should usually be replaced in the child class
    def job_message_klass
      @log.info 'Generic job class'
      Hive::Messages::Job
    end

    def mind_meld_klass
      MindMeld::Device
    end

    def reservation_details
      @log.debug "Reservations details: hive_id=#{@hive_id}, worker_pid=#{Process.pid}, device_id=#{@hive_mind.id}"
      { hive_id: @hive_id, worker_pid: Process.pid, device_id: @hive_mind.id }
    end

    # Execute a job
    def execute_job
      # Ensure that a killed worker cleans up correctly
      Signal.trap('TERM') do |s|
        Signal.trap('TERM') {} # Prevent retry signals
        @log.info "Caught TERM signal"
        @log.info "Terminating script, if running"
        @script.terminate if @script
        @log.info "Post-execution cleanup"
        signal_safe_post_script(@job, @file_system, @script)

        # Upload results
        @file_system.finalise_results_directory
        upload_files(@job, @file_system.results_path, @file_system.logs_path)
        set_job_state_to :completed
        @job.error('Worker killed')
        @log.info "Worker terminated"
        exit
      end

      @log.info('Job starting')
      @job.prepare(@hive_mind.id)
      
      exception = nil
      begin
        @log.info "Setting job paths"
        @file_system = Hive::FileSystem.new(@job.job_id, Hive.config.logging.home, @log)
        set_job_state_to :preparing

        if ! @job.repository.to_s.empty?
          @log.info "Checking out the repository"
          @log.debug "  #{@job.repository}"
          @log.debug "  #{@file_system.testbed_path}"
          checkout_code(@job.repository, @file_system.testbed_path, @job.branch)
        end

        @log.info "Initialising execution script"
        @script = Hive::ExecutionScript.new(
          file_system: @file_system,
          log: @log,
          keep_running: ->() { self.keep_script_running? }
        )
        @script.append_bash_cmd "mkdir -p #{@file_system.testbed_path}/#{@job.execution_directory}"
        @script.append_bash_cmd "cd #{@file_system.testbed_path}/#{@job.execution_directory}"

        @log.info "Setting the execution variables in the environment"
        @script.set_env 'HIVE_RESULTS', @file_system.results_path
        @job.execution_variables.to_h.each_pair do |var, val|
          @script.set_env "HIVE_#{var.to_s}".upcase, val if ! val.kind_of?(Array)
        end
        if @job.execution_variables.retry_urns && !@job.execution_variables.retry_urns.empty?
          @script.set_env "RETRY_URNS", @job.execution_variables.retry_urns
        end
        if @job.execution_variables.tests && @job.execution_variables.tests != [""]
          @script.set_env "TEST_NAMES", @job.execution_variables.tests
        end
        

        @log.info "Appending test script to execution script"
        @script.append_bash_cmd @job.command

        set_job_state_to :running
        @job.start

        @log.info "Pre-execution setup"
        pre_script(@job, @file_system, @script)

        @log.info "Running execution script"
        exit_value = @script.run
        @job.end(exit_value)
      rescue => e
        exception = e
      end

      begin
        @log.info "Post-execution cleanup"
        set_job_state_to :uploading
        post_script(@job, @file_system, @script)

        # Upload results
        @file_system.finalise_results_directory
        upload_results(@job, "#{@file_system.testbed_path}/#{@job.execution_directory}", @file_system.results_path)
      rescue => e
        @log.error( "Post execution failed: " + e.message)
        @log.error("  : #{e.backtrace.join("\n  : ")}")<commit_msg>Set start_time before script execution<commit_after>require 'yaml'

require 'hive'
require 'hive/file_system'
require 'hive/execution_script'
require 'hive/diagnostic_runner'
require 'hive/messages'
require 'hive/port_allocator'
require 'code_cache'
require 'res'

module Hive
  # The generic worker class
  class Worker
    class InvalidJobReservationError < StandardError
    end

    class DeviceNotReady < StandardError
    end

    class NoPortsAvailable < StandardError
    end

    # Device API Object for device associated with this worker
    attr_accessor :device_api, :queues 

    # The main worker process loop
    def initialize(options)
      @options = options
      @parent_pid = @options['parent_pid']
      @device_id = @options['id']
      @hive_id = @options['hive_id']
      @default_component ||= self.class.to_s
      @current_job_start_time = nil
      @hive_mind ||= mind_meld_klass.new(
        url: Chamber.env.network.hive_mind? ? Chamber.env.network.hive_mind : nil,
        pem: Chamber.env.network.cert ? Chamber.env.network.cert : nil,
        ca_file: Chamber.env.network.cafile ? Chamber.env.network.cafile : nil,
        verify_mode: Chamber.env.network.verify_mode ? Chamber.env.network.verify_mode : nil,
        device: hive_mind_device_identifiers
      )
      @device_identity = @options['device_identity'] || 'unknown-device'
      pid = Process.pid
      $PROGRAM_NAME = "#{@options['name_stub'] || 'WORKER'}.#{pid}"
      @log = Hive::Log.new
      @log.add_logger(
        "#{LOG_DIRECTORY}/#{pid}.#{@device_identity}.log",
        Hive.config.logging.worker_level || 'INFO'
      )
      @log.hive_mind = @hive_mind
      @log.default_progname = @default_component

      self.update_queues

      @port_allocator = (@options.has_key?('port_allocator') ? @options['port_allocator'] : Hive::PortAllocator.new(ports: []))
      
      platform = self.class.to_s.scan(/[^:][^:]*/)[2].downcase
      @diagnostic_runner = Hive::DiagnosticRunner.new(@options, Hive.config.diagnostics, platform, @hive_mind) if Hive.config.diagnostics? && Hive.config.diagnostics[platform]

      Hive::Messages.configure do |config|
        config.base_path = Hive.config.network.scheduler
        config.pem_file = Hive.config.network.cert
        config.ssl_verify_mode = OpenSSL::SSL::VERIFY_NONE
      end

      Signal.trap('TERM') do
        @log.info("Worker terminated")
        exit
      end

      @log.info('Starting worker')
      while keep_worker_running?
        begin
          @log.clear({component: @default_component, level: Hive.config.logging.hm_logs_to_delete})
          update_queues
          poll_queue if diagnostics
        rescue DeviceNotReady => e
          @log.warn("#{e.message}\n");
        rescue StandardError => e
          @log.warn("Worker loop aborted: #{e.message}\n  : #{e.backtrace.join("\n  : ")}")
        end
        sleep Hive.config.timings.worker_loop_interval
      end
      @log.info('Exiting worker')
    end

    # Check the queues for work
    def poll_queue
      @job = reserve_job
      if @job.nil?
        @log.info('No job found')
      else
        @log.info('Job starting')
        begin
          @current_job_start_time = Time.now
          execute_job
        rescue => e
          @log.info("Error running test: #{e.message}\n : #{e.backtrace.join("\n :")}")
        end
        cleanup
      end
    end

    # Try to find and reserve a job
    def reserve_job
      @log.info "Trying to reserve job for queues: #{@queues.join(', ')}"
      job = job_message_klass.reserve(@queues, reservation_details)
      @log.debug "Job: #{job.inspect}"
      raise InvalidJobReservationError.new("Invalid Job Reserved") if ! (job.nil? || job.valid?)
      job
    end

    # Get the correct job class
    # This should usually be replaced in the child class
    def job_message_klass
      @log.info 'Generic job class'
      Hive::Messages::Job
    end

    def mind_meld_klass
      MindMeld::Device
    end

    def reservation_details
      @log.debug "Reservations details: hive_id=#{@hive_id}, worker_pid=#{Process.pid}, device_id=#{@hive_mind.id}"
      { hive_id: @hive_id, worker_pid: Process.pid, device_id: @hive_mind.id }
    end

    # Execute a job
    def execute_job
      # Ensure that a killed worker cleans up correctly
      Signal.trap('TERM') do |s|
        Signal.trap('TERM') {} # Prevent retry signals
        @log.info "Caught TERM signal"
        @log.info "Terminating script, if running"
        @script.terminate if @script
        @log.info "Post-execution cleanup"
        signal_safe_post_script(@job, @file_system, @script)

        # Upload results
        @file_system.finalise_results_directory
        upload_files(@job, @file_system.results_path, @file_system.logs_path)
        set_job_state_to :completed
        @job.error('Worker killed')
        @log.info "Worker terminated"
        exit
      end

      @log.info('Job starting')
      @job.prepare(@hive_mind.id)
sleep 20
      exception = nil
      begin
        @log.info "Setting job paths"
        @file_system = Hive::FileSystem.new(@job.job_id, Hive.config.logging.home, @log)
        set_job_state_to :preparing

        if ! @job.repository.to_s.empty?
          @log.info "Checking out the repository"
          @log.debug "  #{@job.repository}"
          @log.debug "  #{@file_system.testbed_path}"
          checkout_code(@job.repository, @file_system.testbed_path, @job.branch)
        end

        @log.info "Initialising execution script"
        @script = Hive::ExecutionScript.new(
          file_system: @file_system,
          log: @log,
          keep_running: ->() { self.keep_script_running? }
        )
        @script.append_bash_cmd "mkdir -p #{@file_system.testbed_path}/#{@job.execution_directory}"
        @script.append_bash_cmd "cd #{@file_system.testbed_path}/#{@job.execution_directory}"

        @log.info "Setting the execution variables in the environment"
        @script.set_env 'HIVE_RESULTS', @file_system.results_path
        @job.execution_variables.to_h.each_pair do |var, val|
          @script.set_env "HIVE_#{var.to_s}".upcase, val if ! val.kind_of?(Array)
        end
        if @job.execution_variables.retry_urns && !@job.execution_variables.retry_urns.empty?
          @script.set_env "RETRY_URNS", @job.execution_variables.retry_urns
        end
        if @job.execution_variables.tests && @job.execution_variables.tests != [""]
          @script.set_env "TEST_NAMES", @job.execution_variables.tests
        end
        

        @log.info "Appending test script to execution script"
        @script.append_bash_cmd @job.command

        set_job_state_to :running

        @log.info "Pre-execution setup"
        pre_script(@job, @file_system, @script)

        @job.start
        @log.info "Running execution script"
        exit_value = @script.run
        @job.end(exit_value)
      rescue => e
        exception = e
      end

      begin
        @log.info "Post-execution cleanup"
        set_job_state_to :uploading
        post_script(@job, @file_system, @script)

        # Upload results
        @file_system.finalise_results_directory
        upload_results(@job, "#{@file_system.testbed_path}/#{@job.execution_directory}", @file_system.results_path)
      rescue => e
        @log.error( "Post execution failed: " + e.message)
        @log.error("  : #{e.backtrace.join("\n  : ")}")
{'fraction_non_alphanumeric': 0.11714690599905526, 'fraction_numerical': 0.000269923746541602, 'mean_word_length': 2.6884021901443504, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 15, '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': '66949', 'n_tokens_mistral': 4742, 'n_tokens_neox': 4659, 'n_words': 940}