Skip to content
This repository has been archived by the owner on Oct 9, 2023. It is now read-only.

Populate TaskNodeMetadata on CreateNodeExecutionModel to correctly display fast cache #483

Merged
merged 3 commits into from
Mar 29, 2023
Merged

Populate TaskNodeMetadata on CreateNodeExecutionModel to correctly display fast cache #483

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6a8ffa8
adding cache information on first node event
hamersaw Oct 10, 2022
6a27535
fixed unit tests
hamersaw Oct 10, 2022
049e37c
Merge branch 'master' into feature/fast-cache
hamersaw Mar 29, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions pkg/manager/impl/node_execution_manager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,9 @@ func TestCreateNodeEvent(t *testing.T) {
StartedAt: occurredAtProto,
CreatedAt: occurredAtProto,
UpdatedAt: occurredAtProto,
TargetMetadata: &admin.NodeExecutionClosure_TaskNodeMetadata{
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is this change necessary? we shouldn't be initializing the metadata in the closure unless it's sourced from the event itself, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So I'm not sure I follow. Currently the TaskNodeMetadata from a node event is only used if when we UpdateNodeExecutionModel here. In the CreateNodeExecutionModel func (which is called on the first node node event, update on all subsequent) we just drop any information in the TaskNodeMetadata. This PR adds support for processing the TaskNodeMetadata in the CreateNodeExecutionModel function - which creates the initial model based on the first event. This change results in setting the TaskNodeMetadata in the CreateNodeEvent function. So I copied the TargetMetadata from the unit tests for updating the node model. Does this make sense? Or am I thinking about this wrong?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah sorry, I didn't realize the event under test actually had a non-empty TargetMetadata - and we should therefore expect to see the closure here have an initialized TargetMetadata

TaskNodeMetadata: &admin.TaskNodeMetadata{},
},
}
closureBytes, _ := proto.Marshal(&expectedClosure)
repository.NodeExecutionRepo().(*repositoryMocks.MockNodeExecutionRepo).SetCreateCallback(
Expand Down
18 changes: 18 additions & 0 deletions pkg/repositories/transformers/node_execution.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,12 +143,30 @@ func CreateNodeExecutionModel(ctx context.Context, input ToNodeExecutionModelInp
return nil, err
}
}

if common.IsNodeExecutionTerminal(input.Request.Event.Phase) {
err := addTerminalState(ctx, input.Request, nodeExecution, &closure, input.InlineEventDataPolicy, input.StorageClient)
if err != nil {
return nil, err
}
}

// Update TaskNodeMetadata, which includes caching information today.
if input.Request.Event.GetTaskNodeMetadata() != nil {
targetMetadata := &admin.NodeExecutionClosure_TaskNodeMetadata{
TaskNodeMetadata: &admin.TaskNodeMetadata{
CheckpointUri: input.Request.Event.GetTaskNodeMetadata().CheckpointUri,
},
}
if input.Request.Event.GetTaskNodeMetadata().CatalogKey != nil {
st := input.Request.Event.GetTaskNodeMetadata().GetCacheStatus().String()
targetMetadata.TaskNodeMetadata.CacheStatus = input.Request.Event.GetTaskNodeMetadata().GetCacheStatus()
targetMetadata.TaskNodeMetadata.CatalogKey = input.Request.Event.GetTaskNodeMetadata().GetCatalogKey()
nodeExecution.CacheStatus = &st
}
closure.TargetMetadata = targetMetadata
}

