jiggablend/internal/runner/tasks/processor.go

// Package tasks provides task processing implementations.
package tasks

import (
	"errors"
	"fmt"
	"jiggablend/internal/runner/api"
	"jiggablend/internal/runner/blender"
	"jiggablend/internal/runner/encoding"
	"jiggablend/internal/runner/sandbox"
	"jiggablend/internal/runner/workspace"
	"jiggablend/pkg/executils"
	"jiggablend/pkg/types"
	"os/exec"
	"sync"
	"time"
)

// Processor handles a specific task type.
type Processor interface {
	Process(ctx *Context) error
}

// Context provides task execution context.
type Context struct {
	TaskID   int64
	JobID    int64
	JobName  string
	Frame    int // frame start (inclusive); kept for backward compat
	FrameEnd int // frame end (inclusive); same as Frame for single-frame
	TaskType string
	WorkDir  string
	JobToken string
	Metadata *types.BlendMetadata

	Manager   *api.ManagerClient
	JobConn   *api.JobConnection
	Workspace *workspace.Manager
	Blender   *blender.Manager
	Encoder   *encoding.Selector
	Processes *executils.ProcessTracker

	// GPULockedOut is set when the runner has detected a GPU error (e.g. HIP) and disables GPU for all jobs.
	GPULockedOut bool
	// GPULockoutArmedThisAttempt is true when this task attempt newly enabled runner GPU lockout.
	// Failures with this set free-requeue without incrementing retry_count (expected ROCm warmup).
	GPULockoutArmedThisAttempt bool
	// HasAMD is true when the runner detected AMD devices at startup.
	HasAMD bool
	// HasNVIDIA is true when the runner detected NVIDIA GPUs at startup.
	HasNVIDIA bool
	// HasIntel is true when the runner detected Intel GPUs (e.g. Arc) at startup.
	HasIntel bool
	// GPUDetectionFailed is true when startup GPU backend detection could not run; we force CPU for all versions (backend availability unknown).
	GPUDetectionFailed bool
	// OnGPUError is called when a GPU error line is seen in render logs; typically sets runner GPU lockout.
	OnGPUError func()
	// ForceCPURendering is a runner-level override that forces CPU rendering for all jobs.
	ForceCPURendering bool
	// DisableRT disables GPU ray tracing acceleration (runner-level flag).
	DisableRT bool
	// HipGPUSampleBatch limits samples per GPU render pass on gfx115x (0 = no batching).
	HipGPUSampleBatch int
	// Sandbox wraps Blender execution (none/podman). Nil means none.
	Sandbox sandbox.Wrapper
}

// ErrJobCancelled indicates the manager-side job was cancelled during execution.
var ErrJobCancelled = errors.New("job cancelled")

// NewContext creates a new task context. frameEnd should be >= frame; if 0 or less than frame, it is treated as single-frame (frameEnd = frame).
// gpuLockedOut is the runner's current GPU lockout state; gpuDetectionFailed means detection failed at startup (force CPU for all versions); onGPUError is called when a GPU error is detected in logs (may be nil).
func NewContext(
	taskID, jobID int64,
	jobName string,
	frameStart, frameEnd int,
	taskType string,
	workDir string,
	jobToken string,
	metadata *types.BlendMetadata,
	manager *api.ManagerClient,
	jobConn *api.JobConnection,
	ws *workspace.Manager,
	blenderMgr *blender.Manager,
	encoder *encoding.Selector,
	processes *executils.ProcessTracker,
	gpuLockedOut bool,
	hasAMD bool,
	hasNVIDIA bool,
	hasIntel bool,
	gpuDetectionFailed bool,
	forceCPURendering bool,
	disableRT bool,
	hipGPUSampleBatch int,
	onGPUError func(),
	sb sandbox.Wrapper,
) *Context {
	if frameEnd < frameStart {
		frameEnd = frameStart
	}
	return &Context{
		TaskID:             taskID,
		JobID:              jobID,
		JobName:            jobName,
		Frame:              frameStart,
		FrameEnd:           frameEnd,
		TaskType:           taskType,
		WorkDir:            workDir,
		JobToken:           jobToken,
		Metadata:           metadata,
		Manager:            manager,
		JobConn:            jobConn,
		Workspace:          ws,
		Blender:            blenderMgr,
		Encoder:            encoder,
		Processes:          processes,
		GPULockedOut:       gpuLockedOut,
		HasAMD:             hasAMD,
		HasNVIDIA:          hasNVIDIA,
		HasIntel:           hasIntel,
		GPUDetectionFailed: gpuDetectionFailed,
		ForceCPURendering:  forceCPURendering,
		DisableRT:           disableRT,
		HipGPUSampleBatch:   hipGPUSampleBatch,
		OnGPUError:          onGPUError,
		Sandbox:             sb,
	}
}

