Real-time Monitoring and Log Streaming

Flyte provides robust mechanisms for real-time monitoring and log streaming, crucial for observing the execution of runs and actions. This functionality is primarily driven by the RunLogsService, which orchestrates log retrieval, and the actionRepo, which provides real-time updates on run and action states through PostgreSQL's LISTEN/NOTIFY feature.

Orchestrating Log Retrieval with RunLogsService

When you need to tail logs for an action attempt, the RunLogsService (runs.service.RunLogsService) is the central component that handles these requests. It acts as an intermediary, fetching the necessary log context from the ActionRepo and then delegating the actual log streaming to a LogStreamer implementation.

The RunLogsService is defined as:

// RunLogsService implements the RunLogsServiceHandler interface.
type RunLogsService struct {
	repo     interfaces.Repository
	streamer LogStreamer
	sem      *semaphore.Weighted
}

Its primary method for log tailing is TailLogs:

func (s *RunLogsService) TailLogs(ctx context.Context, req *connect.Request[workflow.TailLogsRequest], stream *connect.ServerStream[workflow.TailLogsResponse]) error

This service is initialized in runs/setup.go with an ActionRepo and a LogStreamer (typically a K8sLogStreamer). It also employs a semaphore (sem) to limit the number of concurrent log streams, preventing resource exhaustion. If the concurrency limit is exceeded, the service will return a CodeResourceExhausted error.

An example of how RunLogsService processes a TailLogs request can be seen in its test suite. It first retrieves the latest event for a given action attempt from the ActionRepo to get the LogContext, then passes this context to the LogStreamer:

func TestTailLogs_HappyPath(t *testing.T) {
	actionRepo := &repoMocks.ActionRepo{}
	streamer := &mockLogStreamer{}

	logCtx := &core.LogContext{
		PrimaryPodName: "my-pod",
		Pods: []*core.PodLogContext{
			{PodName: "my-pod", Namespace: "ns"},
		},
	}

	eventModel := makeEventWithLogContext(tailLogsActionID, 1, common.ActionPhase_ACTION_PHASE_RUNNING, logCtx)
	actionRepo.On("GetLatestEventByAttempt", mock.Anything, mock.Anything, uint32(1)).Return(eventModel, nil)

	// The streamer should be called with the logContext and send some response.
	streamer.On("TailLogs", mock.Anything, mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
		stream := args.Get(2).(*connect.ServerStream[workflow.TailLogsResponse])
		_ = stream.Send(&workflow.TailLogsResponse{
			Logs: []*workflow.TailLogsResponse_Logs{
				{
					Lines: []*dataplane.LogLine{
						{Message: "hello world", Originator: dataplane.LogLineOriginator_USER},
					},
				},
			},
		})
	}).Return(nil)

	client := newTailLogsTestClient(t, actionRepo, streamer)

	stream, err := client.TailLogs(context.Background(), connect.NewRequest(&workflow.TailLogsRequest{
		ActionId: tailLogsActionID,
		Attempt:  1,
	}))
	assert.NoError(t, err)

	assert.True(t, stream.Receive())
	resp := stream.Msg()
	assert.Len(t, resp.Logs, 1)
	assert.Len(t, resp.Logs[0].Lines, 1)
	assert.Equal(t, "hello world", resp.Logs[0].Lines[0].Message)

	assert.False(t, stream.Receive())
	assert.NoError(t, stream.Err())

	actionRepo.AssertExpectations(t)
	streamer.AssertExpectations(t)
}

Abstracting Log Sources with LogStreamer and K8sLogStreamer

The LogStreamer (runs.service.LogStreamer) is an interface that abstracts the underlying mechanism for fetching logs. This design allows Flyte to support various log backends without modifying the core RunLogsService logic.

// LogStreamer abstracts log fetching from different backends.
type LogStreamer interface {
	TailLogs(ctx context.Context, logContext *core.LogContext, stream *connect.ServerStream[workflow.TailLogsResponse]) error
}

The primary concrete implementation used in Flyte is K8sLogStreamer (runs.service.K8sLogStreamer), which is responsible for streaming logs directly from Kubernetes pods. This implementation leverages the Kubernetes API to access pod logs.

// K8sLogStreamer streams logs directly from Kubernetes pods.
type K8sLogStreamer struct {
	clientset kubernetes.Interface
}

The K8sLogStreamer is initialized with a Kubernetes clientset in runs/setup.go and dataproxy/setup.go. Its TailLogs method handles the specifics of interacting with the Kubernetes API, setting up log options such as Follow, Timestamps, TailLines, and SinceTime. Notably, Follow is set to true only when the pod is actively running; for pending or terminated pods, Follow is disabled to retrieve existing logs immediately. To ensure long-lived log streams are not prematurely terminated, K8sLogStreamer also clears the Kubernetes REST config timeout.

