Implement job metadata extraction and task management features. Add validation for frame range limits, enhance job and task data structures, and introduce new API endpoints for metadata and task retrieval. Update client-side components to handle metadata extraction and display task statuses. Improve error handling in API responses.

2025-11-22 06:37:32 -06:00
parent 27a09aedd6
commit c9ade39ad9
10 changed files with 1078 additions and 88 deletions
--- a/internal/runner/client.go
+++ b/internal/runner/client.go
@@ -27,28 +27,31 @@ import (

 // Client represents a runner client
 type Client struct {
-	managerURL    string
-	name          string
-	hostname      string
-	ipAddress     string
-	httpClient    *http.Client
-	runnerID      int64
-	runnerSecret  string
-	managerSecret string
-	wsConn        *websocket.Conn
-	wsConnMu      sync.RWMutex
-	stopChan      chan struct{}
+	managerURL     string
+	name           string
+	hostname       string
+	ipAddress      string
+	httpClient     *http.Client
+	runnerID       int64
+	runnerSecret   string
+	managerSecret  string
+	wsConn         *websocket.Conn
+	wsConnMu       sync.RWMutex
+	stopChan       chan struct{}
+	stepStartTimes map[string]time.Time // key: "taskID:stepName"
+	stepTimesMu    sync.RWMutex
 }

 // NewClient creates a new runner client
 func NewClient(managerURL, name, hostname, ipAddress string) *Client {
 	return &Client{
-		managerURL: managerURL,
-		name:       name,
-		hostname:   hostname,
-		ipAddress:  ipAddress,
-		httpClient: &http.Client{Timeout: 30 * time.Second},
-		stopChan:   make(chan struct{}),
+		managerURL:     managerURL,
+		name:           name,
+		hostname:       hostname,
+		ipAddress:      ipAddress,
+		httpClient:     &http.Client{Timeout: 30 * time.Second},
+		stopChan:       make(chan struct{}),
+		stepStartTimes: make(map[string]time.Time),
 	}
 }

@@ -259,14 +262,9 @@ func (c *Client) handleTaskAssignment(msg map[string]interface{}) {
 	outputFormat, _ := data["output_format"].(string)
 	frameStart, _ := data["frame_start"].(float64)
 	frameEnd, _ := data["frame_end"].(float64)
+	taskType, _ := data["task_type"].(string)
 	inputFilesRaw, _ := data["input_files"].([]interface{})

-	if len(inputFilesRaw) == 0 {
-		log.Printf("No input files for task %v", taskID)
-		c.sendTaskComplete(int64(taskID), "", false, "No input files")
-		return
-	}
-
 	// Convert to task map format
 	taskMap := map[string]interface{}{
 		"id":          taskID,
@@ -275,9 +273,27 @@ func (c *Client) handleTaskAssignment(msg map[string]interface{}) {
 		"frame_end":   frameEnd,
 	}

-	// Process the task
+	// Process the task based on type
 	go func() {
-		if err := c.processTask(taskMap, jobName, outputFormat, inputFilesRaw); err != nil {
+		var err error
+		if taskType == "metadata" {
+			if len(inputFilesRaw) == 0 {
+				log.Printf("No input files for metadata task %v", taskID)
+				c.sendTaskComplete(int64(taskID), "", false, "No input files")
+				return
+			}
+			err = c.processMetadataTask(taskMap, int64(jobID), inputFilesRaw)
+		} else if taskType == "video_generation" {
+			err = c.processVideoGenerationTask(taskMap, int64(jobID))
+		} else {
+			if len(inputFilesRaw) == 0 {
+				log.Printf("No input files for task %v", taskID)
+				c.sendTaskComplete(int64(taskID), "", false, "No input files")
+				return
+			}
+			err = c.processTask(taskMap, jobName, outputFormat, inputFilesRaw)
+		}
+		if err != nil {
 			log.Printf("Failed to process task %v: %v", taskID, err)
 			c.sendTaskComplete(int64(taskID), "", false, err.Error())
 		}
@@ -334,7 +350,78 @@ func (c *Client) sendLog(taskID int64, logLevel types.LogLevel, message, stepNam

 // sendStepUpdate sends a step start/complete event to the manager
 func (c *Client) sendStepUpdate(taskID int64, stepName string, status types.StepStatus, errorMsg string) {
-	// This would ideally be a separate endpoint, but for now we'll use logs
+	key := fmt.Sprintf("%d:%s", taskID, stepName)
+	var durationMs *int
+
+	// Track step start time
+	if status == types.StepStatusRunning {
+		c.stepTimesMu.Lock()
+		c.stepStartTimes[key] = time.Now()
+		c.stepTimesMu.Unlock()
+	}
+
+	// Calculate duration if step is completing
+	if status == types.StepStatusCompleted || status == types.StepStatusFailed {
+		c.stepTimesMu.RLock()
+		startTime, exists := c.stepStartTimes[key]
+		c.stepTimesMu.RUnlock()
+		if exists {
+			duration := int(time.Since(startTime).Milliseconds())
+			durationMs = &duration
+			c.stepTimesMu.Lock()
+			delete(c.stepStartTimes, key)
+			c.stepTimesMu.Unlock()
+		}
+	}
+
+	// Send step update via HTTP API
+	reqBody := map[string]interface{}{
+		"step_name": stepName,
+		"status":    string(status),
+	}
+	if durationMs != nil {
+		reqBody["duration_ms"] = *durationMs
+	}
+	if errorMsg != "" {
+		reqBody["error_message"] = errorMsg
+	}
+
+	body, _ := json.Marshal(reqBody)
+	path := fmt.Sprintf("/api/runner/tasks/%d/steps?runner_id=%d", taskID, c.runnerID)
+	resp, err := c.doSignedRequest("POST", path, body)
+	if err != nil {
+		log.Printf("Failed to send step update: %v", err)
+		// Fallback to log-based tracking
+		msg := fmt.Sprintf("Step %s: %s", stepName, status)
+		if errorMsg != "" {
+			msg += " - " + errorMsg
+		}
+		logLevel := types.LogLevelInfo
+		if status == types.StepStatusFailed {
+			logLevel = types.LogLevelError
+		}
+		c.sendLog(taskID, logLevel, msg, stepName)
+		return
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		log.Printf("Step update failed: %s", string(body))
+		// Fallback to log-based tracking
+		msg := fmt.Sprintf("Step %s: %s", stepName, status)
+		if errorMsg != "" {
+			msg += " - " + errorMsg
+		}
+		logLevel := types.LogLevelInfo
+		if status == types.StepStatusFailed {
+			logLevel = types.LogLevelError
+		}
+		c.sendLog(taskID, logLevel, msg, stepName)
+		return
+	}
+
+	// Also send log for debugging
 	msg := fmt.Sprintf("Step %s: %s", stepName, status)
 	if errorMsg != "" {
 		msg += " - " + errorMsg
@@ -455,33 +542,20 @@ func (c *Client) processTask(task map[string]interface{}, jobName, outputFormat
 	}
 	c.sendStepUpdate(taskID, "complete", types.StepStatusCompleted, "")

-	// For MP4 format, check if all frames are done and generate video
-	if outputFormat == "MP4" {
-		if err := c.checkAndGenerateMP4(jobID); err != nil {
-			log.Printf("Failed to generate MP4 for job %d: %v", jobID, err)
-			// Don't fail the task if video generation fails - frames are already uploaded
-		}
-	}
-
 	return nil
 }

-// checkAndGenerateMP4 checks if all frames are complete and generates MP4 if so
-func (c *Client) checkAndGenerateMP4(jobID int64) error {
-	// Check job status
-	job, err := c.getJobStatus(jobID)
-	if err != nil {
-		return fmt.Errorf("failed to get job status: %w", err)
-	}
+// processVideoGenerationTask processes a video generation task
+func (c *Client) processVideoGenerationTask(task map[string]interface{}, jobID int64) error {
+	taskID := int64(task["id"].(float64))

-	if job["status"] != "completed" {
-		log.Printf("Job %d not yet complete (%v), skipping MP4 generation", jobID, job["status"])
-		return nil
-	}
+	c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Starting video generation task: job %d", jobID), "")
+	log.Printf("Processing video generation task %d for job %d", taskID, jobID)

 	// Get all output files for this job
 	files, err := c.getJobFiles(jobID)
 	if err != nil {
+		c.sendStepUpdate(taskID, "get_files", types.StepStatusFailed, err.Error())
 		return fmt.Errorf("failed to get job files: %w", err)
 	}

@@ -496,14 +570,24 @@ func (c *Client) checkAndGenerateMP4(jobID int64) error {
 	}

 	if len(pngFiles) == 0 {
-		return fmt.Errorf("no PNG frame files found for MP4 generation")
+		err := fmt.Errorf("no PNG frame files found for MP4 generation")
+		c.sendStepUpdate(taskID, "get_files", types.StepStatusFailed, err.Error())
+		return err
 	}

+	c.sendStepUpdate(taskID, "get_files", types.StepStatusCompleted, "")
+	c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Found %d PNG frames for video generation", len(pngFiles)), "get_files")
+
 	log.Printf("Generating MP4 for job %d from %d PNG frames", jobID, len(pngFiles))

+	// Step: download_frames
+	c.sendStepUpdate(taskID, "download_frames", types.StepStatusRunning, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Downloading PNG frames...", "download_frames")
+
 	// Create work directory for video generation
 	workDir := filepath.Join(os.TempDir(), fmt.Sprintf("fuego-video-%d", jobID))
 	if err := os.MkdirAll(workDir, 0755); err != nil {
+		c.sendStepUpdate(taskID, "download_frames", types.StepStatusFailed, err.Error())
 		return fmt.Errorf("failed to create work directory: %w", err)
 	}
 	defer os.RemoveAll(workDir)
@@ -521,11 +605,19 @@ func (c *Client) checkAndGenerateMP4(jobID int64) error {
 	}

 	if len(frameFiles) == 0 {
-		return fmt.Errorf("failed to download any frame files")
+		err := fmt.Errorf("failed to download any frame files")
+		c.sendStepUpdate(taskID, "download_frames", types.StepStatusFailed, err.Error())
+		return err
 	}

 	// Sort frame files by name to ensure correct order
 	sort.Strings(frameFiles)
+	c.sendStepUpdate(taskID, "download_frames", types.StepStatusCompleted, "")
+	c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Downloaded %d frames", len(frameFiles)), "download_frames")
+
+	// Step: generate_video
+	c.sendStepUpdate(taskID, "generate_video", types.StepStatusRunning, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Generating MP4 video with ffmpeg...", "generate_video")

 	// Generate MP4 using ffmpeg
 	outputMP4 := filepath.Join(workDir, fmt.Sprintf("output_%d.mp4", jobID))
@@ -546,20 +638,42 @@ func (c *Client) checkAndGenerateMP4(jobID int64) error {
 	if err != nil {
 		// Try alternative method with concat demuxer
 		log.Printf("First ffmpeg attempt failed, trying concat method: %s", string(output))
-		return c.generateMP4WithConcat(frameFiles, outputMP4, workDir)
+		err = c.generateMP4WithConcat(frameFiles, outputMP4, workDir)
+		if err != nil {
+			c.sendStepUpdate(taskID, "generate_video", types.StepStatusFailed, err.Error())
+			return err
+		}
 	}

 	// Check if MP4 was created
 	if _, err := os.Stat(outputMP4); os.IsNotExist(err) {
-		return fmt.Errorf("MP4 file not created: %s", outputMP4)
+		err := fmt.Errorf("MP4 file not created: %s", outputMP4)
+		c.sendStepUpdate(taskID, "generate_video", types.StepStatusFailed, err.Error())
+		return err
 	}

+	c.sendStepUpdate(taskID, "generate_video", types.StepStatusCompleted, "")
+	c.sendLog(taskID, types.LogLevelInfo, "MP4 video generated successfully", "generate_video")
+
+	// Step: upload_video
+	c.sendStepUpdate(taskID, "upload_video", types.StepStatusRunning, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Uploading MP4 video...", "upload_video")
+
 	// Upload MP4 file
 	mp4Path, err := c.uploadFile(jobID, outputMP4)
 	if err != nil {
+		c.sendStepUpdate(taskID, "upload_video", types.StepStatusFailed, err.Error())
 		return fmt.Errorf("failed to upload MP4: %w", err)
 	}

+	c.sendStepUpdate(taskID, "upload_video", types.StepStatusCompleted, "")
+	c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Successfully uploaded MP4: %s", mp4Path), "upload_video")
+
+	// Mark task as complete
+	if err := c.completeTask(taskID, mp4Path, true, ""); err != nil {
+		return err
+	}
+
 	log.Printf("Successfully generated and uploaded MP4 for job %d: %s", jobID, mp4Path)
 	return nil
 }
@@ -785,6 +899,197 @@ func (c *Client) uploadFile(jobID int64, filePath string) (string, error) {
 	return result.FilePath, nil
 }

+// processMetadataTask processes a metadata extraction task
+func (c *Client) processMetadataTask(task map[string]interface{}, jobID int64, inputFiles []interface{}) error {
+	taskID := int64(task["id"].(float64))
+
+	c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Starting metadata extraction task: job %d", jobID), "")
+	log.Printf("Processing metadata extraction task %d for job %d", taskID, jobID)
+
+	// Create work directory
+	workDir := filepath.Join(os.TempDir(), fmt.Sprintf("fuego-metadata-%d", taskID))
+	if err := os.MkdirAll(workDir, 0755); err != nil {
+		return fmt.Errorf("failed to create work directory: %w", err)
+	}
+	defer os.RemoveAll(workDir)
+
+	// Step: download
+	c.sendStepUpdate(taskID, "download", types.StepStatusRunning, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Downloading blend file...", "download")
+	blendFile := ""
+	for _, filePath := range inputFiles {
+		filePathStr := filePath.(string)
+		if err := c.downloadFile(filePathStr, workDir); err != nil {
+			c.sendStepUpdate(taskID, "download", types.StepStatusFailed, err.Error())
+			return fmt.Errorf("failed to download file %s: %w", filePathStr, err)
+		}
+		if filepath.Ext(filePathStr) == ".blend" {
+			blendFile = filepath.Join(workDir, filepath.Base(filePathStr))
+		}
+	}
+
+	if blendFile == "" {
+		err := fmt.Errorf("no .blend file found in input files")
+		c.sendStepUpdate(taskID, "download", types.StepStatusFailed, err.Error())
+		return err
+	}
+	c.sendStepUpdate(taskID, "download", types.StepStatusCompleted, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Blend file downloaded successfully", "download")
+
+	// Step: extract_metadata
+	c.sendStepUpdate(taskID, "extract_metadata", types.StepStatusRunning, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Extracting metadata from blend file...", "extract_metadata")
+
+	// Create Python script to extract metadata
+	scriptPath := filepath.Join(workDir, "extract_metadata.py")
+	scriptContent := `import bpy
+import json
+import sys
+
+# Get scene
+scene = bpy.context.scene
+
+# Extract frame range
+frame_start = scene.frame_start
+frame_end = scene.frame_end
+
+# Extract render settings
+render = scene.render
+resolution_x = render.resolution_x
+resolution_y = render.resolution_y
+samples = scene.cycles.samples if scene.cycles else scene.eevee.taa_render_samples
+engine = scene.render.engine.lower()
+
+# Determine output format from file format
+output_format = render.image_settings.file_format
+
+# Extract scene info
+camera_count = len([obj for obj in scene.objects if obj.type == 'CAMERA'])
+object_count = len(scene.objects)
+material_count = len(bpy.data.materials)
+
+# Build metadata dictionary
+metadata = {
+    "frame_start": frame_start,
+    "frame_end": frame_end,
+    "render_settings": {
+        "resolution_x": resolution_x,
+        "resolution_y": resolution_y,
+        "samples": samples,
+        "output_format": output_format,
+        "engine": engine
+    },
+    "scene_info": {
+        "camera_count": camera_count,
+        "object_count": object_count,
+        "material_count": material_count
+    }
+}
+
+# Output as JSON
+print(json.dumps(metadata))
+sys.stdout.flush()
+`
+
+	if err := os.WriteFile(scriptPath, []byte(scriptContent), 0644); err != nil {
+		c.sendStepUpdate(taskID, "extract_metadata", types.StepStatusFailed, err.Error())
+		return fmt.Errorf("failed to create extraction script: %w", err)
+	}
+
+	// Execute Blender with Python script
+	cmd := exec.Command("blender", "-b", blendFile, "--python", scriptPath)
+	cmd.Dir = workDir
+	output, err := cmd.CombinedOutput()
+	if err != nil {
+		errMsg := fmt.Sprintf("blender metadata extraction failed: %w\nOutput: %s", err, string(output))
+		c.sendLog(taskID, types.LogLevelError, errMsg, "extract_metadata")
+		c.sendStepUpdate(taskID, "extract_metadata", types.StepStatusFailed, errMsg)
+		return fmt.Errorf(errMsg)
+	}
+
+	// Parse output (metadata is printed to stdout)
+	metadataJSON := strings.TrimSpace(string(output))
+	// Extract JSON from output (Blender may print other stuff)
+	jsonStart := strings.Index(metadataJSON, "{")
+	jsonEnd := strings.LastIndex(metadataJSON, "}")
+	if jsonStart == -1 || jsonEnd == -1 || jsonEnd <= jsonStart {
+		errMsg := "Failed to extract JSON from Blender output"
+		c.sendLog(taskID, types.LogLevelError, errMsg, "extract_metadata")
+		c.sendStepUpdate(taskID, "extract_metadata", types.StepStatusFailed, errMsg)
+		return fmt.Errorf(errMsg)
+	}
+	metadataJSON = metadataJSON[jsonStart : jsonEnd+1]
+
+	var metadata types.BlendMetadata
+	if err := json.Unmarshal([]byte(metadataJSON), &metadata); err != nil {
+		errMsg := fmt.Sprintf("Failed to parse metadata JSON: %w", err)
+		c.sendLog(taskID, types.LogLevelError, errMsg, "extract_metadata")
+		c.sendStepUpdate(taskID, "extract_metadata", types.StepStatusFailed, errMsg)
+		return fmt.Errorf(errMsg)
+	}
+
+	c.sendLog(taskID, types.LogLevelInfo, fmt.Sprintf("Metadata extracted: frames %d-%d, resolution %dx%d",
+		metadata.FrameStart, metadata.FrameEnd, metadata.RenderSettings.ResolutionX, metadata.RenderSettings.ResolutionY), "extract_metadata")
+	c.sendStepUpdate(taskID, "extract_metadata", types.StepStatusCompleted, "")
+
+	// Step: submit_metadata
+	c.sendStepUpdate(taskID, "submit_metadata", types.StepStatusRunning, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Submitting metadata to manager...", "submit_metadata")
+
+	// Submit metadata to manager
+	if err := c.submitMetadata(jobID, metadata); err != nil {
+		errMsg := fmt.Sprintf("Failed to submit metadata: %w", err)
+		c.sendLog(taskID, types.LogLevelError, errMsg, "submit_metadata")
+		c.sendStepUpdate(taskID, "submit_metadata", types.StepStatusFailed, errMsg)
+		return fmt.Errorf(errMsg)
+	}
+
+	c.sendStepUpdate(taskID, "submit_metadata", types.StepStatusCompleted, "")
+	c.sendLog(taskID, types.LogLevelInfo, "Metadata extraction completed successfully", "")
+
+	// Mark task as complete
+	c.sendTaskComplete(taskID, "", true, "")
+	return nil
+}
+
+// submitMetadata submits extracted metadata to the manager
+func (c *Client) submitMetadata(jobID int64, metadata types.BlendMetadata) error {
+	metadataJSON, err := json.Marshal(metadata)
+	if err != nil {
+		return fmt.Errorf("failed to marshal metadata: %w", err)
+	}
+
+	path := fmt.Sprintf("/api/runner/jobs/%d/metadata?runner_id=%d", jobID, c.runnerID)
+	timestamp := time.Now()
+	message := fmt.Sprintf("POST\n%s\n%s\n%d", path, string(metadataJSON), timestamp.Unix())
+	h := hmac.New(sha256.New, []byte(c.runnerSecret))
+	h.Write([]byte(message))
+	signature := hex.EncodeToString(h.Sum(nil))
+
+	url := fmt.Sprintf("%s%s", c.managerURL, path)
+	req, err := http.NewRequest("POST", url, bytes.NewReader(metadataJSON))
+	if err != nil {
+		return fmt.Errorf("failed to create request: %w", err)
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("X-Runner-Signature", signature)
+	req.Header.Set("X-Runner-Timestamp", fmt.Sprintf("%d", timestamp.Unix()))
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("failed to submit metadata: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("metadata submission failed: %s", string(body))
+	}
+
+	return nil
+}
+
 // completeTask marks a task as complete via WebSocket (or HTTP fallback)
 func (c *Client) completeTask(taskID int64, outputPath string, success bool, errorMsg string) error {
 	return c.sendTaskComplete(taskID, outputPath, success, errorMsg)