[V2] Run Service - Watch action cache by machichima · Pull Request #7017 · flyteorg/flyte

machichima · 2026-03-12T02:24:59Z

Tracking issue

Related to: #6971

Why are the changes needed?

We didn't return the children_phase_counts from WatchAction. This PR is to count and return this info efficiently.

What changes were proposed in this pull request?

Add runs/service/run_state_manager.go as an in memory cache for WatchAction, to store the node tree info

How was this patch tested?

Test locally with sandbox and testing script runs/test/scripts/watch_actions.sh

Labels

Please add one or more of the following labels to categorize your PR:

added: For new features.
changed: For changes in existing functionality.
deprecated: For soon-to-be-removed features.
removed: For features being removed.
fixed: For any bug fixed.
security: In case of vulnerabilities

This is important to improve the readability of release notes.

Setup process

Screenshots

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Related PRs

Docs link

main
- Flyte 2 WIP #6583
  - [V2] Run Service - Watch action cache #7017 👈

Signed-off-by: machichima <nary12321@gmail.com>

Copilot

Pull request overview

Introduces server-side state tracking for WatchActions so the service can stream filtered, tree-aware action updates (including child phase aggregates), and updates persistence/querying to associate actions with a run via a run_name column.

Changes:

Add a per-watch runStateManager to compute action-tree visibility under filters and maintain per-node child phase counts.
Update WatchActions to send an initial filtered/tree-aware snapshot and then stream incremental changes using the state manager.
Add run_name to the actions model and use it in action creation and listing queries; add a helper script for calling WatchActions.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
runs/test/scripts/watch_actions.sh	Adds a `buf curl` helper script to call `RunService/WatchActions`.
runs/service/run_state_manager.go	New in-memory tree/state manager for filtered visibility and child-phase aggregation.
runs/service/run_state_manager_test.go	Unit tests validating phase-count aggregation, filter visibility behavior, and parent-missing errors.
runs/service/run_service.go	Refactors `WatchActions` to use the new state manager; adds helpers for initial listing + streaming updates.
runs/repository/models/action.go	Adds `RunName` column and a `Clone()` helper; updates `GetRunName` to prefer stored `RunName`.
runs/repository/impl/project.go	Changes project creation to use `ON CONFLICT DO NOTHING`.
runs/repository/impl/action.go	Populates `RunName` on inserts and changes `ListActions` to filter by `run_name`.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-12T06:03:14Z

runs/repository/impl/project.go

+	result := r.db.WithContext(ctx).
+		Clauses(clause.OnConflict{DoNothing: true}).
+		Create(project)


CreateProject now uses ON CONFLICT DO NOTHING, which means a duplicate insert will not return an error (RowsAffected=0). As a result, this method will no longer return ErrProjectAlreadyExists for duplicates, and the isDuplicateProjectError branch becomes ineffective. Consider explicitly checking result.RowsAffected == 0 and returning interfaces.ErrProjectAlreadyExists (or reverting to the prior behavior) to preserve the API contract and existing tests.

Copilot · 2026-03-12T06:03:15Z

runs/repository/impl/action.go

 	query := r.db.WithContext(ctx).Model(&models.Action{}).
-		Where("org = ? AND project = ? AND domain = ?",
-							runID.Org, runID.Project, runID.Domain).
-		Where("parent_action_name IS NOT NULL") // Exclude the root action/run itself
+		Where("org = ? AND project = ? AND domain = ? AND run_name = ?",
+			runID.Org, runID.Project, runID.Domain, runID.Name)



ListActions now filters by run_name = runID.Name, but existing rows created before this column was populated will likely have run_name set to the default empty string after AutoMigrate. That will cause ListActions (and WatchActions initial load) to return an empty set for pre-existing runs/actions unless you add a backfill migration/upgrade path to populate run_name for existing records.

Copilot · 2026-03-12T06:03:15Z

runs/repository/impl/action.go

