From 35de5b8d1d9a57287d3bf17f9c1db99022209ce6 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 17 Jan 2024 17:06:15 -0800 Subject: [PATCH] Hydrate failure node (#456) Signed-off-by: Kevin Su --- cmd/register/files_test.go | 25 +++++++++++++++++++++++++ cmd/register/register.go | 4 ++-- cmd/register/register_util.go | 10 ++++++++++ cmd/register/testdata/failure-node.tgz | Bin 0 -> 1671 bytes 4 files changed, 37 insertions(+), 2 deletions(-) create mode 100644 cmd/register/testdata/failure-node.tgz diff --git a/cmd/register/files_test.go b/cmd/register/files_test.go index 0eec77014b..129d4f49c1 100644 --- a/cmd/register/files_test.go +++ b/cmd/register/files_test.go @@ -60,6 +60,31 @@ func TestRegisterFromFiles(t *testing.T) { err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx) assert.Nil(t, err) }) + t.Run("Register a workflow with a failure node", func(t *testing.T) { + s := setup() + testScope := promutils.NewTestScope() + labeled.SetMetricKeys(contextutils.AppNameKey, contextutils.ProjectKey, contextutils.DomainKey) + registerFilesSetup() + rconfig.DefaultFilesConfig.Archive = true + rconfig.DefaultFilesConfig.OutputLocationPrefix = s3Output + rconfig.DefaultFilesConfig.DeprecatedSourceUploadPath = s3Output + mockStorage, err := storage.NewDataStore(&storage.Config{ + Type: storage.TypeMemory, + }, testScope.NewSubScope("flytectl")) + assert.Nil(t, err) + Client = mockStorage + + args := []string{"testdata/failure-node.tgz"} + s.MockAdminClient.OnCreateTaskMatch(mock.Anything, mock.Anything).Return(nil, nil) + s.MockAdminClient.OnCreateWorkflowMatch(mock.Anything, mock.Anything).Return(nil, nil) + s.MockAdminClient.OnCreateLaunchPlanMatch(mock.Anything, mock.Anything).Return(nil, nil) + s.MockAdminClient.OnUpdateLaunchPlanMatch(mock.Anything, mock.Anything).Return(nil, nil) + mockDataProxy := s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient) + mockDataProxy.OnCreateUploadLocationMatch(s.Ctx, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil) + + err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx) + assert.Nil(t, err) + }) t.Run("Failed fast registration while uploading the codebase", func(t *testing.T) { s := setup() registerFilesSetup() diff --git a/cmd/register/register.go b/cmd/register/register.go index a04c99bd19..7caa1e9bbd 100644 --- a/cmd/register/register.go +++ b/cmd/register/register.go @@ -10,7 +10,7 @@ import ( // Long descriptions are whitespace sensitive when generating docs using sphinx. const ( registerCmdShort = "Registers tasks, workflows, and launch plans from a list of generated serialized files." - registercmdLong = ` + registerCmdLong = ` Take input files as serialized versions of the tasks/workflows/launchplans and register them with FlyteAdmin. Currently, these input files are protobuf files generated as output from Flytekit serialize. Project and Domain are mandatory fields to be passed for registration and an optional version which defaults to v1. @@ -23,7 +23,7 @@ func RemoteRegisterCommand() *cobra.Command { registerCmd := &cobra.Command{ Use: "register", Short: registerCmdShort, - Long: registercmdLong, + Long: registerCmdLong, } registerResourcesFuncs := map[string]cmdcore.CommandEntry{ "files": {CmdFunc: registerFromFilesFunc, Aliases: []string{"file"}, PFlagProvider: rconfig.DefaultFilesConfig, diff --git a/cmd/register/register_util.go b/cmd/register/register_util.go index e4dd4355ad..5234e86697 100644 --- a/cmd/register/register_util.go +++ b/cmd/register/register_util.go @@ -491,6 +491,11 @@ func hydrateSpec(message proto.Message, uploadLocation storage.DataReference, co return err } } + if workflowSpec.Template.GetFailureNode() != nil { + if err := hydrateNode(workflowSpec.Template.GetFailureNode(), config.Version, config.Force); err != nil { + return err + } + } hydrateIdentifier(workflowSpec.Template.Id, config.Version, config.Force) for _, subWorkflow := range workflowSpec.SubWorkflows { for _, Noderef := range subWorkflow.Nodes { @@ -498,6 +503,11 @@ func hydrateSpec(message proto.Message, uploadLocation storage.DataReference, co return err } } + if subWorkflow.GetFailureNode() != nil { + if err := hydrateNode(subWorkflow.GetFailureNode(), config.Version, config.Force); err != nil { + return err + } + } hydrateIdentifier(subWorkflow.Id, config.Version, config.Force) } case *admin.TaskSpec: diff --git a/cmd/register/testdata/failure-node.tgz b/cmd/register/testdata/failure-node.tgz new file mode 100644 index 0000000000000000000000000000000000000000..7ac63e86fe471e5f3407ec113b20d18325b6746f GIT binary patch literal 1671 zcmV;226*`&iwFp_WvFEW|8;I@Z*DDNX>w+4Z+9(lVPk7yXJsyQXL9!-vcA|?#I*uq=JZEE?9FhZaX8!M-IYSv6#>SV=y~C%%tBF9FijlsZ z^_LKCzTNKar-V_a6c3dYtlJ5_K;B--rZKm;7dH097t|q^#Z9+dHe6y+v&@kKw(z!* z!aaJO+Bxp;j&1e(Ivecs)U+%{Y-ZbLum4lWREk!b{eM^NH~b0ZWU+dE@!r1y<3Ei5 zL*u_w<`%I`82>Y1gjDLyqfqGL(A#S_g;Dn!&2Atneyb*oY9PE&RH{wcOGHqiMtkc$ z`@e0PXN~`q*nDvOXVfHIDUvk%e@4fD^kxa^YOeHo;#Q!owvWP{mam6iXBTsAL(OUT zs6mOrw=AP3_Nn`5o7r{3YIfZ=YD~Ur`HneW>7Tc*l-AHkNJC4$wZvdo&U2j0O+xtbSxCxzE!VR`6TbmWOVaC*)y zU9awl>TbK$@a46=a0h=P3%}j>Z15hkyu9ZAJ161WvV3pfSEsv9OCwn$(P=gJMS#mE z55TCl0%7!@WyR};txC1!hhD=MfqeRPuiFeOy;oB?JvFcEc(u1n-)q&nO@Z;*FB@Nt z3ie(HzCQTkONzFH^J@hdG=cws{}!A7%8m{GeBLqrRp$E@m~VcVFCS9iJm5dzzfAL=X<7vMFBAO7Ch^}vyTiME0yCaf#?#xc)bLCdIsbBukK zge8?tc{KnG(Eqpk&qwv2g8rwX|IS$bcNt9oWkdgYoc{Ca>p!adSb)Km`#5+Cx-)aR=R0)tA>0`leT}uY;f!AfzfUSLA6eUuaQ5 z$QPQge4)XH0`o|~f53m2lK*V#x-Rfv7WmJanE%}ro?q#Q%{-*Fdy(A@ZY89KZlh$1OCeb|JjrHuhZQ+xOfIo z-29(yIm7dRmMjkZmj?g+RY3REe5n}M3bg#q96mFNyYzFE7^I=0HDl`vzIj|;#Vh^N zjGlMJVUHX0K0PzjIOA}OfC_JjL7pPE|591m-#iFr~ z;Zbt%o?gln=s)QH82xu#w+#B91^sv8^`B2!{}~y{|Ft*+{ZB>z@2iCb;{QWJNKo)S zboW^4ZliY*VHI6PSMYN5m?2FF{D08@xz&GH4%21kfc|Ge|J}*?zjXTRf2eKIk^Emu zK>t(G|NjaJl6&cFz5nkb?|&67N070u5CwmHeBJBl8bVlYH~j5RH$2#Upd|v8D4@VT zSU`a%4k)m>2`Fd~`N?$TXP)>E)_={6|D19BM>&J#-`U{5avc9nfBhF5`Tnm%90&L> z75@t=i=|1fUB_y03tluXutmcsX6Y-@D=H?03kN&o*$68~Wk z16a@)0O&vH{}}yuDg6FJHuRsy>Hm!UfB63Iy!ihtk^j%Y{|Emc^nWq*pR+On{m+E{ zvqbBEFSH6IR{w1_y#CjQ@4uy_|9|Jux)k>Z$y$Mw>{&$t;H$BVtN$?nKQsLwJO9re zZiD=1LjQSE{pZuue?~_2AN>Dx^nYFY|4GIFr*onL!~;P8!T-O6{=Y@#s0;v*1^qWC z$NzZ>>pv)WM(6+G_g~V{|2IkI|2fS6C!YWRu|NNR3$4BbL(!U*ky~s&*3JP`n~Da( z8Zh8L;J>N!pWK1~P$u}#isQc-<$nPGrNe*m`5!z{{s)hb_k0WS9xyO4FfcGMFfcGM RFfg#0*#9H`En)zm006%ic47bk literal 0 HcmV?d00001