Real-time State Updates with ActionRepo

The ActionRepo interface (runs.repository.interfaces.ActionRepo) defines the contract for accessing and managing action and run data. Its concrete implementation, actionRepo (runs.repository.impl.actionRepo), plays a critical role in providing real-time updates on the state of runs and actions. This is achieved through the use of PostgreSQL's LISTEN/NOTIFY mechanism.

The actionRepo implementation is structured to manage subscribers for these real-time updates:

// actionRepo implements actionRepo interface using PostgreSQL
type actionRepo struct {
	db       *sqlx.DB
	dsn      string // stored for pq.Listener
	listener *pq.Listener

	// Subscriber management for LISTEN/NOTIFY
	runSubscribers    map[chan string]bool
	actionSubscribers map[chan string]bool
	mu                sync.RWMutex

	// Dedicated channels for async NOTIFY to avoid pool contention
	actionNotifyCh chan string
	runNotifyCh    chan string
}

Key to its real-time capabilities are methods like WatchActionUpdates and WatchRunUpdates, which allow components to subscribe to changes for specific actions or runs. These methods leverage the pq.Listener to receive notifications from PostgreSQL whenever relevant data changes. This ensures that services like RunLogsService can react promptly to state transitions.

For instance, to watch for updates on a specific action, you would use the WatchActionUpdates method:

func (r *actionRepo) WatchActionUpdates(ctx context.Context, actionID *common.ActionIdentifier, updates chan<- *models.Action, errs chan<- error)

An example from the action_test.go demonstrates how WatchActionUpdates ensures that only updates relevant to the targeted action are streamed:

func TestWatchActionUpdates_OnlyStreamsTargetAction(t *testing.T) {
	db := setupActionDB(t)
	defer func() { db.Exec("DELETE FROM actions") }()
	repo, err := NewActionRepo(db, testDbConfig)
	require.NoError(t, err)
	repoImpl := repo.(*actionRepo)

	runID := &common.RunIdentifier{
		Org:     "org1",
		Project: "proj1",
		Domain:  "domain1",
		Name:    "run1",
	}
	targetActionID := &common.ActionIdentifier{Run: runID, Name: "target"}
	otherActionID := &common.ActionIdentifier{Run: runID, Name: "other"}

	ctx := context.Background()

	// Start watcher before creating actions so we can deterministically
	// drain the creation notification and avoid a race where the async
	// NOTIFY arrives after the subscriber registers.
	watchCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	updates := make(chan *models.Action, 2)
	errs := make(chan error, 1)
	go repo.WatchActionUpdates(watchCtx, targetActionID, updates, errs)

	require.Eventually(t, func() bool {
		repoImpl.mu.RLock()
		defer repoImpl.mu.RUnlock()
		return len(repoImpl.actionSubscribers) > 0
	}, 2*time.Second, 10*time.Millisecond, "timed out waiting for watcher registration")

	_, err = repo.CreateAction(ctx, models.NewActionModel(targetActionID), false)
	require.NoError(t, err)
	_, err = repo.CreateAction(ctx, models.NewActionModel(otherActionID), false)
	require.NoError(t, err)

	// Drain the creation notification for the target action.
	select {
	case <-updates:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for creation notification")
	}

	// Update "other" — should NOT produce an update for "target".
	err = repo.UpdateActionPhase(ctx, otherActionID, common.ActionPhase_ACTION_PHASE_RUNNING, 1, core.CatalogCacheStatus_CACHE_DISABLED, nil)
	require.NoError(t, err)

	select {
	case action := <-updates:
		t.Fatalf("unexpected update for action %s", action.Name)
	case err := <-errs:
		require.NoError(t, err)
	case <-time.After(1200 * time.Millisecond):
	}

	// Update "target" — should produce an update.
	err = repo.UpdateActionPhase(ctx, targetActionID, common.ActionPhase_ACTION_PHASE_RUNNING, 1, core.CatalogCacheStatus_CACHE_DISABLED, nil)
	require.NoError(t, err)

	select {
	case action := <-updates:
		require.Equal(t, targetActionID.Name, action.Name)
	case err := <-errs:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for target action update")
	}
}

It's important to note that the UpdateActionPhase method in actionRepo allows specific phase transitions, such as moving from FAILED or TIMED_OUT to QUEUED for retries. However, it explicitly blocks backward transitions from other non-retryable phases (e.g., RUNNING to QUEUED) to maintain state integrity. Additionally, when an action transitions to ABORTED, the actionRepo does not insert a synthetic event row; this event is expected to be emitted by the controller responsible for the abortion process. This design ensures that the database accurately reflects the state as managed by the broader Flyte system.