// Log sends a log entry to the manager.
func (c *Context) Log(level types.LogLevel, message string) {
	if c.JobConn != nil {
		c.JobConn.Log(c.TaskID, level, message)
	}
}

// Info logs an info message.
func (c *Context) Info(message string) {
	c.Log(types.LogLevelInfo, message)
}

// Warn logs a warning message.
func (c *Context) Warn(message string) {
	c.Log(types.LogLevelWarn, message)
}

// Error logs an error message.
func (c *Context) Error(message string) {
	c.Log(types.LogLevelError, message)
}

// Progress sends a progress update.
func (c *Context) Progress(progress float64) {
	if c.JobConn != nil {
		c.JobConn.Progress(c.TaskID, progress)
	}
}

// OutputUploaded notifies that an output file was uploaded.
func (c *Context) OutputUploaded(fileName string) {
	if c.JobConn != nil {
		c.JobConn.OutputUploaded(c.TaskID, fileName)
	}
}

// Complete sends task completion.
// When the failure newly armed GPU lockout, free_requeue is set so the manager
// requeues without burning retry_count.
func (c *Context) Complete(success bool, errorMsg error) {
	if c.JobConn != nil {
		freeRequeue := !success && c.GPULockoutArmedThisAttempt
		c.JobConn.Complete(c.TaskID, success, errorMsg, freeRequeue)
	}
}

// GetOutputFormat returns the output format from metadata or default.
// Product default is EXR (Blender always renders EXR; deliverable may add video).
func (c *Context) GetOutputFormat() string {
	if c.Metadata != nil && c.Metadata.RenderSettings.OutputFormat != "" {
		return c.Metadata.RenderSettings.OutputFormat
	}
	return "EXR"
}

// GetFrameRate returns the frame rate from metadata or default.
func (c *Context) GetFrameRate() float64 {
	if c.Metadata != nil && c.Metadata.RenderSettings.FrameRate > 0 {
		return c.Metadata.RenderSettings.FrameRate
	}
	return 24.0
}

// GetBlenderVersion returns the Blender version from metadata.
func (c *Context) GetBlenderVersion() string {
	if c.Metadata != nil {
		return c.Metadata.BlenderVersion
	}
	return ""
}

// ShouldUnhideObjects returns whether to unhide objects.
func (c *Context) ShouldUnhideObjects() bool {
	return c.Metadata != nil && c.Metadata.UnhideObjects != nil && *c.Metadata.UnhideObjects
}

// ShouldEnableExecution returns whether to pass Blender --enable-autoexec
// (Python drivers/scripts inside the .blend). Job-bundled blender_addons/ install
// independently whenever that folder is present in the context.
func (c *Context) ShouldEnableExecution() bool {
	return c.Metadata != nil && c.Metadata.EnableExecution != nil && *c.Metadata.EnableExecution
}

