Optionally constrain to connected, approved compute nodes when selecting/ranking them

cmd/cli/serve/util.go

@@ -11,6 +11,7 @@ import (
 	"github.com/bacalhau-project/bacalhau/pkg/compute/store/boltdb"
 	"github.com/bacalhau-project/bacalhau/pkg/jobstore"
 	boltjobstore "github.com/bacalhau-project/bacalhau/pkg/jobstore/boltdb"
+	"github.com/bacalhau-project/bacalhau/pkg/models"
 	"github.com/bacalhau-project/bacalhau/pkg/util/idgen"
 	pkgerrors "github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
@@ -91,7 +92,8 @@ func GetRequesterConfig(ctx context.Context, createJobStore bool) (node.Requeste
 			return node.RequesterConfig{}, pkgerrors.Wrapf(err, "failed to create job store")
 		}
 	}
-	return node.NewRequesterConfigWith(node.RequesterConfigParams{
+
+	requesterConfig, err := node.NewRequesterConfigWith(node.RequesterConfigParams{
 		JobDefaults: transformer.JobDefaults{
 			ExecutionTimeout: time.Duration(cfg.JobDefaults.ExecutionTimeout),
 		},
@@ -120,6 +122,17 @@ func GetRequesterConfig(ctx context.Context, createJobStore bool) (node.Requeste
 		JobStore:                       jobStore,
 		DefaultPublisher:               cfg.DefaultPublisher,
 	})
+	if err != nil {
+		return node.RequesterConfig{}, err
+	}
+
+	if cfg.ManualNodeApproval {
+		requesterConfig.DefaultApprovalState = models.NodeApprovals.PENDING
+	} else {
+		requesterConfig.DefaultApprovalState = models.NodeApprovals.APPROVED
+	}
+
+	return requesterConfig, nil
 }
 
 func getNodeType() (requester, compute bool, err error) {

pkg/config/types/generated_constants.go

@@ -142,6 +142,7 @@ const NodeRequesterControlPlaneSettings = "Node.Requester.ControlPlaneSettings"
 const NodeRequesterControlPlaneSettingsHeartbeatCheckFrequency = "Node.Requester.ControlPlaneSettings.HeartbeatCheckFrequency"
 const NodeRequesterControlPlaneSettingsHeartbeatTopic = "Node.Requester.ControlPlaneSettings.HeartbeatTopic"
 const NodeRequesterControlPlaneSettingsNodeDisconnectedAfter = "Node.Requester.ControlPlaneSettings.NodeDisconnectedAfter"
+const NodeRequesterManualNodeApproval = "Node.Requester.ManualNodeApproval"
 const NodeBootstrapAddresses = "Node.BootstrapAddresses"
 const NodeDownloadURLRequestRetries = "Node.DownloadURLRequestRetries"
 const NodeDownloadURLRequestTimeout = "Node.DownloadURLRequestTimeout"

pkg/config/types/generated_viper_defaults.go

@@ -164,6 +164,7 @@ func SetDefaults(cfg BacalhauConfig, opts ...SetOption) {
 	p.Viper.SetDefault(NodeRequesterControlPlaneSettingsHeartbeatCheckFrequency, cfg.Node.Requester.ControlPlaneSettings.HeartbeatCheckFrequency.AsTimeDuration())
 	p.Viper.SetDefault(NodeRequesterControlPlaneSettingsHeartbeatTopic, cfg.Node.Requester.ControlPlaneSettings.HeartbeatTopic)
 	p.Viper.SetDefault(NodeRequesterControlPlaneSettingsNodeDisconnectedAfter, cfg.Node.Requester.ControlPlaneSettings.NodeDisconnectedAfter.AsTimeDuration())
+	p.Viper.SetDefault(NodeRequesterManualNodeApproval, cfg.Node.Requester.ManualNodeApproval)
 	p.Viper.SetDefault(NodeBootstrapAddresses, cfg.Node.BootstrapAddresses)
 	p.Viper.SetDefault(NodeDownloadURLRequestRetries, cfg.Node.DownloadURLRequestRetries)
 	p.Viper.SetDefault(NodeDownloadURLRequestTimeout, cfg.Node.DownloadURLRequestTimeout.AsTimeDuration())
@@ -361,6 +362,7 @@ func Set(cfg BacalhauConfig, opts ...SetOption) {
 	p.Viper.Set(NodeRequesterControlPlaneSettingsHeartbeatCheckFrequency, cfg.Node.Requester.ControlPlaneSettings.HeartbeatCheckFrequency.AsTimeDuration())
 	p.Viper.Set(NodeRequesterControlPlaneSettingsHeartbeatTopic, cfg.Node.Requester.ControlPlaneSettings.HeartbeatTopic)
 	p.Viper.Set(NodeRequesterControlPlaneSettingsNodeDisconnectedAfter, cfg.Node.Requester.ControlPlaneSettings.NodeDisconnectedAfter.AsTimeDuration())
+	p.Viper.Set(NodeRequesterManualNodeApproval, cfg.Node.Requester.ManualNodeApproval)
 	p.Viper.Set(NodeBootstrapAddresses, cfg.Node.BootstrapAddresses)
 	p.Viper.Set(NodeDownloadURLRequestRetries, cfg.Node.DownloadURLRequestRetries)
 	p.Viper.Set(NodeDownloadURLRequestTimeout, cfg.Node.DownloadURLRequestTimeout.AsTimeDuration())

pkg/config/types/requester.go

@@ -27,6 +27,11 @@ type RequesterConfig struct {
 	DefaultPublisher string            `yaml:"DefaultPublisher"`
 
 	ControlPlaneSettings RequesterControlPlaneConfig `yaml:"ControlPlaneSettings"`
+
+	// ManualNodeApproval is a flag that determines if nodes should be manually approved or not.
+	// By default, nodes are auto-approved to simplify upgrades, by setting this property to
+	// true, nodes will need to be manually approved before they are included in node selection.
+	ManualNodeApproval bool `yaml:"ManualNodeApproval"`
 }
 
 type EvaluationBrokerConfig struct {

pkg/devstack/devstack.go

@@ -306,6 +306,10 @@ func Setup(
 			}
 		}
 
+		// Set the default approval state from the config provided, either PENDING if the user has
+		// chosen manual approval, or the default otherwise.
+		nodeConfig.RequesterNodeConfig.DefaultApprovalState = stackConfig.RequesterConfig.DefaultApprovalState
+
 		// Create dedicated store paths for each node
 		err = setStorePaths(ctx, fsRepo, &nodeConfig)
 		if err != nil {

pkg/devstack/option.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/rs/zerolog"
 
+	"github.com/bacalhau-project/bacalhau/pkg/models"
 	"github.com/bacalhau-project/bacalhau/pkg/node"
 	"github.com/bacalhau-project/bacalhau/pkg/routing"
 )
@@ -121,6 +122,12 @@ func (o *DevStackConfig) Validate() error {
 	return errs
 }
 
+func WithAutoNodeApproval() ConfigOption {
+	return func(cfg *DevStackConfig) {
+		cfg.RequesterConfig.DefaultApprovalState = models.NodeApprovals.APPROVED
+	}
+}
+
 func WithNodeOverrides(overrides ...node.NodeConfig) ConfigOption {
 	return func(cfg *DevStackConfig) {
 		cfg.NodeOverrides = overrides

pkg/node/config_defaults.go

@@ -74,6 +74,8 @@ var DefaultRequesterConfig = RequesterConfigParams{
 		HeartbeatTopic:          "heartbeat",
 		NodeDisconnectedAfter:   types.Duration(30 * time.Second), //nolint:gomnd
 	},
+
+	DefaultApprovalState: models.NodeApprovals.APPROVED,
 }
 
 var TestRequesterConfig = RequesterConfigParams{
@@ -108,6 +110,8 @@ var TestRequesterConfig = RequesterConfigParams{
 		HeartbeatTopic:          "heartbeat",
 		NodeDisconnectedAfter:   types.Duration(30 * time.Second), //nolint:gomnd
 	},
+
+	DefaultApprovalState: models.NodeApprovals.APPROVED,
 }
 
 func getRequesterConfigParams() RequesterConfigParams {

pkg/node/config_requester.go

@@ -53,6 +53,11 @@ type RequesterConfigParams struct {
 
 	DefaultPublisher string
 
+	// When new nodes join the cluster, what state do they have? By default, APPROVED, and
+	// for tests, APPROVED. We will provide an option to set this to PENDING for production
+	// or for when operators are ready to control node approval.
+	DefaultApprovalState models.NodeApproval
+
 	ControlPlaneSettings types.RequesterControlPlaneConfig
 }
 

pkg/node/manager/node_manager.go

@@ -24,23 +24,26 @@ const (
 // also provides operations for querying and managing compute
 // node information.
 type NodeManager struct {
-	nodeInfo    routing.NodeInfoStore
-	resourceMap *concurrency.StripedMap[models.Resources]
-	heartbeats  *heartbeat.HeartbeatServer
+	nodeInfo             routing.NodeInfoStore
+	resourceMap          *concurrency.StripedMap[models.Resources]
+	heartbeats           *heartbeat.HeartbeatServer
+	defaultApprovalState models.NodeApproval
 }
 
 type NodeManagerParams struct {
-	NodeInfo   routing.NodeInfoStore
-	Heartbeats *heartbeat.HeartbeatServer
+	NodeInfo             routing.NodeInfoStore
+	Heartbeats           *heartbeat.HeartbeatServer
+	DefaultApprovalState models.NodeApproval
 }
 
 // NewNodeManager constructs a new node manager and returns a pointer
 // to the structure.
 func NewNodeManager(params NodeManagerParams) *NodeManager {
 	return &NodeManager{
-		resourceMap: concurrency.NewStripedMap[models.Resources](resourceMapLockCount),
-		nodeInfo:    params.NodeInfo,
-		heartbeats:  params.Heartbeats,
+		resourceMap:          concurrency.NewStripedMap[models.Resources](resourceMapLockCount),
+		nodeInfo:             params.NodeInfo,
+		heartbeats:           params.Heartbeats,
+		defaultApprovalState: params.DefaultApprovalState,
 	}
 }
 
@@ -83,9 +86,7 @@ func (n *NodeManager) Register(ctx context.Context, request requests.RegisterReq
 		}, nil
 	}
 
-	// TODO: We will default to PENDING, but once we start filtering on NodeApprovals.APPROVED we will need to
-	// make a decision on how this is determined.
-	request.Info.Approval = models.NodeApprovals.PENDING
+	request.Info.Approval = n.defaultApprovalState
 
 	if err := n.nodeInfo.Add(ctx, request.Info); err != nil {
 		return nil, errors.Wrap(err, "failed to save nodeinfo during node registration")

pkg/node/node.go

@@ -298,8 +298,9 @@ func NewNode(
 		// to the network. Provide it with a mechanism to lookup (and enhance)
 		// node info, and a reference to the heartbeat server if running NATS.
 		nodeManager := manager.NewNodeManager(manager.NodeManagerParams{
-			NodeInfo:   tracingInfoStore,
-			Heartbeats: heartbeatSvr,
+			NodeInfo:             tracingInfoStore,
+			Heartbeats:           heartbeatSvr,
+			DefaultApprovalState: config.RequesterNodeConfig.DefaultApprovalState,
 		})
 
 		// Start the nodemanager, ensuring it doesn't block the main thread and

pkg/node/requester.go

@@ -80,6 +80,10 @@ func NewRequesterNode(
 		}),
 	)
 
+	log.Ctx(ctx).
+		Info().
+		Msgf("Nodes joining the cluster will be assigned approval state: %s", requesterConfig.DefaultApprovalState.String())
+
 	// compute node ranker
 	nodeRankerChain := ranking.NewChain()
 	nodeRankerChain.Add(

pkg/orchestrator/interfaces.go

@@ -102,11 +102,13 @@ type NodeRanker interface {
 type NodeSelector interface {
 	// AllNodes returns all nodes in the network.
 	AllNodes(ctx context.Context) ([]models.NodeInfo, error)
+
 	// AllMatchingNodes returns all nodes that match the job constrains and selection criteria.
-	AllMatchingNodes(ctx context.Context, job *models.Job) ([]models.NodeInfo, error)
+	AllMatchingNodes(ctx context.Context, job *models.Job, constraints *NodeSelectionConstraints) ([]models.NodeInfo, error)
+
 	// TopMatchingNodes return the top ranked desiredCount number of nodes that match job constraints
 	// ordered in descending order based on their rank, or error if not enough nodes match.
-	TopMatchingNodes(ctx context.Context, job *models.Job, desiredCount int) ([]models.NodeInfo, error)
+	TopMatchingNodes(ctx context.Context, job *models.Job, desiredCount int, constraints *NodeSelectionConstraints) ([]models.NodeInfo, error)
 }
 
 type RetryStrategy interface {

pkg/orchestrator/scheduler/batch_job_test.go

@@ -303,10 +303,20 @@ func (s *BatchJobSchedulerTestSuite) TestProcess_ShouldMarkJobAsFailed_NoRetry()
 }
 
 func (s *BatchJobSchedulerTestSuite) mockNodeSelection(job *models.Job, nodeInfos []models.NodeInfo, desiredCount int) {
+	constraints := &orchestrator.NodeSelectionConstraints{
+		RequireApproval:  false,
+		RequireConnected: false,
+	}
+
 	if len(nodeInfos) < desiredCount {
-		s.nodeSelector.EXPECT().TopMatchingNodes(gomock.Any(), job, desiredCount).Return(nil, orchestrator.ErrNotEnoughNodes{})
+		s.nodeSelector.EXPECT().TopMatchingNodes(gomock.Any(), job, desiredCount, constraints).Return(nil, orchestrator.ErrNotEnoughNodes{})
 	} else {
-		s.nodeSelector.EXPECT().TopMatchingNodes(gomock.Any(), job, desiredCount).Return(nodeInfos, nil)
+		s.nodeSelector.EXPECT().TopMatchingNodes(
+			gomock.Any(),
+			job,
+			desiredCount,
+			constraints,
+		).Return(nodeInfos, nil)
 	}
 }
 

pkg/orchestrator/scheduler/batch_service_job.go

@@ -153,7 +153,16 @@ func (b *BatchServiceJobScheduler) createMissingExecs(
 // placeExecs places the executions
 func (b *BatchServiceJobScheduler) placeExecs(ctx context.Context, execs execSet, job *models.Job) error {
 	if len(execs) > 0 {
-		selectedNodes, err := b.nodeSelector.TopMatchingNodes(ctx, job, len(execs))
+		// TODO: Remove the options once we are ready to enforce that only connected/approved nodes can be used
+		selectedNodes, err := b.nodeSelector.TopMatchingNodes(
+			ctx,
+			job,
+			len(execs),
+			&orchestrator.NodeSelectionConstraints{
+				RequireApproval:  false,
+				RequireConnected: false,
+			},
+		)
 		if err != nil {
 			return err
 		}

pkg/orchestrator/scheduler/daemon_job.go

@@ -85,7 +85,13 @@ func (b *DaemonJobScheduler) Process(ctx context.Context, evaluation *models.Eva
 func (b *DaemonJobScheduler) createMissingExecs(
 	ctx context.Context, job *models.Job, plan *models.Plan, existingExecs execSet) (execSet, error) {
 	newExecs := execSet{}
-	nodes, err := b.nodeSelector.AllMatchingNodes(ctx, job)
+
+	// Require approval when selecting nodes, but do not require them to be connected.
+	nodes, err := b.nodeSelector.AllMatchingNodes(
+		ctx,
+		job,
+		&orchestrator.NodeSelectionConstraints{RequireApproval: true, RequireConnected: false},
+	)
 	if err != nil {
 		return newExecs, err
 	}

pkg/orchestrator/scheduler/daemon_job_test.go

@@ -52,7 +52,7 @@ func (s *DaemonJobSchedulerTestSuite) TestProcess_ShouldCreateNewExecutions() {
 		*mockNodeInfo(s.T(), nodeIDs[1]),
 		*mockNodeInfo(s.T(), nodeIDs[2]),
 	}
-	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), gomock.Any()).Return(nodeInfos, nil)
+	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), gomock.Any(), gomock.Any()).Return(nodeInfos, nil)
 
 	matcher := NewPlanMatcher(s.T(), PlanMatcherParams{
 		Evaluation:               evaluation,
@@ -77,7 +77,7 @@ func (s *DaemonJobSchedulerTestSuite) TestProcess_ShouldNOTMarkJobAsCompleted()
 	executions[1].ComputeState = models.NewExecutionState(models.ExecutionStateCompleted) // Simulate a completed execution
 	s.jobStore.EXPECT().GetJob(gomock.Any(), job.ID).Return(*job, nil)
 	s.jobStore.EXPECT().GetExecutions(gomock.Any(), jobstore.GetExecutionsOptions{JobID: job.ID}).Return(executions, nil)
-	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), gomock.Any()).Return([]models.NodeInfo{}, nil)
+	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), gomock.Any(), gomock.Any()).Return([]models.NodeInfo{}, nil)
 
 	// Noop plan
 	matcher := NewPlanMatcher(s.T(), PlanMatcherParams{
@@ -100,7 +100,7 @@ func (s *DaemonJobSchedulerTestSuite) TestProcess_ShouldMarkLostExecutionsOnUnhe
 		*mockNodeInfo(s.T(), executions[1].NodeID),
 	}
 	s.nodeSelector.EXPECT().AllNodes(gomock.Any()).Return(nodeInfos, nil)
-	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job).Return(nodeInfos, nil)
+	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job, gomock.Any()).Return(nodeInfos, nil)
 
 	matcher := NewPlanMatcher(s.T(), PlanMatcherParams{
 		Evaluation:         evaluation,
@@ -129,7 +129,7 @@ func (s *DaemonJobSchedulerTestSuite) TestProcess_ShouldNOTMarkJobAsFailed() {
 		*mockNodeInfo(s.T(), executions[1].NodeID),
 	}
 	s.nodeSelector.EXPECT().AllNodes(gomock.Any()).Return(nodeInfos, nil)
-	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job).Return(nodeInfos, nil)
+	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job, gomock.Any()).Return(nodeInfos, nil)
 
 	matcher := NewPlanMatcher(s.T(), PlanMatcherParams{
 		Evaluation: evaluation,
@@ -165,7 +165,7 @@ func (s *DaemonJobSchedulerTestSuite) TestProcessFail_NoMatchingNodes() {
 	executions := []models.Execution{} // no executions yet
 	s.jobStore.EXPECT().GetJob(gomock.Any(), job.ID).Return(*job, nil)
 	s.jobStore.EXPECT().GetExecutions(gomock.Any(), jobstore.GetExecutionsOptions{JobID: job.ID}).Return(executions, nil)
-	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job).Return([]models.NodeInfo{}, nil)
+	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job, gomock.Any()).Return([]models.NodeInfo{}, nil)
 
 	// Noop plan
 	matcher := NewPlanMatcher(s.T(), PlanMatcherParams{

pkg/orchestrator/scheduler/ops_job.go

@@ -103,7 +103,12 @@ func (b *OpsJobScheduler) Process(ctx context.Context, evaluation *models.Evalua
 func (b *OpsJobScheduler) createMissingExecs(
 	ctx context.Context, job *models.Job, plan *models.Plan) (execSet, error) {
 	newExecs := execSet{}
-	nodes, err := b.nodeSelector.AllMatchingNodes(ctx, job)
+	nodes, err := b.nodeSelector.AllMatchingNodes(
+		ctx,
+		job, &orchestrator.NodeSelectionConstraints{
+			RequireApproval:  true,
+			RequireConnected: false,
+		})
 	if err != nil {
 		return newExecs, err
 	}

pkg/orchestrator/scheduler/ops_job_test.go

@@ -167,7 +167,7 @@ func (s *OpsJobSchedulerTestSuite) TestProcessFail_NoMatchingNodes() {
 }
 
 func (s *OpsJobSchedulerTestSuite) mockNodeSelection(job *models.Job, nodeInfos []models.NodeInfo) {
-	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job).Return(nodeInfos, nil)
+	s.nodeSelector.EXPECT().AllMatchingNodes(gomock.Any(), job, gomock.Any()).Return(nodeInfos, nil)
 }
 
 func mockOpsJob() (*models.Job, []models.Execution, *models.Evaluation) {