package services import ( "context" "fmt" "github.com/jfraeys/fetch_ml/cmd/tui/internal/config" "github.com/jfraeys/fetch_ml/cmd/tui/internal/model" "github.com/jfraeys/fetch_ml/internal/experiment" "github.com/jfraeys/fetch_ml/internal/network" "github.com/jfraeys/fetch_ml/internal/queue" ) // TaskQueue wraps the internal queue.TaskQueue for TUI compatibility type TaskQueue struct { internal *queue.TaskQueue expManager *experiment.Manager ctx context.Context } func NewTaskQueue(cfg *config.Config) (*TaskQueue, error) { // Create internal queue config queueCfg := queue.Config{ RedisAddr: cfg.RedisAddr, RedisPassword: cfg.RedisPassword, RedisDB: cfg.RedisDB, } internalQueue, err := queue.NewTaskQueue(queueCfg) if err != nil { return nil, fmt.Errorf("failed to create task queue: %w", err) } // Initialize experiment manager // TODO: Get base path from config expManager := experiment.NewManager("./experiments") return &TaskQueue{ internal: internalQueue, expManager: expManager, ctx: context.Background(), }, nil } func (tq *TaskQueue) EnqueueTask(jobName, args string, priority int64) (*model.Task, error) { // Create internal task internalTask := &queue.Task{ JobName: jobName, Args: args, Priority: priority, } // Use internal queue to enqueue err := tq.internal.AddTask(internalTask) if err != nil { return nil, err } // Convert to TUI model return &model.Task{ ID: internalTask.ID, JobName: internalTask.JobName, Args: internalTask.Args, Status: "queued", Priority: int64(internalTask.Priority), CreatedAt: internalTask.CreatedAt, Metadata: internalTask.Metadata, }, nil } func (tq *TaskQueue) GetNextTask() (*model.Task, error) { internalTask, err := tq.internal.GetNextTask() if err != nil { return nil, err } if internalTask == nil { return nil, nil } // Convert to TUI model return &model.Task{ ID: internalTask.ID, JobName: internalTask.JobName, Args: internalTask.Args, Status: internalTask.Status, Priority: internalTask.Priority, CreatedAt: internalTask.CreatedAt, Metadata: internalTask.Metadata, }, nil } func (tq *TaskQueue) GetTask(taskID string) (*model.Task, error) { internalTask, err := tq.internal.GetTask(taskID) if err != nil { return nil, err } // Convert to TUI model return &model.Task{ ID: internalTask.ID, JobName: internalTask.JobName, Args: internalTask.Args, Status: internalTask.Status, Priority: internalTask.Priority, CreatedAt: internalTask.CreatedAt, Metadata: internalTask.Metadata, }, nil } func (tq *TaskQueue) UpdateTask(task *model.Task) error { // Convert to internal task internalTask := &queue.Task{ ID: task.ID, JobName: task.JobName, Args: task.Args, Status: task.Status, Priority: task.Priority, CreatedAt: task.CreatedAt, Metadata: task.Metadata, } return tq.internal.UpdateTask(internalTask) } func (tq *TaskQueue) GetQueuedTasks() ([]*model.Task, error) { internalTasks, err := tq.internal.GetAllTasks() if err != nil { return nil, err } // Convert to TUI models tasks := make([]*model.Task, len(internalTasks)) for i, task := range internalTasks { tasks[i] = &model.Task{ ID: task.ID, JobName: task.JobName, Args: task.Args, Status: task.Status, Priority: task.Priority, CreatedAt: task.CreatedAt, Metadata: task.Metadata, } } return tasks, nil } func (tq *TaskQueue) GetJobStatus(jobName string) (map[string]string, error) { // This method doesn't exist in internal queue, implement basic version task, err := tq.internal.GetTaskByName(jobName) if err != nil { return nil, err } if task == nil { return map[string]string{"status": "not_found"}, nil } return map[string]string{ "status": task.Status, "task_id": task.ID, }, nil } func (tq *TaskQueue) RecordMetric(jobName, metric string, value float64) error { return tq.internal.RecordMetric(jobName, metric, value) } func (tq *TaskQueue) GetMetrics(jobName string) (map[string]string, error) { // This method doesn't exist in internal queue, return empty for now return map[string]string{}, nil } func (tq *TaskQueue) ListDatasets() ([]model.DatasetInfo, error) { // This method doesn't exist in internal queue, return empty for now return []model.DatasetInfo{}, nil } func (tq *TaskQueue) CancelTask(taskID string) error { return tq.internal.CancelTask(taskID) } func (tq *TaskQueue) ListExperiments() ([]string, error) { return tq.expManager.ListExperiments() } func (tq *TaskQueue) GetExperimentDetails(commitID string) (string, error) { meta, err := tq.expManager.ReadMetadata(commitID) if err != nil { return "", err } metrics, err := tq.expManager.GetMetrics(commitID) if err != nil { return "", err } output := fmt.Sprintf("Experiment: %s\n", meta.JobName) output += fmt.Sprintf("Commit ID: %s\n", meta.CommitID) output += fmt.Sprintf("User: %s\n", meta.User) output += fmt.Sprintf("Timestamp: %d\n\n", meta.Timestamp) output += "Metrics:\n" if len(metrics) == 0 { output += " No metrics logged.\n" } else { for _, m := range metrics { output += fmt.Sprintf(" %s: %.4f (Step: %d)\n", m.Name, m.Value, m.Step) } } return output, nil } func (tq *TaskQueue) Close() error { return tq.internal.Close() } // MLServer wraps network.SSHClient for backward compatibility type MLServer struct { *network.SSHClient addr string } func NewMLServer(cfg *config.Config) (*MLServer, error) { // Local mode: skip SSH entirely if cfg.Host == "" { client, _ := network.NewSSHClient("", "", "", 0, "") return &MLServer{SSHClient: client, addr: "localhost"}, nil } client, err := network.NewSSHClient(cfg.Host, cfg.User, cfg.SSHKey, cfg.Port, cfg.KnownHosts) if err != nil { return nil, err } addr := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port) return &MLServer{SSHClient: client, addr: addr}, nil }