marshaledClosure, err := proto.Marshal(&closure)
if err != nil {
return nil, errors.NewFlyteAdminErrorf(
Expand Down
73 changes: 49 additions & 24 deletions pkg/repositories/transformers/node_execution_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,35 +199,51 @@ func TestAddTerminalState_Error(t *testing.T) {

func TestCreateNodeExecutionModel(t *testing.T) {
parentTaskExecID := uint(8)
nodeExecutionModel, err := CreateNodeExecutionModel(context.TODO(), ToNodeExecutionModelInput{
Request: &admin.NodeExecutionEventRequest{
Event: &event.NodeExecutionEvent{
Id: &core.NodeExecutionIdentifier{
NodeId: "node id",
ExecutionId: &core.WorkflowExecutionIdentifier{
Project: "project",
Domain: "domain",
Name: "name",
},
},
Phase: core.NodeExecution_RUNNING,
InputValue: &event.NodeExecutionEvent_InputUri{
InputUri: testInputURI,
},
OutputResult: &event.NodeExecutionEvent_OutputUri{
OutputUri: "output uri",
request := &admin.NodeExecutionEventRequest{
Event: &event.NodeExecutionEvent{
Id: &core.NodeExecutionIdentifier{
NodeId: "node id",
ExecutionId: &core.WorkflowExecutionIdentifier{
Project: "project",
Domain: "domain",
Name: "name",
},
OccurredAt: occurredAtProto,
ParentTaskMetadata: &event.ParentTaskExecutionMetadata{
Id: &core.TaskExecutionIdentifier{
RetryAttempt: 1,
},
Phase: core.NodeExecution_RUNNING,
InputValue: &event.NodeExecutionEvent_InputUri{
InputUri: testInputURI,
},
OutputResult: &event.NodeExecutionEvent_OutputUri{
OutputUri: "output uri",
},
OccurredAt: occurredAtProto,
TargetMetadata: &event.NodeExecutionEvent_TaskNodeMetadata{
TaskNodeMetadata: &event.TaskNodeMetadata{
CacheStatus: core.CatalogCacheStatus_CACHE_POPULATED,
CatalogKey: &core.CatalogMetadata{
DatasetId: &core.Identifier{
ResourceType: core.ResourceType_DATASET,
Name: "x",
Project: "proj",
Domain: "domain",
},
},
CheckpointUri: "last checkpoint uri",
},
IsParent: true,
IsDynamic: true,
EventVersion: 2,
},
ParentTaskMetadata: &event.ParentTaskExecutionMetadata{
Id: &core.TaskExecutionIdentifier{
RetryAttempt: 1,
},
},
IsParent: true,
IsDynamic: true,
EventVersion: 2,
},
}

nodeExecutionModel, err := CreateNodeExecutionModel(context.TODO(), ToNodeExecutionModelInput{
Request: request,
ParentTaskExecutionID: &parentTaskExecID,
})
assert.Nil(t, err)
Expand All @@ -237,6 +253,13 @@ func TestCreateNodeExecutionModel(t *testing.T) {
StartedAt: occurredAtProto,
CreatedAt: occurredAtProto,
UpdatedAt: occurredAtProto,
TargetMetadata: &admin.NodeExecutionClosure_TaskNodeMetadata{
TaskNodeMetadata: &admin.TaskNodeMetadata{
CacheStatus: request.Event.GetTaskNodeMetadata().CacheStatus,
CatalogKey: request.Event.GetTaskNodeMetadata().CatalogKey,
CheckpointUri: request.Event.GetTaskNodeMetadata().CheckpointUri,
},
},
}
var closureBytes, _ = proto.Marshal(closure)
var nodeExecutionMetadata, _ = proto.Marshal(&admin.NodeExecutionMetaData{
Expand All @@ -247,6 +270,7 @@ func TestCreateNodeExecutionModel(t *testing.T) {
EventVersion: 2,
}
internalDataBytes, _ := proto.Marshal(internalData)
cacheStatus := request.Event.GetTaskNodeMetadata().CacheStatus.String()
assert.Equal(t, &models.NodeExecution{
NodeExecutionKey: models.NodeExecutionKey{
NodeID: "node id",
Expand All @@ -264,6 +288,7 @@ func TestCreateNodeExecutionModel(t *testing.T) {
NodeExecutionUpdatedAt: &occurredAt,
NodeExecutionMetadata: nodeExecutionMetadata,
ParentTaskExecutionID: &parentTaskExecID,
CacheStatus: &cacheStatus,
InternalData: internalDataBytes,
}, nodeExecutionModel)
}
Expand Down