Execution and Partition-Level Environment Variables (#4801)

This PR introduces two key features:

## 1. Enhanced Execution Environment Variables

Added support for passing rich job metadata to execution engines via
environment variables including:

- `BACALHAU_PARTITION_INDEX`: Current partition index (0 to N-1)
- `BACALHAU_PARTITION_COUNT`: Total number of partitions
- `BACALHAU_JOB_ID`: Unique job identifier
- `BACALHAU_JOB_NAME`: User-provided job name
- `BACALHAU_JOB_NAMESPACE`: Job namespace 
- `BACALHAU_JOB_TYPE`: Job type (Batch/Service)
- `BACALHAU_EXECUTION_ID`: Unique execution identifier
- `BACALHAU_NODE_ID`: ID of executing compute node

This allows jobs to:
- Be partition-aware and handle their specific partition's work
- Access their execution context
- Track node assignment

## 2. Test Suite for Partition Scheduling

Added comprehensive test suite that validates:
- Environment variable propagation to executors
- Partition scheduling behavior:
  - Unique partition indices
  - Node distribution
  - Retry behavior
  - Service job continuous execution

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

Here are the release notes for this pull request:

**New Features**
- Added environment variable management for job executions
- Enhanced support for system and task-level environment variables
- Improved job partitioning and execution context handling

**Bug Fixes**
- Fixed potential nil slice access in job task retrieval
- Added validation for environment variable naming conventions

**Improvements**
- Streamlined executor and job handling interfaces
- Added utility functions for environment variable manipulation
- Enhanced test coverage for job execution scenarios

**Technical Enhancements**
- Refactored execution context management
- Improved error handling in task and job validation
- Added robust environment variable sanitization and merging
capabilities

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: wdbaruni

pkg/compute/envvars.go

@@ -0,0 +1,44 @@
+// File: pkg/compute/envvars.go
+
+package compute
+
+import (
+	"fmt"
+
+	"github.com/bacalhau-project/bacalhau/pkg/lib/envvar"
+	"github.com/bacalhau-project/bacalhau/pkg/models"
+)
+
+// GetExecutionEnvVars returns a map of environment variables that should be passed to the execution.
+// Task-level environment variables are included but system variables take precedence.
+func GetExecutionEnvVars(execution *models.Execution) map[string]string {
+	if execution == nil {
+		return make(map[string]string)
+	}
+
+	// Start with task-level environment variables if they exist
+	taskEnv := make(map[string]string)
+	if execution.Job != nil && execution.Job.Task() != nil && execution.Job.Task().Env != nil {
+		taskEnv = execution.Job.Task().Env
+	}
+
+	// Build system environment variables
+	sysEnv := make(map[string]string)
+	sysEnv[models.EnvVarPrefix+"EXECUTION_ID"] = execution.ID
+	sysEnv[models.EnvVarPrefix+"NODE_ID"] = envvar.Sanitize(execution.NodeID)
+
+	// Add job-related environment variables if job exists
+	if execution.Job != nil {
+		sysEnv[models.EnvVarPrefix+"JOB_ID"] = execution.JobID
+		sysEnv[models.EnvVarPrefix+"JOB_NAME"] = envvar.Sanitize(execution.Job.Name)
+		sysEnv[models.EnvVarPrefix+"JOB_NAMESPACE"] = envvar.Sanitize(execution.Job.Namespace)
+		sysEnv[models.EnvVarPrefix+"JOB_TYPE"] = execution.Job.Type
+
+		// Add partition-related environment variables
+		sysEnv[models.EnvVarPrefix+"PARTITION_INDEX"] = fmt.Sprintf("%d", execution.PartitionIndex)
+		sysEnv[models.EnvVarPrefix+"PARTITION_COUNT"] = fmt.Sprintf("%d", execution.Job.Count)
+	}
+
+	// Merge task and system variables, with system taking precedence
+	return envvar.Merge(taskEnv, sysEnv)
+}

pkg/compute/envvars_test.go

@@ -0,0 +1,140 @@
+//go:build unit || !integration
+
+package compute
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/bacalhau-project/bacalhau/pkg/models"
+)
+
+func TestGetExecutionEnvVars(t *testing.T) {
+	tests := []struct {
+		name      string
+		execution *models.Execution
+		want      map[string]string
+	}{
+		{
+			name:      "nil execution returns empty map",
+			execution: nil,
+			want:      map[string]string{},
+		},
+		{
+			name: "basic execution without job",
+			execution: &models.Execution{
+				ID:     "exec-1",
+				NodeID: "node-1",
+			},
+			want: map[string]string{
+				"BACALHAU_EXECUTION_ID": "exec-1",
+				"BACALHAU_NODE_ID":      "node-1",
+			},
+		},
+		{
+			name: "execution with job but no task env",
+			execution: &models.Execution{
+				ID:             "exec-1",
+				NodeID:         "node-1",
+				JobID:          "job-1",
+				PartitionIndex: 0,
+				Job: &models.Job{
+					Name:      "test-job",
+					Namespace: "default",
+					Type:      "batch",
+					Tasks: []*models.Task{
+						{
+							Name: "task-1",
+						},
+					},
+					Count: 3,
+				},
+			},
+			want: map[string]string{
+				"BACALHAU_EXECUTION_ID":    "exec-1",
+				"BACALHAU_NODE_ID":         "node-1",
+				"BACALHAU_JOB_ID":          "job-1",
+				"BACALHAU_JOB_NAME":        "test-job",
+				"BACALHAU_JOB_NAMESPACE":   "default",
+				"BACALHAU_JOB_TYPE":        "batch",
+				"BACALHAU_PARTITION_INDEX": "0",
+				"BACALHAU_PARTITION_COUNT": "3",
+			},
+		},
+		{
+			name: "execution with task env",
+			execution: &models.Execution{
+				ID:             "exec-1",
+				NodeID:         "node-1",
+				JobID:          "job-1",
+				PartitionIndex: 0,
+				Job: &models.Job{
+					Name:      "test-job",
+					Namespace: "default",
+					Type:      "batch",
+					Tasks: []*models.Task{
+						{
+							Name: "task-1",
+							Env: map[string]string{
+								"MY_VAR":           "my-value",
+								"BACALHAU_NODE_ID": "should-not-override", // Should not override system env
+								"OTHER_VAR":        "other-value",
+							},
+						},
+					},
+					Count: 3,
+				},
+			},
+			want: map[string]string{
+				"BACALHAU_EXECUTION_ID":    "exec-1",
+				"BACALHAU_NODE_ID":         "node-1", // System value takes precedence
+				"BACALHAU_JOB_ID":          "job-1",
+				"BACALHAU_JOB_NAME":        "test-job",
+				"BACALHAU_JOB_NAMESPACE":   "default",
+				"BACALHAU_JOB_TYPE":        "batch",
+				"BACALHAU_PARTITION_INDEX": "0",
+				"BACALHAU_PARTITION_COUNT": "3",
+				"MY_VAR":                   "my-value",
+				"OTHER_VAR":                "other-value",
+			},
+		},
+		{
+			name: "execution with special characters in names",
+			execution: &models.Execution{
+				ID:             "exec-1",
+				NodeID:         "node-1",
+				JobID:          "job-1",
+				PartitionIndex: 0,
+				Job: &models.Job{
+					Name:      "test=job with spaces",
+					Namespace: "test=namespace",
+					Type:      "batch",
+					Tasks: []*models.Task{
+						{
+							Name: "task-1",
+						},
+					},
+					Count: 1,
+				},
+			},
+			want: map[string]string{
+				"BACALHAU_EXECUTION_ID":    "exec-1",
+				"BACALHAU_NODE_ID":         "node-1",
+				"BACALHAU_JOB_ID":          "job-1",
+				"BACALHAU_JOB_NAME":        "test_job_with_spaces", // Sanitized
+				"BACALHAU_JOB_NAMESPACE":   "test_namespace",       // Sanitized
+				"BACALHAU_JOB_TYPE":        "batch",
+				"BACALHAU_PARTITION_INDEX": "0",
+				"BACALHAU_PARTITION_COUNT": "1",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := GetExecutionEnvVars(tt.execution)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

pkg/compute/executor.go

@@ -183,6 +183,7 @@ func PrepareRunArguments(
 			Inputs:       inputVolumes,
 			ResultsDir:   resultsDir,
 			EngineParams: engineArgs,
+			Env:          GetExecutionEnvVars(execution),
 			OutputLimits: executor.OutputLimits{
 				MaxStdoutFileLength:   system.MaxStdoutFileLength,
 				MaxStdoutReturnLength: system.MaxStdoutReturnLength,

pkg/executor/docker/executor.go

@@ -21,6 +21,7 @@ import (
 	"github.com/bacalhau-project/bacalhau/pkg/config/types"
 	"github.com/bacalhau-project/bacalhau/pkg/config_legacy"
 	dockermodels "github.com/bacalhau-project/bacalhau/pkg/executor/docker/models"
+	"github.com/bacalhau-project/bacalhau/pkg/lib/envvar"
 	"github.com/bacalhau-project/bacalhau/pkg/models"
 	"github.com/bacalhau-project/bacalhau/pkg/models/messages"
 	pkgUtil "github.com/bacalhau-project/bacalhau/pkg/util"
@@ -137,16 +138,7 @@ func (e *Executor) Start(ctx context.Context, request *executor.RunCommandReques
 			}
 		}
 
-		jobContainer, err := e.newDockerJobContainer(ctx, &dockerJobContainerParams{
-			ExecutionID:   request.ExecutionID,
-			JobID:         request.JobID,
-			EngineSpec:    request.EngineParams,
-			NetworkConfig: request.Network,
-			Resources:     request.Resources,
-			Inputs:        request.Inputs,
-			Outputs:       request.Outputs,
-			ResultsDir:    request.ResultsDir,
-		})
+		jobContainer, err := e.newDockerJobContainer(ctx, request)
 		if err != nil {
 			return err
 		}
@@ -309,32 +301,28 @@ func (e *Executor) Run(
 	}
 }
 
-type dockerJobContainerParams struct {
-	ExecutionID   string
-	JobID         string
-	EngineSpec    *models.SpecConfig
-	NetworkConfig *models.NetworkConfig
-	Resources     *models.Resources
-	Inputs        []storage.PreparedStorage
-	Outputs       []*models.ResultPath
-	ResultsDir    string
-}
-
 // newDockerJobContainer is an internal method called by Start to set up a new Docker container
 // for the job execution. It configures the container based on the provided dockerJobContainerParams.
 // This includes decoding engine specifications, setting up environment variables, mounts, resource
 // constraints, and network configurations. It then creates the container but does not start it.
 // The method returns a container.CreateResponse and an error if any part of the setup fails.
-func (e *Executor) newDockerJobContainer(ctx context.Context, params *dockerJobContainerParams) (container.CreateResponse, error) {
+func (e *Executor) newDockerJobContainer(ctx context.Context, params *executor.RunCommandRequest) (container.CreateResponse, error) {
 	// decode the request arguments, bail if they are invalid.
-	dockerArgs, err := dockermodels.DecodeSpec(params.EngineSpec)
+	dockerArgs, err := dockermodels.DecodeSpec(params.EngineParams)
 	if err != nil {
 		return container.CreateResponse{}, fmt.Errorf("decoding engine spec: %w", err)
 	}
+
+	// merge both the job level and engine level environment variables
+	envVars := envvar.MergeSlices(
+		envvar.ToSlice(params.Env),
+		dockerArgs.EnvironmentVariables,
+	)
+
 	containerConfig := &container.Config{
 		Image:      dockerArgs.Image,
 		Tty:        false,
-		Env:        dockerArgs.EnvironmentVariables,
+		Env:        envVars,
 		Entrypoint: dockerArgs.Entrypoint,
 		Cmd:        dockerArgs.Parameters,
 		Labels:     e.containerLabels(params.ExecutionID, params.JobID),
@@ -379,7 +367,7 @@ func (e *Executor) newDockerJobContainer(ctx context.Context, params *dockerJobC
 	}
 	log.Ctx(ctx).Trace().Msgf("Container: %+v %+v", containerConfig, mounts)
 	// Create a network if the job requests it, modifying the containerConfig and hostConfig.
-	err = e.setupNetworkForJob(ctx, params.JobID, params.ExecutionID, params.NetworkConfig, containerConfig, hostConfig)
+	err = e.setupNetworkForJob(ctx, params.JobID, params.ExecutionID, params.Network, containerConfig, hostConfig)
 	if err != nil {
 		return container.CreateResponse{}, fmt.Errorf("setting up network: %w", err)
 	}

pkg/executor/docker/executor_test.go

@@ -172,6 +172,7 @@ func (s *ExecutorTestSuite) runJobWithContext(ctx context.Context, spec *models.
 			Inputs:       nil,
 			ResultsDir:   result,
 			EngineParams: spec.Engine,
+			Env:          spec.Env,
 			OutputLimits: executor.OutputLimits{
 				MaxStdoutFileLength:   system.MaxStdoutFileLength,
 				MaxStdoutReturnLength: system.MaxStdoutReturnLength,
@@ -577,3 +578,76 @@ func (s *ExecutorTestSuite) TestDockerOOM() {
 	require.NoError(s.T(), err)
 	require.Contains(s.T(), result.ErrorMsg, "memory limit exceeded")
 }
+
+func (s *ExecutorTestSuite) TestDockerEnvironmentVariables() {
+	tests := []struct {
+		name      string
+		taskEnv   map[string]string
+		engineEnv []string
+		checkVars []string // variables to check in order
+		want      string
+	}{
+		{
+			name: "task environment variables",
+			taskEnv: map[string]string{
+				"TEST_VAR":    "test_value",
+				"ANOTHER_VAR": "another_value",
+			},
+			checkVars: []string{"TEST_VAR", "ANOTHER_VAR"},
+			want:      "test_value\nanother_value",
+		},
+		{
+			name: "engine environment variables",
+			engineEnv: []string{
+				"TEST_VAR=engine_value",
+				"ENGINE_VAR=engine_only",
+			},
+			checkVars: []string{"TEST_VAR", "ENGINE_VAR"},
+			want:      "engine_value\nengine_only",
+		},
+		{
+			name: "merged environment variables with engine precedence",
+			taskEnv: map[string]string{
+				"TEST_VAR": "task_value",
+				"TASK_VAR": "task_only",
+			},
+			engineEnv: []string{
+				"TEST_VAR=engine_value",
+				"ENGINE_VAR=engine_only",
+			},
+			checkVars: []string{"TEST_VAR", "TASK_VAR", "ENGINE_VAR"},
+			want:      "engine_value\ntask_only\nengine_only",
+		},
+	}
+
+	for _, tt := range tests {
+		s.Run(tt.name, func() {
+			// Create simple script that prints vars in order
+			script := strings.Builder{}
+			for _, v := range tt.checkVars {
+				script.WriteString(fmt.Sprintf("echo $%s\n", v))
+			}
+
+			builder := dockermodels.NewDockerEngineBuilder("busybox:1.37.0").
+				WithEntrypoint("sh", "-c", script.String())
+
+			if len(tt.engineEnv) > 0 {
+				builder = builder.WithEnvironmentVariables(tt.engineEnv...)
+			}
+
+			es, err := builder.Build()
+			s.Require().NoError(err)
+
+			task := mock.Task()
+			task.Engine = es
+			task.Env = tt.taskEnv
+
+			result, err := s.runJob(task, uuid.New().String())
+			require.NoError(s.T(), err)
+			require.Zero(s.T(), result.ExitCode, result.STDERR)
+
+			output := strings.TrimSpace(result.STDOUT)
+			s.Equal(tt.want, output)
+		})
+	}
+}

pkg/executor/docker/models/types.go

@@ -39,6 +39,16 @@ func (c EngineSpec) Validate() error {
 				"must contain absolute path", c.WorkingDirectory)
 		}
 	}
+	// Validate environment variables
+	for _, env := range c.EnvironmentVariables {
+		parts := strings.SplitN(env, "=", 2)
+		if len(parts) > 0 {
+			if strings.HasPrefix(strings.ToUpper(parts[0]), models.EnvVarPrefix) {
+				return fmt.Errorf("invalid docker engine param: environment variable '%s' cannot start with %s",
+					parts[0], models.EnvVarPrefix)
+			}
+		}
+	}
 	return nil
 }