From faa901c0eb1dd315130f89f13e2b6bead8e1c0dc Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 17 Jan 2024 16:14:53 -0800 Subject: [PATCH] Hydrate failure node 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 0eec7701..129d4f49 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 a04c99bd..7caa1e9b 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 e4dd4355..5234e866 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