Refactor installation and runner scripts for enhanced usability and security

- Updated the installation script to clarify that production wrappers do not include fixed test secrets, improving user guidance for local testing.
- Modified the jiggablend-manager and jiggablend-runner scripts to remove fixed test configurations, emphasizing the use of local test credentials via `make init-test`.
- Enhanced error handling in the runner script to ensure that an API key is provided, improving security and user feedback.
- Added new flags and options for the runner, including sandboxing capabilities and GPU ray tracing control, enhancing flexibility for users.
- Improved README documentation to reflect changes in usage and configuration, ensuring users have clear instructions for setup and execution.
This commit is contained in:
2026-07-12 10:01:15 -05:00
parent a3defe5cf6
commit 1a69fcfd04
47 changed files with 2718 additions and 402 deletions
+271 -89
View File
@@ -63,21 +63,20 @@ func (s *Manager) runnerAuthMiddleware(next http.HandlerFunc) http.HandlerFunc {
return
}
// For fixed API keys, skip database verification
// Always verify the runner exists. For non-fixed keys, also check API key ownership.
var dbAPIKeyID sql.NullInt64
err = s.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT api_key_id FROM runners WHERE id = ?", runnerID).Scan(&dbAPIKeyID)
})
if err == sql.ErrNoRows {
s.respondError(w, http.StatusNotFound, "runner not found")
return
}
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to query runner API key: %v", err))
return
}
if apiKeyID != -1 {
// Verify runner exists and uses this API key
var dbAPIKeyID sql.NullInt64
err = s.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT api_key_id FROM runners WHERE id = ?", runnerID).Scan(&dbAPIKeyID)
})
if err == sql.ErrNoRows {
s.respondError(w, http.StatusNotFound, "runner not found")
return
}
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("failed to query runner API key: %v", err))
return
}
if !dbAPIKeyID.Valid || dbAPIKeyID.Int64 != apiKeyID {
s.respondError(w, http.StatusForbidden, "runner does not belong to this API key")
return
@@ -1000,6 +999,32 @@ func (s *Manager) handleUploadFileWithToken(w http.ResponseWriter, r *http.Reque
})
}
// runnerCanAccessJob returns true if the runner has been assigned at least one task on the job.
// Used to scope file/metadata/status APIs so a valid API key cannot read arbitrary jobs.
func (s *Manager) runnerCanAccessJob(runnerID, jobID int64) bool {
if runnerID <= 0 || jobID <= 0 {
return false
}
var n int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
`SELECT COUNT(*) FROM tasks WHERE job_id = ? AND runner_id = ?`,
jobID, runnerID,
).Scan(&n)
})
return err == nil && n > 0
}
// requireRunnerJobAccess enforces that the authenticated runner may access the job.
func (s *Manager) requireRunnerJobAccess(w http.ResponseWriter, r *http.Request, jobID int64) bool {
runnerID, _ := r.Context().Value(runnerIDContextKey).(int64)
if !s.runnerCanAccessJob(runnerID, jobID) {
s.respondError(w, http.StatusForbidden, "runner is not assigned to this job")
return false
}
return true
}
// handleGetJobStatusForRunner allows runners to check job status
func (s *Manager) handleGetJobStatusForRunner(w http.ResponseWriter, r *http.Request) {
jobID, err := parseID(r, "jobId")
@@ -1007,6 +1032,9 @@ func (s *Manager) handleGetJobStatusForRunner(w http.ResponseWriter, r *http.Req
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
var job types.Job
var startedAt, completedAt sql.NullTime
@@ -1069,6 +1097,9 @@ func (s *Manager) handleGetJobFilesForRunner(w http.ResponseWriter, r *http.Requ
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
runnerID := r.URL.Query().Get("runner_id")
log.Printf("GetJobFiles request for job %d from runner %s", jobID, runnerID)
@@ -1127,6 +1158,9 @@ func (s *Manager) handleGetJobMetadataForRunner(w http.ResponseWriter, r *http.R
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
var blendMetadataJSON sql.NullString
err = s.db.With(func(conn *sql.DB) error {
@@ -1166,6 +1200,9 @@ func (s *Manager) handleDownloadFileForRunner(w http.ResponseWriter, r *http.Req
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
if !s.requireRunnerJobAccess(w, r, jobID) {
return
}
// Get fileName from URL path (may need URL decoding)
fileName := chi.URLParam(r, "fileName")
@@ -1259,11 +1296,14 @@ type WSLogEntry struct {
}
type WSTaskUpdate struct {
TaskID int64 `json:"task_id"`
Status string `json:"status"`
OutputPath string `json:"output_path,omitempty"`
Success bool `json:"success"`
Error string `json:"error,omitempty"`
TaskID int64 `json:"task_id"`
Status string `json:"status"`
OutputPath string `json:"output_path,omitempty"`
Success bool `json:"success"`
Error string `json:"error,omitempty"`
// FreeRequeue asks the manager to requeue without incrementing retry_count
// (e.g. this attempt newly armed GPU lockout after a ROCm fault).
FreeRequeue bool `json:"free_requeue,omitempty"`
}
// handleRunnerJobWebSocket handles per-job WebSocket connections from runners
@@ -1362,25 +1402,18 @@ func (s *Manager) handleRunnerJobWebSocket(w http.ResponseWriter, r *http.Reques
delete(s.runnerJobConnsWriteMu, connKey)
s.runnerJobConnsMu.Unlock()
// Check if task is still running - if so, mark as failed
// If the runner never delivered task_complete (crash, bad payload, network drop),
// treat this like a task failure with retries. Prefer the last ERROR log so
// segfaults surface as such instead of a generic WebSocket message.
var currentStatus string
s.db.With(func(conn *sql.DB) error {
return conn.QueryRow("SELECT status FROM tasks WHERE id = ?", taskID).Scan(&currentStatus)
})
if currentStatus == string(types.TaskStatusRunning) {
log.Printf("Job WebSocket disconnected unexpectedly for task %d, marking as failed", taskID)
s.db.With(func(conn *sql.DB) error {
_, err := conn.Exec(
`UPDATE tasks SET status = ?, runner_id = NULL, error_message = ?, completed_at = ? WHERE id = ?`,
types.TaskStatusFailed, "WebSocket connection lost", time.Now(), taskID,
)
return err
})
s.broadcastTaskUpdate(jobID, taskID, "task_update", map[string]interface{}{
"status": types.TaskStatusFailed,
"error_message": "WebSocket connection lost",
})
s.updateJobStatusFromTasks(jobID)
errorMsg := s.inferTaskFailureFromLogs(taskID)
freeRequeue := s.taskLogsIndicateGPULockoutArm(taskID)
log.Printf("Job WebSocket disconnected unexpectedly for task %d: %s (free_requeue=%v)", taskID, errorMsg, freeRequeue)
s.handleTaskFailureWithRetry(runnerID, taskID, jobID, errorMsg, freeRequeue)
}
log.Printf("Job WebSocket closed: job=%d, runner=%d, task=%d", jobID, runnerID, taskID)
@@ -1467,13 +1500,19 @@ func (s *Manager) handleRunnerJobWebSocket(w http.ResponseWriter, r *http.Reques
}
case "task_complete":
var taskUpdate WSTaskUpdate
if err := json.Unmarshal(msg.Data, &taskUpdate); err == nil {
if taskUpdate.TaskID == taskID {
s.handleWebSocketTaskComplete(runnerID, taskUpdate)
// Task is done, close connection
return
}
taskUpdate, err := parseWSTaskUpdate(msg.Data)
if err != nil {
log.Printf("Job WebSocket task_complete for task %d: %v", taskID, err)
// Still close; disconnect handler will infer failure from ERROR logs.
return
}
if taskUpdate.TaskID == 0 {
taskUpdate.TaskID = taskID
}
if taskUpdate.TaskID == taskID {
s.handleWebSocketTaskComplete(runnerID, taskUpdate)
// Task is done, close connection
return
}
case "runner_heartbeat":
s.handleWSRunnerHeartbeat(conn, jobID)
@@ -1576,12 +1615,11 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
// Verify task belongs to runner and get task info
var taskRunnerID sql.NullInt64
var jobID int64
var retryCount, maxRetries int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
"SELECT runner_id, job_id, retry_count, max_retries FROM tasks WHERE id = ?",
"SELECT runner_id, job_id FROM tasks WHERE id = ?",
taskUpdate.TaskID,
).Scan(&taskRunnerID, &jobID, &retryCount, &maxRetries)
).Scan(&taskRunnerID, &jobID)
})
if err != nil {
log.Printf("Failed to get task %d info: %v", taskUpdate.TaskID, err)
@@ -1597,7 +1635,10 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
// Handle successful completion
if taskUpdate.Success {
err = s.db.WithTx(func(tx *sql.Tx) error {
_, err := tx.Exec(`UPDATE tasks SET status = ? WHERE id = ?`, types.TaskStatusCompleted, taskUpdate.TaskID)
_, err := tx.Exec(
`UPDATE tasks SET status = ?, error_message = NULL, completed_at = ? WHERE id = ?`,
types.TaskStatusCompleted, now, taskUpdate.TaskID,
)
if err != nil {
return err
}
@@ -1607,8 +1648,7 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
return err
}
}
_, err = tx.Exec(`UPDATE tasks SET completed_at = ? WHERE id = ?`, now, taskUpdate.TaskID)
return err
return nil
})
if err != nil {
log.Printf("Failed to update task %d: %v", taskUpdate.TaskID, err)
@@ -1617,89 +1657,231 @@ func (s *Manager) handleWebSocketTaskComplete(runnerID int64, taskUpdate WSTaskU
// Broadcast task update
s.broadcastTaskUpdate(jobID, taskUpdate.TaskID, "task_update", map[string]interface{}{
"status": types.TaskStatusCompleted,
"output_path": taskUpdate.OutputPath,
"completed_at": now,
"status": types.TaskStatusCompleted,
"output_path": taskUpdate.OutputPath,
"completed_at": now,
"error_message": nil,
})
s.updateJobStatusFromTasks(jobID)
return
}
// Handle task failure - this is an actual task failure (e.g., Blender crash)
// Check if we have retries remaining
if retryCount < maxRetries {
// Reset to pending for retry - increment retry_count
// Actual task failure (e.g. Blender segfault) — retry when retries remain.
// free_requeue: failure newly armed GPU lockout; requeue without burning a retry.
errorMsg := taskUpdate.Error
if errorMsg == "" {
errorMsg = s.inferTaskFailureFromLogs(taskUpdate.TaskID)
}
freeRequeue := taskUpdate.FreeRequeue || s.taskLogsIndicateGPULockoutArm(taskUpdate.TaskID)
s.handleTaskFailureWithRetry(runnerID, taskUpdate.TaskID, jobID, errorMsg, freeRequeue)
}
// defaultUnexpectedDisconnectError is used when the runner drops the job WebSocket
// without a task_complete message and without any ERROR logs.
const defaultUnexpectedDisconnectError = "WebSocket connection lost"
// parseWSTaskUpdate decodes a task_complete payload. Older runners marshaled
// the error field as a Go error interface (JSON object {}), which fails strict
// unmarshaling into a string — fall back to extracting fields loosely so
// success/failure is still recognized.
func parseWSTaskUpdate(data json.RawMessage) (WSTaskUpdate, error) {
var update WSTaskUpdate
if err := json.Unmarshal(data, &update); err == nil {
return update, nil
}
var loose map[string]interface{}
if err := json.Unmarshal(data, &loose); err != nil {
return WSTaskUpdate{}, err
}
if v, ok := loose["task_id"].(float64); ok {
update.TaskID = int64(v)
}
if v, ok := loose["success"].(bool); ok {
update.Success = v
}
if v, ok := loose["output_path"].(string); ok {
update.OutputPath = v
}
if v, ok := loose["status"].(string); ok {
update.Status = v
}
if v, ok := loose["free_requeue"].(bool); ok {
update.FreeRequeue = v
}
switch v := loose["error"].(type) {
case string:
update.Error = v
case map[string]interface{}:
// Legacy runner bug: error interface marshaled as {}
update.Error = ""
}
return update, nil
}
// taskLogsIndicateGPULockoutArm is true when this attempt's logs show the runner
// newly armed GPU lockout (free-requeue path when task_complete lacked free_requeue).
func (s *Manager) taskLogsIndicateGPULockoutArm(taskID int64) bool {
var n int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
`SELECT COUNT(*) FROM task_logs
WHERE task_id = ?
AND (
message LIKE '%GPU error detected%'
OR message LIKE '%free-requeue%'
OR message LIKE '%GPU disabled for subsequent jobs%'
)`,
taskID,
).Scan(&n)
})
return err == nil && n > 0
}
// inferTaskFailureFromLogs returns the best failure reason for a task that died
// without a usable task_complete payload. Prefers the latest ERROR log so Blender
// segfaults (and similar crashes) are stored instead of a generic disconnect message.
func (s *Manager) inferTaskFailureFromLogs(taskID int64) string {
var lastError string
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
`SELECT message FROM task_logs
WHERE task_id = ? AND UPPER(log_level) = ?
ORDER BY id DESC LIMIT 1`,
taskID, string(types.LogLevelError),
).Scan(&lastError)
})
if err != nil || strings.TrimSpace(lastError) == "" {
return defaultUnexpectedDisconnectError
}
return formatInferredTaskFailure(lastError)
}
// formatInferredTaskFailure normalizes streamed runner ERROR lines into a concise
// task error_message. Blender crashes often arrive as "Task failed: blender failed: signal: ..."
// or "Blender render failed: ..."; keep the underlying failure text.
func formatInferredTaskFailure(logMessage string) string {
msg := strings.TrimSpace(logMessage)
for _, prefix := range []string{
"Task failed: ",
"Blender render failed: ",
} {
if strings.HasPrefix(msg, prefix) {
msg = strings.TrimSpace(strings.TrimPrefix(msg, prefix))
}
}
if msg == "" {
return defaultUnexpectedDisconnectError
}
return msg
}
// handleTaskFailureWithRetry applies the task failure policy used for Blender
// crashes and unexpected job-WebSocket drops: requeue as pending while
// retry_count < max_retries, otherwise mark failed permanently with error_message.
// When freeRequeue is true (GPU lockout newly armed this attempt), the task is
// always requeued as pending without incrementing retry_count.
func (s *Manager) handleTaskFailureWithRetry(runnerID, taskID, jobID int64, errorMsg string, freeRequeue bool) {
var retryCount, maxRetries int
err := s.db.With(func(conn *sql.DB) error {
return conn.QueryRow(
"SELECT retry_count, max_retries FROM tasks WHERE id = ?",
taskID,
).Scan(&retryCount, &maxRetries)
})
if err != nil {
log.Printf("Failed to get retry info for task %d: %v", taskID, err)
return
}
if errorMsg == "" {
errorMsg = defaultUnexpectedDisconnectError
}
if freeRequeue {
// Make the reason obvious in the UI without looking like a permanent fail.
if !strings.Contains(strings.ToLower(errorMsg), "gpu lockout") {
errorMsg = "GPU lockout armed (free requeue): " + errorMsg
}
}
now := time.Now()
// Free requeue always goes back to pending (expected ROCm GPU fault that armed lockout).
// Otherwise only requeue while retries remain.
canRequeue := freeRequeue || retryCount < maxRetries
if canRequeue {
newRetryCount := retryCount
retrySQL := `UPDATE tasks SET status = ?, runner_id = NULL, current_step = NULL,
error_message = ?, started_at = NULL, completed_at = NULL
WHERE id = ?`
args := []interface{}{types.TaskStatusPending, errorMsg, taskID}
if !freeRequeue {
retrySQL = `UPDATE tasks SET status = ?, runner_id = NULL, current_step = NULL,
error_message = ?, retry_count = retry_count + 1, started_at = NULL, completed_at = NULL
WHERE id = ?`
newRetryCount = retryCount + 1
}
err = s.db.WithTx(func(tx *sql.Tx) error {
_, err := tx.Exec(
`UPDATE tasks SET status = ?, runner_id = NULL, current_step = NULL,
retry_count = retry_count + 1, started_at = NULL, completed_at = NULL
WHERE id = ?`,
types.TaskStatusPending, taskUpdate.TaskID,
)
_, err := tx.Exec(retrySQL, args...)
if err != nil {
return err
}
// Clear steps and logs for fresh retry
_, err = tx.Exec(`DELETE FROM task_steps WHERE task_id = ?`, taskUpdate.TaskID)
_, err = tx.Exec(`DELETE FROM task_steps WHERE task_id = ?`, taskID)
if err != nil {
return err
}
_, err = tx.Exec(`DELETE FROM task_logs WHERE task_id = ?`, taskUpdate.TaskID)
_, err = tx.Exec(`DELETE FROM task_logs WHERE task_id = ?`, taskID)
return err
})
if err != nil {
log.Printf("Failed to reset task %d for retry: %v", taskUpdate.TaskID, err)
log.Printf("Failed to reset task %d for retry: %v", taskID, err)
return
}
// Broadcast task reset to clients (includes steps_cleared and logs_cleared flags)
s.broadcastTaskUpdate(jobID, taskUpdate.TaskID, "task_reset", map[string]interface{}{
s.broadcastTaskUpdate(jobID, taskID, "task_reset", map[string]interface{}{
"status": types.TaskStatusPending,
"retry_count": retryCount + 1,
"error_message": taskUpdate.Error,
"retry_count": newRetryCount,
"error_message": errorMsg,
"steps_cleared": true,
"logs_cleared": true,
"free_requeue": freeRequeue,
})
log.Printf("Task %d failed but has retries remaining (%d/%d), reset to pending", taskUpdate.TaskID, retryCount+1, maxRetries)
if freeRequeue {
log.Printf("Task %d failed (%s); free-requeued without burning retry (still %d/%d used)",
taskID, errorMsg, retryCount, maxRetries)
} else {
log.Printf("Task %d failed (%s) but has retries remaining (%d/%d), reset to pending",
taskID, errorMsg, newRetryCount, maxRetries)
}
} else {
// No retries remaining - mark as failed
err = s.db.WithTx(func(tx *sql.Tx) error {
_, err := tx.Exec(
`UPDATE tasks SET status = ?, runner_id = NULL, completed_at = ? WHERE id = ?`,
types.TaskStatusFailed, now, taskUpdate.TaskID,
`UPDATE tasks SET status = ?, runner_id = NULL, completed_at = ?, error_message = ? WHERE id = ?`,
types.TaskStatusFailed, now, errorMsg, taskID,
)
if err != nil {
return err
}
if taskUpdate.Error != "" {
_, err = tx.Exec(`UPDATE tasks SET error_message = ? WHERE id = ?`, taskUpdate.Error, taskUpdate.TaskID)
if err != nil {
return err
}
}
return nil
return err
})
if err != nil {
log.Printf("Failed to mark task %d as failed: %v", taskUpdate.TaskID, err)
log.Printf("Failed to mark task %d as failed: %v", taskID, err)
return
}
// Log the final failure
s.logTaskEvent(taskUpdate.TaskID, &runnerID, types.LogLevelError,
fmt.Sprintf("Task failed permanently after %d retries: %s", maxRetries, taskUpdate.Error), "")
s.logTaskEvent(taskID, &runnerID, types.LogLevelError,
fmt.Sprintf("Task failed permanently after %d retries: %s", maxRetries, errorMsg), "")
// Broadcast task update
s.broadcastTaskUpdate(jobID, taskUpdate.TaskID, "task_update", map[string]interface{}{
s.broadcastTaskUpdate(jobID, taskID, "task_update", map[string]interface{}{
"status": types.TaskStatusFailed,
"completed_at": now,
"error_message": taskUpdate.Error,
"error_message": errorMsg,
})
log.Printf("Task %d failed permanently after %d retries", taskUpdate.TaskID, maxRetries)
log.Printf("Task %d failed permanently after %d retries: %s", taskID, maxRetries, errorMsg)
}
// Update job status and progress
s.updateJobStatusFromTasks(jobID)
}