Cache Serialize API

go.mod

@@ -33,3 +33,7 @@ require (
 )
 
 replace github.com/aws/amazon-sagemaker-operator-for-k8s => github.com/aws/amazon-sagemaker-operator-for-k8s v1.0.1-0.20210303003444-0fb33b1fd49d
+
+replace github.com/flyteorg/flyteidl => github.com/hamersaw/flyteidl v0.19.26-0.20211118002926-987e338423ab
+
+replace github.com/flyteorg/flyteplugins => github.com/hamersaw/flyteplugins v0.6.2-0.20211118004720-b9783595a7ad

go.sum

@@ -236,10 +236,6 @@ github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5Kwzbycv
 github.com/fatih/color v1.10.0 h1:s36xzo75JdqLaaWoiEHk767eHiwo0598uUxyfiPkDsg=
 github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGEBuJM=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
-github.com/flyteorg/flyteidl v0.21.4 h1:gtJK5rX2ydLAo2xLRHHznOSLuLHrRRdXDbpEAlxluhk=
-github.com/flyteorg/flyteidl v0.21.4/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
-github.com/flyteorg/flyteplugins v0.7.2 h1:WLftq5+/RYt3HLc/mjRprqexot3brlMJHMH1yS7jee4=
-github.com/flyteorg/flyteplugins v0.7.2/go.mod h1:kOiuXk1ddIEVSPoHcc4kBfVQcLuyf8jw3vWJT2Was90=
 github.com/flyteorg/flytestdlib v0.3.13/go.mod h1:Tz8JCECAbX6VWGwFT6cmEQ+RJpZ/6L9pswu3fzWs220=
 github.com/flyteorg/flytestdlib v0.3.36/go.mod h1:7cDWkY3v7xsoesFcDdu6DSW5Q2U2W5KlHUbUHSwBG1Q=
 github.com/flyteorg/flytestdlib v0.4.1 h1:tEi/B9Q2T2VUOUQiFGoynmSJ8xNjdBgcZ4tgsCzpHAk=
@@ -444,6 +440,14 @@ github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t
 github.com/grpc-ecosystem/grpc-gateway v1.9.5/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0 h1:gmcG1KaJ57LophUzW0Hy8NmPhnMZb4M0+kPpLofRdBo=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
+github.com/hamersaw/flyteidl v0.19.26-0.20211103115633-100abab11c51 h1:e9zvtfNKr+K84a7du4wJgC+MXYcEsA13yYGMsEGsjQs=
+github.com/hamersaw/flyteidl v0.19.26-0.20211103115633-100abab11c51/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
+github.com/hamersaw/flyteidl v0.19.26-0.20211118002926-987e338423ab h1:1FC+p4ilWvLJZDm6DJeKWPzu+L+DRN0Y27osXgD5dzE=
+github.com/hamersaw/flyteidl v0.19.26-0.20211118002926-987e338423ab/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
+github.com/hamersaw/flyteplugins v0.6.2-0.20211103120611-e74f13d04613 h1:NSj2O4QOX5Y4t/d7ZuXcxSG2H5sHplcVuFLOmlv4/4o=
+github.com/hamersaw/flyteplugins v0.6.2-0.20211103120611-e74f13d04613/go.mod h1:ERtmZ9DEtbc9eM8v2X+bqyLJ7eER4sGOCpYPeEIizso=
+github.com/hamersaw/flyteplugins v0.6.2-0.20211118004720-b9783595a7ad h1:bUWcGrrmuiy+p+9KsMHxm3hilFNjSa4pLF0W+z/R7Lk=
+github.com/hamersaw/flyteplugins v0.6.2-0.20211118004720-b9783595a7ad/go.mod h1:UBxUP9suj5xYnue88DyQj4Oneh2tZbQ/MAGWBs4wsg0=
 github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBtguAZLlVdkD9Q=
 github.com/hashicorp/consul/api v1.3.0/go.mod h1:MmDNSzIMUjNpY/mQ398R4bk2FnqQLoPndWW5VkKPlCE=
 github.com/hashicorp/consul/sdk v0.1.1/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8=

pkg/controller/nodes/task/catalog/datacatalog/datacatalog.go

@@ -264,6 +264,81 @@ func (m *CatalogClient) Put(ctx context.Context, key catalog.Key, reader io.Outp
 	return catalog.NewStatus(core.CatalogCacheStatus_CACHE_POPULATED, EventCatalogMetadata(datasetID, tag, nil)), nil
 }
 
