iop/apps/edge/internal/service/run_submit.go

package service

import (
	"context"
	"fmt"

	edgenode "iop/apps/edge/internal/node"
	iop "iop/proto/gen/iop"
)

// runSubscription holds the request-bound event channels of a dispatched run.
// A background run has no subscriber, so every field stays nil and close is a
// no-op.
type runSubscription struct {
	runEvents      <-chan *iop.RunEvent
	nodeEvents     <-chan *iop.EdgeNodeEvent
	unregisterRun  func()
	unregisterNode func()
}

// close unregisters both subscriptions. It is safe on a background (empty)
// subscription.
func (sub *runSubscription) close() {
	if sub.unregisterRun != nil {
		sub.unregisterRun()
	}
	if sub.unregisterNode != nil {
		sub.unregisterNode()
	}
}

// subscribeRun opens the run and node event channels a foreground run needs.
// It runs before the send so events emitted immediately after dispatch are not
// missed.
func (s *Service) subscribeRun(runID, nodeID string, background bool) (*runSubscription, error) {
	sub := &runSubscription{}
	if background {
		return sub, nil
	}
	if s.events == nil {
		return nil, fmt.Errorf("event bus is not configured")
	}
	sub.runEvents, sub.unregisterRun = s.events.SubscribeRun(runID, 4096)
	sub.nodeEvents, sub.unregisterNode = s.events.SubscribeNode(nodeID, 16)
	return sub, nil
}

func (s *Service) SubmitRun(ctx context.Context, req SubmitRunRequest) (RunResult, error) {
	if req.ProviderPool && req.ModelGroupKey != "" && s.queue != nil {
		return s.submitRunQueued(ctx, req)
	}
	return s.submitRunDirect(req)
}

func (s *Service) submitRunDirect(req SubmitRunRequest) (RunResult, error) {
	entry, err := s.ResolveDispatchReady(req.NodeRef)
	if err != nil {
		return nil, err
	}
	return s.dispatchToEntry(entry, req)
}

func (s *Service) submitRunQueued(ctx context.Context, req SubmitRunRequest) (RunResult, error) {
	candidates, returnedPolicy, err := s.resolveQueueCandidates(req)
	if err != nil {
		return nil, err
	}

	long := req.ContextClass == contextClassLong

	// For provider-pool requests the canonical policy is owned by the atomic
	// runtime snapshot, not by the resolution path. Legacy paths use the
	// policy derived from the request or store.
	providerPool := req.ProviderPool
	var policy groupPolicy
	if providerPool {
		_, _, policy = s.runtimeConfigSnapshot()
	} else {
		policy = returnedPolicy
	}

	selected, queueReason, err := s.queue.admitWithReason(ctx, req.ModelGroupKey, req.Adapter, req.Target, candidates, policy, s.resolveQueueCandidatesClosure(req), long, providerPool)
	if err != nil {
		return nil, err
	}

	// The admitted slot is owned by one reservation from here on: every failure
	// path below releases through it, and a dispatched run hands it off to the
	// event watcher.
	reservation := newQueueReservation(s.queue, selected)

	// Rewrite adapter and target for provider-pool dispatch: the winning candidate
	// carries the concrete adapter and served model name determined at selection time.
	if selected.adapter != "" {
		req.Adapter = selected.adapter
	}
	if selected.servedTarget != "" {
		req.Target = selected.servedTarget
	}

	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		reservation.release("build-error")
		return nil, err
	}

	// Track inflight before send so the event watcher can release the slot
	// even if a terminal event arrives before the Send call completes.
	reservation.track(runID)

	// Pre-send generation fence: if the selected connection was superseded by a
	// disconnect/reconnect since admission, release the lease and fail instead of
	// dispatching to a dead client.
	if !s.candidateIsCurrentOwner(selected) {
		reservation.release("stale-generation")
		return nil, staleGenerationError(selected)
	}

	sub, err := s.subscribeRun(runID, selected.entry.NodeID, runReq.GetBackground())
	if err != nil {
		reservation.release("no-event-bus")
		return nil, err
	}

	if err := selected.entry.Client.Send(runReq); err != nil {
		reservation.release("send-error")
		sub.close()
		return nil, err
	}
	// The run now owns the slot: the event watcher releases it on the terminal
	// run event or on node disconnect.
	reservation.handOff()

	return newRunHandle(RunDispatch{
		RunID:                runID,
		NodeID:               selected.entry.NodeID,
		NodeLabel:            nodeLabel(selected.entry),
		ModelGroupKey:        req.ModelGroupKey,
		Adapter:              runReq.GetAdapter(),
		Target:               runReq.GetTarget(),
		SessionID:            runReq.GetSessionId(),
		Background:           runReq.GetBackground(),
		TimeoutSec:           int(runReq.GetTimeoutSec()),
		EstimatedInputTokens: req.EstimatedInputTokens,
		ContextClass:         req.ContextClass,
		ProviderID:           selected.providerID,
		ProviderType:         selected.providerType,
		ExecutionPath:        string(selected.executionPath),
		QueueReason:          queueReason,
	}, sub), nil
}

// newRunHandle wraps a dispatched run's identity and event subscription into
// the surface-neutral handle callers read.
func newRunHandle(disp RunDispatch, sub *runSubscription) *RunHandle {
	return &RunHandle{
		RunDispatch: disp,
		RunStream: RunStream{
			Events:     sub.runEvents,
			NodeEvents: sub.nodeEvents,
		},
		close: sub.close,
	}
}

func (s *Service) dispatchToEntry(entry *edgenode.NodeEntry, req SubmitRunRequest) (RunResult, error) {
	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		return nil, err
	}

	sub, err := s.subscribeRun(runID, entry.NodeID, runReq.GetBackground())
	if err != nil {
		return nil, err
	}

	if err := entry.Client.Send(runReq); err != nil {
		sub.close()
		return nil, err
	}

	return newRunHandle(RunDispatch{
		RunID:                runID,
		NodeID:               entry.NodeID,
		NodeLabel:            nodeLabel(entry),
		ModelGroupKey:        req.ModelGroupKey,
		Adapter:              runReq.GetAdapter(),
		Target:               runReq.GetTarget(),
		SessionID:            runReq.GetSessionId(),
		Background:           runReq.GetBackground(),
		TimeoutSec:           int(runReq.GetTimeoutSec()),
		EstimatedInputTokens: req.EstimatedInputTokens,
		ContextClass:         req.ContextClass,
		QueueReason:          "dispatched",
	}, sub), nil
}