fyp/logic/tasks/runner/runner.go

package task_runner

import (
	"database/sql"
	"fmt"
	"os"
	"time"

	"github.com/charmbracelet/log"

	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/utils"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/tasks/utils"
    . "git.andr3h3nriqu3s.com/andr3/fyp/logic/models"
)

/**
* Actually runs the code
 */
func runner(db *sql.DB, task_channel chan Task, index int, back_channel chan int) {
	logger := log.NewWithOptions(os.Stdout, log.Options{
		ReportCaller:    true,
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          fmt.Sprintf("Runner %d", index),
	})

	defer func() {
		if r := recover(); r != nil {
			logger.Error("Recovered in file processor", "processor id", index, "due to", r)
			back_channel <- -index
		}
	}()

	logger.Info("Started up")

    var err error

    base := BasePackStruct{
        Db: db, 
        Logger: logger,
    }

	for task := range task_channel {
		logger.Info("Got task", "task", task)
        
        if task.TaskType == int(TASK_TYPE_CLASSIFICATION) {
            logger.Info("Classification Task")
            if err = ClassifyTask(base, task); err != nil {
                logger.Error("Classification task failed", "error", "err")
            }

            back_channel <- index
            continue
        }


        logger.Error("Do not know how to route task", "task", task)
        back_channel <- index
	}
}

/**
* Tells the orcchestator to look at the task list from time to time
 */
func attentionSeeker(config Config, back_channel chan int) {
	logger := log.NewWithOptions(os.Stdout, log.Options{
		ReportCaller:    true,
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          "Runner Orchestrator Logger [Attention]",
	})

	logger.Info("Started up")

    t, err := time.ParseDuration(config.GpuWorker.Pulling)
    if err != nil {
        logger.Error("Failed to load", "error", err)
        return
    }

	for true {
		back_channel <- 0
        
        time.Sleep(t)
	}
}

/**
* Manages what worker should to Work
 */
func RunnerOrchestrator(db *sql.DB, config Config) {
	logger := log.NewWithOptions(os.Stdout, log.Options{
		ReportCaller:    true,
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          "Runner Orchestrator Logger",
	})

	gpu_workers := config.GpuWorker.NumberOfWorkers

	logger.Info("Starting runners")

	task_runners := make([]chan Task, gpu_workers)
	task_runners_used := make([]bool, gpu_workers)
	// One more to accomudate the Attention Seeker channel
	back_channel := make(chan int, gpu_workers+1)

	go attentionSeeker(config, back_channel)

	// Start the runners
	for i := 0; i < gpu_workers; i++ {
		task_runners[i] = make(chan Task, 10)
		task_runners_used[i] = false
		go runner(db, task_runners[i], i+1, back_channel)
	}

    var task_to_dispatch *Task = nil

	for i := range back_channel {

		if i > 0 {
			logger.Info("Runner freed", "runner", i)
			task_runners_used[i-1] = false
		} else if i < 0 {
			logger.Error("Runner died! Restarting!", "runner", i)
			task_runners_used[i-1] = false
			go runner(db, task_runners[i-1], i, back_channel)
		}

        if task_to_dispatch == nil {
            var task Task
            err := GetDBOnce(db, &task, "tasks where status=$1 limit 1", TASK_TODO)
            if err != NotFoundError && err != nil{
                log.Error("Failed to get tasks from db")
                continue
            }
            if err == NotFoundError {
                task_to_dispatch = nil
            } else {
                task_to_dispatch = &task
            }
        }

        if task_to_dispatch != nil {
            for i := 0; i < len(task_runners_used); i += 1 {
                if !task_runners_used[i] {
                    task_runners[i] <- *task_to_dispatch
                    task_runners_used[i] = true
                    task_to_dispatch = nil
                    break
                }
            }
        }

	}
}

func StartRunners(db *sql.DB, config Config) {
	go RunnerOrchestrator(db, config)
}