Implemented Status Table

src/api.go

@@ -2,12 +2,14 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"net/http"
 	"os"
 	"os/signal"
+	"strings"
 	"sync"
 	"syscall"
 	"time"
@@ -16,34 +18,43 @@ import (
 )
 
 type App struct {
-	redisClient *redis.Client
-	scheduler   *Scheduler
-	supervisor  *Supervisor
-	httpServer  *http.Server
-	wg          sync.WaitGroup
-	log         *slog.Logger
+	redisClient    *redis.Client
+	scheduler      *Scheduler
+	supervisor     *Supervisor
+	httpServer     *http.Server
+	wg             sync.WaitGroup
+	log            *slog.Logger
+	statusRegistry *StatusRegistry
 }
 
 func NewApp(redisAddr, gpuType string, log *slog.Logger) *App {
 	client := redis.NewClient(&redis.Options{Addr: redisAddr})
 	scheduler := NewScheduler(redisAddr, log)
+	statusRegistry := NewStatusRegistry(client, log)
 
 	consumerID := fmt.Sprintf("worker_%d", os.Getpid())
 	supervisor := NewSupervisor(redisAddr, consumerID, gpuType, log)
 
+	// Add dummy supervisors for testing
+	addDummySupervisors(statusRegistry, log)
+
 	mux := http.NewServeMux()
 	a := &App{
-		redisClient: client,
-		scheduler:   scheduler,
-		supervisor:  supervisor,
-		httpServer:  &http.Server{Addr: ":3000", Handler: mux},
-		log:         log,
+		redisClient:    client,
+		scheduler:      scheduler,
+		supervisor:     supervisor,
+		httpServer:     &http.Server{Addr: ":3000", Handler: mux},
+		log:            log,
+		statusRegistry: statusRegistry,
 	}
 
 	mux.HandleFunc("/auth/login", a.login)
 	mux.HandleFunc("/auth/refresh", a.refresh)
 	mux.HandleFunc("/jobs", a.enqueueJob)
 	mux.HandleFunc("/jobs/status", a.getJobStatus)
+	mux.HandleFunc("/supervisors/status", a.getSupervisorStatus)
+	mux.HandleFunc("/supervisors/status/", a.getSupervisorStatusByID)
+	mux.HandleFunc("/supervisors", a.getAllSupervisors)
 
 	a.log.Info("new app initialized", "redis_address", redisAddr,
 		"gpu_type", gpuType, "http_address", a.httpServer.Addr)
@@ -58,6 +69,12 @@ func (a *App) Start() error {
 		return err
 	}
 
+	// Start supervisor
+	if err := a.supervisor.Start(); err != nil {
+		a.log.Error("supervisor start failed", "err", err)
+		return err
+	}
+
 	// Launch HTTP server
 	a.wg.Add(1)
 	go func() {
@@ -164,3 +181,76 @@ func (a *App) getJobStatus(w http.ResponseWriter, r *http.Request) {
 	id := r.URL.Query().Get("id")
 	fmt.Fprintln(w, "job id=", id)
 }
+
+func (a *App) getSupervisorStatus(w http.ResponseWriter, r *http.Request) {
+	supervisors, err := a.statusRegistry.GetAllSupervisors()
+	if err != nil {
+		a.log.Error("failed to get supervisor status", "error", err)
+		http.Error(w, "failed to get supervisor status", http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(map[string]interface{}{
+		"supervisors": supervisors,
+		"count":       len(supervisors),
+	}); err != nil {
+		a.log.Error("failed to encode supervisor status response", "error", err)
+		http.Error(w, "failed to encode response", http.StatusInternalServerError)
+		return
+	}
+}
+
+func (a *App) getSupervisorStatusByID(w http.ResponseWriter, r *http.Request) {
+	// extract consumer ID from URL path
+	path := strings.TrimPrefix(r.URL.Path, "/supervisors/status/")
+	if path == "" {
+		http.Error(w, "consumer ID required", http.StatusBadRequest)
+		return
+	}
+
+	supervisor, err := a.statusRegistry.GetSupervisor(path)
+	if err != nil {
+		a.log.Error("failed to get supervisor status", "consumer_id", path, "error", err)
+		http.Error(w, "supervisor not found", http.StatusNotFound)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(supervisor); err != nil {
+		a.log.Error("failed to encode supervisor status response", "error", err)
+		http.Error(w, "failed to encode response", http.StatusInternalServerError)
+		return
+	}
+}
+
+func (a *App) getAllSupervisors(w http.ResponseWriter, r *http.Request) {
+	// check if we want only active supervisors
+	activeOnly := r.URL.Query().Get("active") == "true"
+
+	var supervisors []SupervisorStatus
+	var err error
+
+	if activeOnly {
+		supervisors, err = a.statusRegistry.GetActiveSupervisors()
+	} else {
+		supervisors, err = a.statusRegistry.GetAllSupervisors()
+	}
+
+	if err != nil {
+		a.log.Error("failed to get supervisors", "active_only", activeOnly, "error", err)
+		http.Error(w, "failed to get supervisors", http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(map[string]interface{}{
+		"supervisors": supervisors,
+		"count":       len(supervisors),
+		"active_only": activeOnly,
+	}); err != nil {
+		a.log.Error("failed to encode supervisors response", "error", err)
+		http.Error(w, "failed to encode response", http.StatusInternalServerError)
+		return
+	}
+}

src/status.go

@@ -0,0 +1,128 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+type StatusRegistry struct {
+	redisClient *redis.Client
+	log         *slog.Logger
+}
+
+func NewStatusRegistry(redisClient *redis.Client, log *slog.Logger) *StatusRegistry {
+	return &StatusRegistry{
+		redisClient: redisClient,
+		log:         log,
+	}
+}
+
+func addDummySupervisors(statusRegistry *StatusRegistry, log *slog.Logger) {
+	now := time.Now()
+
+	// three dummy supervisors with different statuses
+	dummySupervisors := []SupervisorStatus{
+		{
+			ConsumerID: "worker_amd_001",
+			GPUType:    "AMD",
+			Status:     "active",
+			LastSeen:   now,                     // now
+			StartedAt:  now.Add(-2 * time.Hour), // 2hours ago
+		},
+		{
+			ConsumerID: "worker_nvidia_002",
+			GPUType:    "NVIDIA",
+			Status:     "active",
+			LastSeen:   now.Add(-30 * time.Second), // 30 seconds ago
+			StartedAt:  now.Add(-1 * time.Hour),    // 1 hour ago
+		},
+		{
+			ConsumerID: "worker_tt_003",
+			GPUType:    "TT",
+			Status:     "inactive",
+			LastSeen:   now.Add(-5 * time.Minute), // seen 5 minutes ago
+			StartedAt:  now.Add(-3 * time.Hour),   // 3 hours ago
+		},
+	}
+
+	// Add each dummy supervisor to the registry
+	for _, supervisor := range dummySupervisors {
+		if err := statusRegistry.UpdateStatus(supervisor.ConsumerID, supervisor); err != nil {
+			log.Error("failed to add dummy supervisor", "consumer_id", supervisor.ConsumerID, "error", err)
+		} else {
+			log.Info("added dummy supervisor", "consumer_id", supervisor.ConsumerID, "gpu_type", supervisor.GPUType, "status", supervisor.Status)
+		}
+	}
+}
+
+func (sr *StatusRegistry) GetAllSupervisors() ([]SupervisorStatus, error) {
+	ctx := context.Background()
+	result := sr.redisClient.HGetAll(ctx, SupervisorStatusKey)
+	if result.Err() != nil {
+		return nil, fmt.Errorf("failed to get supervisor status: %w", result.Err())
+	}
+
+	var supervisors []SupervisorStatus
+	for consumerID, statusJSON := range result.Val() {
+		var status SupervisorStatus
+		if err := json.Unmarshal([]byte(statusJSON), &status); err != nil {
+			sr.log.Error("failed to unmarshal supervisor status", "consumer_id", consumerID, "error", err)
+			continue
+		}
+		supervisors = append(supervisors, status)
+	}
+
+	return supervisors, nil
+}
+
+func (sr *StatusRegistry) GetSupervisor(consumerID string) (*SupervisorStatus, error) {
+	ctx := context.Background()
+	result := sr.redisClient.HGet(ctx, SupervisorStatusKey, consumerID)
+	if result.Err() != nil {
+		return nil, fmt.Errorf("failed to get supervisor status: %w", result.Err())
+	}
+
+	var status SupervisorStatus
+	if err := json.Unmarshal([]byte(result.Val()), &status); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal supervisor status: %w", err)
+	}
+
+	return &status, nil
+}
+
+func (sr *StatusRegistry) GetActiveSupervisors() ([]SupervisorStatus, error) {
+	allSupervisors, err := sr.GetAllSupervisors()
+	if err != nil {
+		return nil, err
+	}
+
+	var activeSupervisors []SupervisorStatus
+	for _, supervisor := range allSupervisors {
+		if supervisor.Status == "active" {
+			activeSupervisors = append(activeSupervisors, supervisor)
+		}
+	}
+
+	return activeSupervisors, nil
+}
+
+func (sr *StatusRegistry) UpdateStatus(consumerID string, status SupervisorStatus) error {
+	ctx := context.Background()
+	statusJSON, err := json.Marshal(status)
+	if err != nil {
+		return fmt.Errorf("failed to marshal supervisor status: %w", err)
+	}
+
+	result := sr.redisClient.HSet(ctx, SupervisorStatusKey, consumerID, string(statusJSON))
+	if result.Err() != nil {
+		return fmt.Errorf("failed to update supervisor status: %w", result.Err())
+	}
+
+	sr.log.Info("supervisor status updated", "consumer_id", consumerID, "status", status.Status)
+	return nil
+}

src/util.go

@@ -22,6 +22,14 @@ type Job struct {
 	RequiredGPU string                 `json:"gpu"`
 }
 
+type SupervisorStatus struct {
+	ConsumerID string    `json:"consumer_id"`
+	GPUType    string    `json:"gpu_type"`
+	Status     string    `json:"status"` // "active", "inactive", "failed"
+	LastSeen   time.Time `json:"last_seen"`
+	StartedAt  time.Time `json:"started_at"`
+}
+
 func generateJobID() string {
 	return fmt.Sprintf("job_%d_%d", time.Now().UnixNano(), os.Getpid())
 }