+// Attempts to get a reservation for the cachable task. If you have previously acquired a reservation
+// it will be extended. If another entity holds the reservation that is returned.
+func (m *CatalogClient) GetOrExtendReservation(ctx context.Context, key catalog.Key, ownerID string, heartbeatInterval time.Duration) (*datacatalog.Reservation, error) {
+	datasetID, err := GenerateDatasetIDForTask(ctx, key)
+	if err != nil {
+		return nil, err
+	}
+
+	inputs := &core.LiteralMap{}
+	if key.TypedInterface.Inputs != nil {
+		retInputs, err := key.InputReader.Get(ctx)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to read inputs when trying to query catalog")
+		}
+		inputs = retInputs
+	}
+
+	tag, err := GenerateArtifactTagName(ctx, inputs)
+	if err != nil {
+		return nil, err
+	}
+
+	reservationQuery := &datacatalog.GetOrExtendReservationRequest{
+		ReservationId: &datacatalog.ReservationID{
+			DatasetId: datasetID,
+			TagName:   tag,
+		},
+		OwnerId:           ownerID,
+		HeartbeatInterval: ptypes.DurationProto(heartbeatInterval),
+	}
+
+	response, err := m.client.GetOrExtendReservation(ctx, reservationQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	return response.Reservation, nil
+}
+
+func (m *CatalogClient) ReleaseReservation(ctx context.Context, key catalog.Key, ownerID string) error {
+	datasetID, err := GenerateDatasetIDForTask(ctx, key)
+	if err != nil {
+		return err
+	}
+
+	inputs := &core.LiteralMap{}
+	if key.TypedInterface.Inputs != nil {
+		retInputs, err := key.InputReader.Get(ctx)
+		if err != nil {
+			return errors.Wrap(err, "failed to read inputs when trying to query catalog")
+		}
+		inputs = retInputs
+	}
+
+	tag, err := GenerateArtifactTagName(ctx, inputs)
+	if err != nil {
+		return err
+	}
+
+	reservationQuery := &datacatalog.ReleaseReservationRequest{
+		ReservationId: &datacatalog.ReservationID{
+			DatasetId: datasetID,
+			TagName:   tag,
+		},
+		OwnerId: ownerID,
+	}
+
+	_, err = m.client.ReleaseReservation(ctx, reservationQuery)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // Create a new Datacatalog client for task execution caching
 func NewDataCatalog(ctx context.Context, endpoint string, insecureConnection bool, maxCacheAge time.Duration) (*CatalogClient, error) {
 	var opts []grpc.DialOption

pkg/controller/nodes/task/catalog/datacatalog/datacatalog_test.go

@@ -539,3 +539,132 @@ func TestCatalog_Put(t *testing.T) {
 	})
 
 }
