From 0ffab4a7116bd17eab2481fe5c684036d5110f4e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 22 Apr 2024 15:17:43 -0700 Subject: [PATCH] fix(webapi): Ensure cache deletion on abort workflow (#5235) Signed-off-by: Kevin Su --- .../go/tasks/pluginmachinery/internal/webapi/core.go | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/flyteplugins/go/tasks/pluginmachinery/internal/webapi/core.go b/flyteplugins/go/tasks/pluginmachinery/internal/webapi/core.go index 9c98521897..1d0e31f71c 100644 --- a/flyteplugins/go/tasks/pluginmachinery/internal/webapi/core.go +++ b/flyteplugins/go/tasks/pluginmachinery/internal/webapi/core.go @@ -117,13 +117,19 @@ func (c CorePlugin) Abort(ctx context.Context, tCtx core.TaskExecutionContext) e } func (c CorePlugin) Finalize(ctx context.Context, tCtx core.TaskExecutionContext) error { + cacheItemID := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName() + err := c.cache.DeleteDelayed(cacheItemID) + if err != nil { + logger.Errorf(ctx, "Failed to delete resource [%v] from cache. Error: %v", cacheItemID, err) + return fmt.Errorf("failed to delete resource [%v] from cache. Error: %v", cacheItemID, err) + } + if len(c.p.GetConfig().ResourceQuotas) == 0 { // If there are no defined quotas, there is nothing to cleanup. return nil } - logger.Infof(ctx, "Attempting to finalize resource [%v].", - tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName()) + logger.Infof(ctx, "Attempting to finalize resource [%v].", cacheItemID) return c.tokenAllocator.releaseToken(ctx, c.p, tCtx, c.metrics) }