[Backend]Cache - Max cache staleness support

backend/src/cache/server/mutation.go

@@ -97,9 +97,14 @@ func MutatePodIfCached(req *v1beta1.AdmissionRequest, clientMgr ClientManagerInt
 
 	annotations[ExecutionKey] = executionHashKey
 	labels[CacheIDLabelKey] = ""
+	var maxCacheStalenessInSeconds int64 = -1
+	maxCacheStaleness, exists := annotations[MaxCacheStalenessKey]
+	if exists {
+		maxCacheStalenessInSeconds = getMaxCacheStaleness(maxCacheStaleness)
+	}
 
 	var cachedExecution *model.ExecutionCache
-	cachedExecution, err = clientMgr.CacheStore().GetExecutionCache(executionHashKey)
+	cachedExecution, err = clientMgr.CacheStore().GetExecutionCache(executionHashKey, maxCacheStalenessInSeconds)
 	if err != nil {
 		log.Println(err.Error())
 	}

backend/src/cache/server/mutation_test.go

@@ -129,6 +129,7 @@ func TestMutatePodIfCachedWithCacheEntryExist(t *testing.T) {
 		ExecutionCacheKey: "f98b62e4625b9f96bac478ac72d88181a37e4f1d6bfd3bd5f53e29286b2ca034",
 		ExecutionOutput:   "testOutput",
 		ExecutionTemplate: `{"name": "test_template"}`,
+		MaxCacheStaleness: -1,
 	}
 	fakeClientManager.CacheStore().CreateExecutionCache(executionCache)
 

backend/src/cache/server/watcher.go

@@ -6,9 +6,11 @@ import (
 	"log"
 	"reflect"
 	"strconv"
+	"time"
 
 	"github.com/kubeflow/pipelines/backend/src/cache/client"
 	"github.com/kubeflow/pipelines/backend/src/cache/model"
+	"github.com/peterhellberg/duration"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
@@ -18,6 +20,7 @@ import (
 const (
 	ArgoCompleteLabelKey   string = "workflows.argoproj.io/completed"
 	MetadataExecutionIDKey string = "pipelines.kubeflow.org/metadata_execution_id"
+	MaxCacheStalenessKey   string = "pipelines.kubeflow.org/max_cache_staleness"
 )
 
 func WatchPods(namespaceToWatch string, clientManager ClientManagerInterface) {
@@ -61,12 +64,18 @@ func WatchPods(namespaceToWatch string, clientManager ClientManagerInterface) {
 			executionOutputMap[MetadataExecutionIDKey] = pod.ObjectMeta.Labels[MetadataExecutionIDKey]
 			executionOutputJSON, _ := json.Marshal(executionOutputMap)
 
+			executionMaxCacheStaleness, exists := pod.ObjectMeta.Annotations[MaxCacheStalenessKey]
+			var maxCacheStalenessInSeconds int64 = -1
+			if exists {
+				maxCacheStalenessInSeconds = getMaxCacheStaleness(executionMaxCacheStaleness)
+			}
+
 			executionTemplate := pod.ObjectMeta.Annotations[ArgoWorkflowTemplate]
 			executionToPersist := model.ExecutionCache{
 				ExecutionCacheKey: executionKey,
 				ExecutionTemplate: executionTemplate,
 				ExecutionOutput:   string(executionOutputJSON),
-				MaxCacheStaleness: -1,
+				MaxCacheStaleness: maxCacheStalenessInSeconds,
 			}
 
 			cacheEntryCreated, err := clientManager.CacheStore().CreateExecutionCache(&executionToPersist)
@@ -112,3 +121,12 @@ func patchCacheID(k8sCore client.KubernetesCoreInterface, podToPatch *corev1.Pod
 	log.Printf("Cache id patched.")
 	return nil
 }
+
+// Convert RFC3339 Duration(Eg. "P1DT30H4S") to int64 seconds.
+func getMaxCacheStaleness(maxCacheStaleness string) int64 {
+	var seconds int64 = -1
+	if d, err := duration.Parse(maxCacheStaleness); err == nil {
+		seconds = int64(d / time.Second)
+	}
+	return seconds
+}

backend/src/cache/storage/execution_cache_store.go

@@ -25,7 +25,7 @@ import (
 )
 
 type ExecutionCacheStoreInterface interface {
-	GetExecutionCache(executionCacheKey string) (*model.ExecutionCache, error)
+	GetExecutionCache(executionCacheKey string, maxCacheStaleness int64) (*model.ExecutionCache, error)
 	CreateExecutionCache(*model.ExecutionCache) (*model.ExecutionCache, error)
 	DeleteExecutionCache(executionCacheKey string) error
 }
@@ -35,13 +35,16 @@ type ExecutionCacheStore struct {
 	time util.TimeInterface
 }
 
-func (s *ExecutionCacheStore) GetExecutionCache(executionCacheKey string) (*model.ExecutionCache, error) {
+func (s *ExecutionCacheStore) GetExecutionCache(executionCacheKey string, maxCacheStaleness int64) (*model.ExecutionCache, error) {
+	if maxCacheStaleness == 0 {
+		return nil, fmt.Errorf("MaxCacheStaleness=0, Cache is disabled.")
+	}
 	r, err := s.db.Table("execution_caches").Where("ExecutionCacheKey = ?", executionCacheKey).Rows()
 	if err != nil {
 		return nil, fmt.Errorf("Failed to get execution cache: %q", executionCacheKey)
 	}
 	defer r.Close()
-	executionCaches, err := s.scanRows(r)
+	executionCaches, err := s.scanRows(r, maxCacheStaleness)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to get execution cache: %q", executionCacheKey)
 	}
@@ -55,7 +58,7 @@ func (s *ExecutionCacheStore) GetExecutionCache(executionCacheKey string) (*mode
 	return latestCache, nil
 }
 
-func (s *ExecutionCacheStore) scanRows(rows *sql.Rows) ([]*model.ExecutionCache, error) {
+func (s *ExecutionCacheStore) scanRows(rows *sql.Rows, podMaxCacheStaleness int64) ([]*model.ExecutionCache, error) {
 	var executionCaches []*model.ExecutionCache
 	for rows.Next() {
 		var executionCacheKey, executionTemplate, executionOutput string
@@ -73,15 +76,18 @@ func (s *ExecutionCacheStore) scanRows(rows *sql.Rows) ([]*model.ExecutionCache,
 		}
 		log.Println("Get id: " + strconv.FormatInt(id, 10))
 		log.Println("Get template: " + executionTemplate)
-		executionCaches = append(executionCaches, &model.ExecutionCache{
-			ID:                id,
-			ExecutionCacheKey: executionCacheKey,
-			ExecutionTemplate: executionTemplate,
-			ExecutionOutput:   executionOutput,
-			MaxCacheStaleness: maxCacheStaleness,
-			StartedAtInSec:    startedAtInSec,
-			EndedAtInSec:      endedAtInSec,
-		})
+		if maxCacheStaleness == -1 || s.time.Now().UTC().Unix()-startedAtInSec <= podMaxCacheStaleness {
+			executionCaches = append(executionCaches, &model.ExecutionCache{
+				ID:                id,
+				ExecutionCacheKey: executionCacheKey,
+				ExecutionTemplate: executionTemplate,
+				ExecutionOutput:   executionOutput,
+				MaxCacheStaleness: maxCacheStaleness,
+				StartedAtInSec:    startedAtInSec,
+				EndedAtInSec:      endedAtInSec,
+			})
+		}
+
 	}
 	return executionCaches, nil
 }
@@ -112,7 +118,6 @@ func (s *ExecutionCacheStore) CreateExecutionCache(executionCache *model.Executi
 	newExecutionCache.StartedAtInSec = now
 	// TODO: ended time need to be modified after demo version.
 	newExecutionCache.EndedAtInSec = now
-	newExecutionCache.MaxCacheStaleness = -1
 
 	ok := s.db.NewRecord(newExecutionCache)
 	if !ok {

backend/src/cache/storage/execution_cache_store_test.go

@@ -51,6 +51,7 @@ func TestCreateExecutionCache(t *testing.T) {
 		ExecutionCacheKey: "test",
 		ExecutionTemplate: "testTemplate",
 		ExecutionOutput:   "testOutput",
+		MaxCacheStaleness: -1,
 	}
 	executionCache, err := executionCacheStore.CreateExecutionCache(executionCache)
 	assert.Nil(t, err)
@@ -93,7 +94,7 @@ func TestGetExecutionCache(t *testing.T) {
 	}
 
 	var executionCache *model.ExecutionCache
-	executionCache, err := executionCacheStore.GetExecutionCache("testKey")
+	executionCache, err := executionCacheStore.GetExecutionCache("testKey", -1)
 	require.Nil(t, err)
 	require.Equal(t, &executionCacheExpected, executionCache)
 }
@@ -105,7 +106,7 @@ func TestGetExecutionCacheWithEmptyCacheEntry(t *testing.T) {
 
 	executionCacheStore.CreateExecutionCache(createExecutionCache("testKey", "testOutput"))
 	var executionCache *model.ExecutionCache
-	executionCache, err := executionCacheStore.GetExecutionCache("wrongKey")
+	executionCache, err := executionCacheStore.GetExecutionCache("wrongKey", -1)
 	require.Nil(t, executionCache)
 	require.Contains(t, err.Error(), `Execution cache not found with cache key: "wrongKey"`)
 }
@@ -128,23 +129,41 @@ func TestGetExecutionCacheWithLatestCacheEntry(t *testing.T) {
 		EndedAtInSec:      2,
 	}
 	var executionCache *model.ExecutionCache
-	executionCache, err := executionCacheStore.GetExecutionCache("testKey")
+	executionCache, err := executionCacheStore.GetExecutionCache("testKey", -1)
 	require.Nil(t, err)
 	require.Equal(t, &executionCacheExpected, executionCache)
 }
 
+func TestGetExecutionCacheWithExpiredMaxCacheStaleness(t *testing.T) {
+	db := NewFakeDbOrFatal()
+	defer db.Close()
+	executionCacheStore := NewExecutionCacheStore(db, util.NewFakeTimeForEpoch())
+	executionCacheToPersist := &model.ExecutionCache{
+		ExecutionCacheKey: "testKey",
+		ExecutionTemplate: "testTemplate",
+		ExecutionOutput:   "testOutput",
+		MaxCacheStaleness: 0,
+	}
+	executionCacheStore.CreateExecutionCache(executionCacheToPersist)
+
+	var executionCache *model.ExecutionCache
+	executionCache, err := executionCacheStore.GetExecutionCache("testKey", -1)
+	require.Contains(t, err.Error(), "Execution cache not found")
+	require.Nil(t, executionCache)
+}
+
 func TestDeleteExecutionCache(t *testing.T) {
 	db := NewFakeDbOrFatal()
 	defer db.Close()
 	executionCacheStore := NewExecutionCacheStore(db, util.NewFakeTimeForEpoch())
 	executionCacheStore.CreateExecutionCache(createExecutionCache("testKey", "testOutput"))
-	executionCache, err := executionCacheStore.GetExecutionCache("testKey")
+	executionCache, err := executionCacheStore.GetExecutionCache("testKey", -1)
 	assert.Nil(t, err)
 	assert.NotNil(t, executionCache)
 
 	err = executionCacheStore.DeleteExecutionCache("1")
 	assert.Nil(t, err)
-	_, err = executionCacheStore.GetExecutionCache("testKey")
+	_, err = executionCacheStore.GetExecutionCache("testKey", -1)
 	assert.NotNil(t, err)
 	assert.Contains(t, err.Error(), "not found")
 }

go.mod

@@ -54,6 +54,7 @@ require (
 	github.com/onsi/ginkgo v1.7.0 // indirect
 	github.com/onsi/gomega v1.4.3 // indirect
 	github.com/peterbourgon/diskv v2.0.1+incompatible // indirect
+	github.com/peterhellberg/duration v0.0.0-20191119133758-ec6baeebcd10
 	github.com/pkg/errors v0.8.0
 	github.com/prometheus/client_golang v0.9.2 // indirect
 	github.com/robfig/cron v0.0.0-20180505203441-b41be1df6967