+
+var tagName = "flyte_cached-BE6CZsMk6N3ExR_4X9EuwBgj2Jh2UwasXK3a_pM9xlY"
+var reservationID = datacatalog.ReservationID{
+	DatasetId: datasetID,
+	TagName:   tagName,
+}
+var prevOwner = "prevOwner"
+var currentOwner = "currentOwner"
+
+func TestCatalog_GetOrExtendReservation(t *testing.T) {
+	ctx := context.Background()
+
+	heartbeatInterval := time.Second * 5
+	prevReservation := datacatalog.Reservation{
+		ReservationId: &reservationID,
+		OwnerId:       prevOwner,
+	}
+
+	currentReservation := datacatalog.Reservation{
+		ReservationId: &reservationID,
+		OwnerId:       currentOwner,
+	}
+
+	t.Run("CreateOrUpdateReservation", func(t *testing.T) {
+		ir := &mocks2.InputReader{}
+		ir.On("Get", mock.Anything).Return(sampleParameters, nil, nil)
+
+		mockClient := &mocks.DataCatalogClient{}
+		catalogClient := &CatalogClient{
+			client: mockClient,
+		}
+
+		mockClient.On("GetOrExtendReservation",
+			ctx,
+			mock.MatchedBy(func(o *datacatalog.GetOrExtendReservationRequest) bool {
+				assert.EqualValues(t, datasetID.String(), o.ReservationId.DatasetId.String())
+				assert.EqualValues(t, tagName, o.ReservationId.TagName)
+				return true
+			}),
+		).Return(&datacatalog.GetOrExtendReservationResponse{Reservation: &currentReservation}, nil, "")
+
+		newKey := sampleKey
+		newKey.InputReader = ir
+		reservation, err := catalogClient.GetOrExtendReservation(ctx, newKey, currentOwner, heartbeatInterval)
+
+		assert.NoError(t, err)
+		assert.Equal(t, reservation.OwnerId, currentOwner)
+	})
+
+	t.Run("ExistingReservation", func(t *testing.T) {
+		ir := &mocks2.InputReader{}
+		ir.On("Get", mock.Anything).Return(sampleParameters, nil, nil)
+
+		mockClient := &mocks.DataCatalogClient{}
+		catalogClient := &CatalogClient{
+			client: mockClient,
+		}
+
+		mockClient.On("GetOrExtendReservation",
+			ctx,
+			mock.MatchedBy(func(o *datacatalog.GetOrExtendReservationRequest) bool {
+				assert.EqualValues(t, datasetID.String(), o.ReservationId.DatasetId.String())
+				assert.EqualValues(t, tagName, o.ReservationId.TagName)
+				return true
+			}),
+		).Return(&datacatalog.GetOrExtendReservationResponse{Reservation: &prevReservation}, nil, "")
+
+		newKey := sampleKey
+		newKey.InputReader = ir
+		reservation, err := catalogClient.GetOrExtendReservation(ctx, newKey, currentOwner, heartbeatInterval)
+
+		assert.NoError(t, err)
+		assert.Equal(t, reservation.OwnerId, prevOwner)
+	})
+}
+
+func TestCatalog_ReleaseReservation(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("ReleaseReservation", func(t *testing.T) {
+		ir := &mocks2.InputReader{}
+		ir.On("Get", mock.Anything).Return(sampleParameters, nil, nil)
+
+		mockClient := &mocks.DataCatalogClient{}
+		catalogClient := &CatalogClient{
+			client: mockClient,
+		}
+
+		mockClient.On("ReleaseReservation",
+			ctx,
+			mock.MatchedBy(func(o *datacatalog.ReleaseReservationRequest) bool {
+				assert.EqualValues(t, datasetID.String(), o.ReservationId.DatasetId.String())
+				assert.EqualValues(t, tagName, o.ReservationId.TagName)
+				return true
+			}),
+		).Return(&datacatalog.ReleaseReservationResponse{}, nil, "")
+
+		newKey := sampleKey
+		newKey.InputReader = ir
+		err := catalogClient.ReleaseReservation(ctx, newKey, currentOwner)
+
+		assert.NoError(t, err)
+	})
+
+	t.Run("ReleaseReservationFailure", func(t *testing.T) {
+		ir := &mocks2.InputReader{}
+		ir.On("Get", mock.Anything).Return(sampleParameters, nil, nil)
+
+		mockClient := &mocks.DataCatalogClient{}
+		catalogClient := &CatalogClient{
+			client: mockClient,
+		}
+
+		mockClient.On("ReleaseReservation",
+			ctx,
+			mock.MatchedBy(func(o *datacatalog.ReleaseReservationRequest) bool {
+				assert.EqualValues(t, datasetID.String(), o.ReservationId.DatasetId.String())
+				assert.EqualValues(t, tagName, o.ReservationId.TagName)
+				return true
+			}),
+		).Return(nil, status.Error(codes.NotFound, "reservation not found"))
+
+		newKey := sampleKey
+		newKey.InputReader = ir
+		err := catalogClient.ReleaseReservation(ctx, newKey, currentOwner)
+
+		assertGrpcErr(t, err, codes.NotFound)
+	})
+}

pkg/controller/nodes/task/catalog/noop_catalog.go

@@ -2,8 +2,10 @@ package catalog
 
 import (
 	"context"
+	"time"
 
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog"
 	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/catalog"
 	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/io"
 )
@@ -20,3 +22,11 @@ func (n NOOPCatalog) Get(_ context.Context, _ catalog.Key) (catalog.Entry, error
 func (n NOOPCatalog) Put(_ context.Context, _ catalog.Key, _ io.OutputReader, _ catalog.Metadata) (catalog.Status, error) {
 	return disabledStatus, nil
 }
+
+func (n NOOPCatalog) GetOrExtendReservation(_ context.Context, _ catalog.Key, _ string, _ time.Duration) (*datacatalog.Reservation, error) {
+	return nil, nil
+}
+
+func (n NOOPCatalog) ReleaseReservation(_ context.Context, _ catalog.Key, _ string) error {
+	return nil
+}