// ShouldForceCPU returns true if GPU should be disabled and CPU rendering forced
// (runner GPU lockout, GPU detection failed at startup, or metadata force_cpu).
func (c *Context) ShouldForceCPU() bool {
	if c.ForceCPURendering {
		return true
	}
	if c.GPULockedOut {
		return true
	}
	// Detection failed at startup: backend availability unknown, so force CPU for all versions.
	if c.GPUDetectionFailed {
		return true
	}
	if c.Metadata != nil && c.Metadata.RenderSettings.EngineSettings != nil {
		if v, ok := c.Metadata.RenderSettings.EngineSettings["force_cpu"]; ok {
			if b, ok := v.(bool); ok && b {
				return true
			}
		}
	}
	return false
}

// GetCyclesSamples returns the target Cycles sample count from metadata (default 128).
func (c *Context) GetCyclesSamples() int {
	if c.Metadata != nil && c.Metadata.RenderSettings.EngineSettings != nil {
		if v, ok := c.Metadata.RenderSettings.EngineSettings["samples"]; ok {
			switch n := v.(type) {
			case float64:
				if int(n) > 0 {
					return int(n)
				}
			case int:
				if n > 0 {
					return n
				}
			}
		}
	}
	return 128
}

// GetCyclesSeed returns the Cycles seed from metadata (default 0).
func (c *Context) GetCyclesSeed() int {
	if c.Metadata != nil && c.Metadata.RenderSettings.EngineSettings != nil {
		if v, ok := c.Metadata.RenderSettings.EngineSettings["seed"]; ok {
			switch n := v.(type) {
			case float64:
				return int(n)
			case int:
				return n
			}
		}
	}
	return 0
}

// ShouldBatchGPUSamples returns true when HIP GPU renders should be split into sample batches.
func (c *Context) ShouldBatchGPUSamples() bool {
	if c.ShouldForceCPU() || c.HipGPUSampleBatch <= 0 {
		return false
	}
	return c.GetCyclesSamples() > c.HipGPUSampleBatch
}

// ShouldDisableRT returns true when GPU ray tracing acceleration should be disabled.
func (c *Context) ShouldDisableRT() bool {
	if c.DisableRT {
		return true
	}
	if c.Metadata != nil && c.Metadata.RenderSettings.EngineSettings != nil {
		settings := c.Metadata.RenderSettings.EngineSettings
		if v, ok := settings["disable_rt"]; ok {
			if b, ok := v.(bool); ok && b {
				return true
			}
		}
		if v, ok := settings["disable_hiprt"]; ok {
			if b, ok := v.(bool); ok && b {
				return true
			}
		}
	}
	return false
}

// IsJobCancelled checks whether the manager marked this job as cancelled.
func (c *Context) IsJobCancelled() (bool, error) {
	if c.Manager == nil {
		return false, nil
	}
	status, err := c.Manager.GetJobStatus(c.JobID)
	if err != nil {
		return false, err
	}
	return status == types.JobStatusCancelled, nil
}

// CheckCancelled returns ErrJobCancelled if the job was cancelled.
func (c *Context) CheckCancelled() error {
	cancelled, err := c.IsJobCancelled()
	if err != nil {
		return fmt.Errorf("failed to check job status: %w", err)
	}
	if cancelled {
		return ErrJobCancelled
	}
	return nil
}

// StartCancellationMonitor polls manager status and kills cmd if job is cancelled.
// Caller must invoke returned stop function when cmd exits.
func (c *Context) StartCancellationMonitor(cmd *exec.Cmd, taskLabel string) func() {
	stop := make(chan struct{})
	var once sync.Once

	go func() {
		ticker := time.NewTicker(2 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-stop:
				return
			case <-ticker.C:
				cancelled, err := c.IsJobCancelled()
				if err != nil {
					c.Warn(fmt.Sprintf("Could not check cancellation for %s task: %v", taskLabel, err))
					continue
				}
				if !cancelled {
					continue
				}
				c.Warn(fmt.Sprintf("Job %d was cancelled, stopping %s task early", c.JobID, taskLabel))
				if cmd != nil && cmd.Process != nil {
					_ = cmd.Process.Kill()
				}
				return
			}
		}
	}()

	return func() {
		once.Do(func() {
			close(stop)
		})
	}
}