-		Where("org = ? AND project = ? AND domain = ?",
-							runID.Org, runID.Project, runID.Domain).
-		Where("parent_action_name IS NOT NULL") // Exclude the root action/run itself
+		Where("org = ? AND project = ? AND domain = ? AND run_name = ?",


This query no longer excludes the root action/run itself (previously parent_action_name IS NOT NULL). If ListActions is intended to return only non-root actions (as the removed comment indicated), this will change the ListActions API behavior and may break clients that assume the run action is not included. If WatchActions needs the root for tree rendering, consider keeping ListActions semantics by reintroducing the exclusion here and adding a separate repo method/flag for WatchActions to include the root when needed.

Suggested change

Where("org = ? AND project = ? AND domain = ? AND run_name = ?",

Where("org = ? AND project = ? AND domain = ? AND run_name = ? AND parent_action_name IS NOT NULL",

Copilot · 2026-03-12T06:03:15Z

runs/service/run_service.go

 		case updated, ok := <-updatesCh:
 			if !ok {
 				return nil
 			}
-			if err := stream.Send(&workflow.WatchActionsResponse{
-				EnrichedActions: []*workflow.EnrichedAction{s.convertActionToEnrichedProto(updated)},
-			}); err != nil {
+			updates, err := rsm.upsertActions(ctx, []*models.Action{updated})
+			if err != nil {
+				return err
+			}
+			if err := s.sendChangedActions(runID, updates, stream); err != nil {


WatchActions now assumes that every updated action received from WatchActionUpdates belongs to the requested run, and it constructs the outgoing ActionIdentifier using the request runID. However, the SQLite polling implementation of WatchActionUpdates currently only filters by org/project/domain (not run), so updates from other runs in the same org/project/domain can leak into this stream. That can produce incorrect action IDs on the wire and may also trigger runStateManager errors (e.g., missing parents). Consider adding a server-side guard here (drop updates whose run doesn't match) and/or updating the repository watcher to include run_name = runID.Name (now that the column exists).

Copilot · 2026-03-12T06:03:15Z

runs/repository/models/action.go

 // GetRunName extracts the run name from the action
 // For root actions (runs), returns the action's own name
 // For child actions, extracts from ActionSpec JSON
 func (a *Action) GetRunName() string {
+	if a.RunName != "" {
+		return a.RunName
+	}
+
 	if a.ParentActionName == nil {
 		// Root action - the run name is the action name
 		return a.Name
 	}
-
-	// TODO: Extract run name from ActionSpec JSON
-	// For now, return empty string as placeholder
 	return ""
 }


GetRunName now returns an empty string for child actions when RunName isn't populated. With the repository now relying on run_name for ListActions filtering and convertActionToEnrichedProto using GetRunName to populate ActionIdentifier.Run.Name, legacy records (or any code path that forgets to set RunName) will produce actions that can't be associated with a run and may be impossible to query/stream correctly. Consider ensuring RunName is always set for all actions (including backfill for existing DB rows) or implementing a fallback extraction for child actions if that’s still required.

pingsutw · 2026-03-12T06:40:58Z

runs/service/run_service.go

+
+	enriched := make([]*workflow.EnrichedAction, 0, len(updates))
+	for _, update := range updates {
+		enriched = append(enriched, s.convertNodeUpdateToEnrichedProto(runID, update))


Should we check if s.convertNodeUpdateToEnrichedProto() is nil

Updated in f2b2858

pingsutw · 2026-03-12T06:44:25Z

runs/repository/impl/project.go


-	result := r.db.WithContext(ctx).Create(project)
+	result := r.db.WithContext(ctx).
+		Clauses(clause.OnConflict{DoNothing: true}).


Shouldn't we return an error on conflict?

This is to prevent duplicated key not allowed logs when start up (adding seed/default projects)

2026/03/13 09:15:49 /Users/nary/workData/open-source/flyte-v2/runs/repository/impl/project.go:32 duplicated key not allowed [0.433ms] [rows:0] INSERT INTO `projects` (`identifier`,`name`,`description`,`labels`,`state`,`created_at`,`updated_at`) VALUES ("flytesnacks","flytesnacks","","",0,"2026-03-13 01:15:49.009","2026-03-13 01:15:49.009")

Fixed in 1b1ff43 to only raise error and log in CreateProject endpoint

Signed-off-by: machichima <nary12321@gmail.com>

* feat: simple state manager to store actions Signed-off-by: machichima <nary12321@gmail.com> * feat: run state manager record node tree Signed-off-by: machichima <nary12321@gmail.com> * fix: also include root action in ListActions Signed-off-by: machichima <nary12321@gmail.com> * fix: prevent error when create default project on startup Signed-off-by: machichima <nary12321@gmail.com> * feat: add watch_actions.sh test script Signed-off-by: machichima <nary12321@gmail.com> * fix: continue for nil enrichedAction for defense Signed-off-by: machichima <nary12321@gmail.com> * fix: raise error in CreateProject when already exists Signed-off-by: machichima <nary12321@gmail.com> --------- Signed-off-by: machichima <nary12321@gmail.com>

machichima added 6 commits March 11, 2026 18:00

feat: simple state manager to store actions

5206594

Signed-off-by: machichima <nary12321@gmail.com>

feat: run state manager record node tree

a4e8f5b

Signed-off-by: machichima <nary12321@gmail.com>

Merge branch 'v2' of github.com:flyteorg/flyte into watch-action-cache

2506550

Signed-off-by: machichima <nary12321@gmail.com>

fix: also include root action in ListActions

cff91b0

Signed-off-by: machichima <nary12321@gmail.com>

fix: prevent error when create default project on startup

b24e51f

Signed-off-by: machichima <nary12321@gmail.com>

feat: add watch_actions.sh test script

a186a37

Signed-off-by: machichima <nary12321@gmail.com>

github-actions bot mentioned this pull request Mar 12, 2026

Flyte 2 WIP #6583

Draft

3 tasks

pingsutw requested a review from Copilot March 12, 2026 05:56

pingsutw assigned machichima Mar 12, 2026

pingsutw added the flyte2 label Mar 12, 2026

Copilot started reviewing on behalf of pingsutw March 12, 2026 05:57 View session

Copilot AI reviewed Mar 12, 2026

View reviewed changes

pingsutw reviewed Mar 12, 2026

View reviewed changes

pingsutw added this to the V2 GA milestone Mar 12, 2026

machichima added 2 commits March 13, 2026 08:54

fix: continue for nil enrichedAction for defense

f2b2858

Signed-off-by: machichima <nary12321@gmail.com>

fix: raise error in CreateProject when already exists

1b1ff43

Signed-off-by: machichima <nary12321@gmail.com>

pingsutw approved these changes Mar 13, 2026

View reviewed changes

machichima merged commit 8d586bd into flyteorg:v2 Mar 13, 2026
13 of 14 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[V2] Run Service - Watch action cache#7017

[V2] Run Service - Watch action cache#7017
machichima merged 8 commits intoflyteorg:v2from
machichima:watch-action-cache

machichima commented Mar 12, 2026 •

edited by github-actions bot

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

pingsutw Mar 12, 2026

Uh oh!

machichima Mar 13, 2026

Uh oh!

pingsutw Mar 12, 2026

Uh oh!

machichima Mar 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	Where("org = ? AND project = ? AND domain = ? AND run_name = ?",
	Where("org = ? AND project = ? AND domain = ? AND run_name = ? AND parent_action_name IS NOT NULL",

Conversation

machichima commented Mar 12, 2026 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Labels

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

pingsutw Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

machichima Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

pingsutw Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

machichima Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

machichima commented Mar 12, 2026 •

edited by github-actions bot

Loading