Migrate Hive Plugin to Service Plugins Interface

go.mod

@@ -32,7 +32,7 @@ require (
 	k8s.io/api v0.17.3
 	k8s.io/apimachinery v0.17.3
 	k8s.io/client-go v11.0.1-0.20190918222721-c0e3722d5cf0+incompatible
-	k8s.io/utils v0.0.0-20200124190032-861946025e34 // indirect
+	k8s.io/utils v0.0.0-20200124190032-861946025e34
 	sigs.k8s.io/controller-runtime v0.5.1
 )
 

go/tasks/pluginmachinery/core/resource_manager.go

@@ -4,21 +4,23 @@ import (
 	"context"
 )
 
-type AllocationStatus string
+//go:generate enumer -type=AllocationStatus -trimprefix=AllocationStatus
+
+type AllocationStatus int
 
 const (
 	// This is the enum returned when there's an error
-	AllocationUndefined AllocationStatus = "ResourceGranted"
+	AllocationUndefined AllocationStatus = iota
 
 	// Go for it
-	AllocationStatusGranted AllocationStatus = "ResourceGranted"
+	AllocationStatusGranted
 
 	// This means that no resources are available globally.  This is the only rejection message we use right now.
-	AllocationStatusExhausted AllocationStatus = "ResourceExhausted"
+	AllocationStatusExhausted
 
 	// We're not currently using this - but this would indicate that things globally are okay, but that your
 	// own namespace is too busy
-	AllocationStatusNamespaceQuotaExceeded AllocationStatus = "NamespaceQuotaExceeded"
+	AllocationStatusNamespaceQuotaExceeded
 )
 
 const namespaceSeparator = ":"

go/tasks/pluginmachinery/internal/remote/allocation_token.go

@@ -0,0 +1,101 @@
+package remote
+
+import (
+	"context"
+	"fmt"
+
+	clock2 "k8s.io/utils/clock"
+
+	"github.com/lyft/flytestdlib/logger"
+
+	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/remote"
+
+	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core"
+)
+
+var (
+	clock clock2.Clock
+)
+
+func SetClockForTest(clck clock2.Clock) {
+	clock = clck
+}
+
+func init() {
+	clock = clock2.RealClock{}
+}
+
+func allocateToken(ctx context.Context, p remote.Plugin, tCtx core.TaskExecutionContext, state *State, metrics Metrics) (
+	newState *State, phaseInfo core.PhaseInfo, err error) {
+	if len(p.GetPluginProperties().ResourceQuotas) == 0 {
+		// No quota, return success
+		return &State{
+			AllocationTokenRequestStartTime: clock.Now(),
+			Phase:                           PhaseAllocationTokenAcquired,
+		}, core.PhaseInfo{}, nil
+	}
+
+	ns, constraints, err := p.ResourceRequirements(ctx, tCtx)
+	if err != nil {
+		logger.Errorf(ctx, "Failed to calculate resource requirements for task. Error: %v", err)
+		return nil, core.PhaseInfo{}, err
+	}
+
+	token := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName()
+	allocationStatus, err := tCtx.ResourceManager().AllocateResource(ctx, ns, token, constraints)
+	if err != nil {
+		logger.Errorf(ctx, "Failed to allocate resources for task. Error: %v", err)
+		return nil, core.PhaseInfo{}, err
+	}
+
+	switch allocationStatus {
+	case core.AllocationStatusGranted:
+		metrics.AllocationGranted.Inc(ctx)
+		metrics.ResourceWaitTime.Observe(float64(clock.Since(state.AllocationTokenRequestStartTime).Milliseconds()))
+		return &State{
+			AllocationTokenRequestStartTime: clock.Now(),
+			Phase:                           PhaseAllocationTokenAcquired,
+		}, core.PhaseInfo{}, nil
+	case core.AllocationStatusNamespaceQuotaExceeded:
+	case core.AllocationStatusExhausted:
+		metrics.AllocationNotGranted.Inc(ctx)
+		logger.Infof(ctx, "Couldn't allocate token because allocation status is [%v].", allocationStatus.String())
+		startTime := state.AllocationTokenRequestStartTime
+		if startTime.IsZero() {
+			startTime = clock.Now()
+		}
+
+		return &State{
+			AllocationTokenRequestStartTime: startTime,
+			Phase:                           PhaseNotStarted,
+		}, core.PhaseInfo{}, nil
+	default:
+		return nil, core.PhaseInfo{}, fmt.Errorf("allocation status undefined")
+	}
+
+	return state, core.PhaseInfo{}, nil
+}
+
+func releaseToken(ctx context.Context, p remote.Plugin, tCtx core.TaskExecutionContext, metrics Metrics) error {
+	if len(p.GetPluginProperties().ResourceQuotas) == 0 {
+		// No quota, return success
+		return nil
+	}
+
+	ns, _, err := p.ResourceRequirements(ctx, tCtx)
+	if err != nil {
+		logger.Errorf(ctx, "Failed to calculate resource requirements for task. Error: %v", err)
+		return err
+	}
+
+	token := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName()
+	err = tCtx.ResourceManager().ReleaseResource(ctx, ns, token)
+	if err != nil {
+		metrics.ResourceReleaseFailed.Inc(ctx)
+		logger.Errorf(ctx, "Failed to release resources for task. Error: %v", err)
+		return err
+	}
+
+	metrics.ResourceReleased.Inc(ctx)
+	return nil
+}

go/tasks/pluginmachinery/internal/remote/allocation_token_test.go

@@ -0,0 +1,149 @@
+package remote
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/mock"
+
+	mocks2 "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core/mocks"
+	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/remote/mocks"
+
+	"github.com/stretchr/testify/assert"
+
+	testing2 "k8s.io/utils/clock/testing"
+
+	"github.com/lyft/flytestdlib/contextutils"
+	"github.com/lyft/flytestdlib/promutils/labeled"
+
+	"github.com/go-test/deep"
+	"github.com/lyft/flytestdlib/promutils"
+
+	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core"
+	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/remote"
+)
+
+func init() {
+	labeled.SetMetricKeys(contextutils.NamespaceKey)
+}
+
+func newPluginWithProperties(properties remote.PluginProperties) *mocks.Plugin {
+	m := &mocks.Plugin{}
+	m.OnGetPluginProperties().Return(properties)
+	return m
+}
+
+func Test_allocateToken(t *testing.T) {
+	ctx := context.Background()
+	metrics := newMetrics(promutils.NewTestScope())
+
+	tNow := time.Now()
+	clck := testing2.NewFakeClock(tNow)
+	SetClockForTest(clck)
+
+	tID := &mocks2.TaskExecutionID{}
+	tID.OnGetGeneratedName().Return("abc")
+
+	tMeta := &mocks2.TaskExecutionMetadata{}
+	tMeta.OnGetTaskExecutionID().Return(tID)
+
+	rm := &mocks2.ResourceManager{}
+	rm.OnAllocateResourceMatch(ctx, core.ResourceNamespace("ns"), "abc", mock.Anything).Return(core.AllocationStatusGranted, nil)
+	rm.OnAllocateResourceMatch(ctx, core.ResourceNamespace("ns"), "abc2", mock.Anything).Return(core.AllocationStatusExhausted, nil)
+
+	tCtx := &mocks2.TaskExecutionContext{}
+	tCtx.OnTaskExecutionMetadata().Return(tMeta)
+	tCtx.OnResourceManager().Return(rm)
+
+	state := &State{}
+
+	p := newPluginWithProperties(remote.PluginProperties{
+		ResourceQuotas: map[core.ResourceNamespace]int{
+			"ns": 1,
+		},
+	})
+
+	t.Run("no quota", func(t *testing.T) {
+		p := newPluginWithProperties(remote.PluginProperties{ResourceQuotas: nil})
+		gotNewState, _, err := allocateToken(ctx, p, nil, nil, metrics)
+		assert.NoError(t, err)
+		if diff := deep.Equal(gotNewState, &State{
+			AllocationTokenRequestStartTime: tNow,
+			Phase:                           PhaseAllocationTokenAcquired,
+		}); len(diff) > 0 {
+			t.Errorf("allocateToken() gotNewState = %v, Diff: %v", gotNewState, diff)
+		}
+	})
+
+	t.Run("Allocation Successful", func(t *testing.T) {
+		p.OnResourceRequirements(ctx, tCtx).Return("ns", core.ResourceConstraintsSpec{}, nil)
+		gotNewState, _, err := allocateToken(ctx, p, tCtx, state, metrics)
+		assert.NoError(t, err)
+		if diff := deep.Equal(gotNewState, &State{
+			AllocationTokenRequestStartTime: tNow,
+			Phase:                           PhaseAllocationTokenAcquired,
+		}); len(diff) > 0 {
+			t.Errorf("allocateToken() gotNewState = %v, Diff: %v", gotNewState, diff)
+		}
+	})
+
+	t.Run("Allocation Failed", func(t *testing.T) {
+		tID := &mocks2.TaskExecutionID{}
+		tID.OnGetGeneratedName().Return("abc2")
+
+		tMeta := &mocks2.TaskExecutionMetadata{}
+		tMeta.OnGetTaskExecutionID().Return(tID)
+
+		rm := &mocks2.ResourceManager{}
+		rm.OnAllocateResourceMatch(ctx, core.ResourceNamespace("ns"), "abc", mock.Anything).Return(core.AllocationStatusGranted, nil)
+		rm.OnAllocateResourceMatch(ctx, core.ResourceNamespace("ns"), "abc2", mock.Anything).Return(core.AllocationStatusExhausted, nil)
+
+		tCtx := &mocks2.TaskExecutionContext{}
+		tCtx.OnTaskExecutionMetadata().Return(tMeta)
+		tCtx.OnResourceManager().Return(rm)
+
+		p.OnResourceRequirements(ctx, tCtx).Return("ns", core.ResourceConstraintsSpec{}, nil)
+		gotNewState, _, err := allocateToken(ctx, p, tCtx, state, metrics)
+		assert.NoError(t, err)
+		if diff := deep.Equal(gotNewState, &State{
+			AllocationTokenRequestStartTime: tNow,
+			Phase:                           PhaseNotStarted,
+		}); len(diff) > 0 {
+			t.Errorf("allocateToken() gotNewState = %v, Diff: %v", gotNewState, diff)
+		}
+	})
+}
+
+func Test_releaseToken(t *testing.T) {
+	ctx := context.Background()
+	metrics := newMetrics(promutils.NewTestScope())
+
+	tNow := time.Now()
+	clck := testing2.NewFakeClock(tNow)
+	SetClockForTest(clck)
+
+	tID := &mocks2.TaskExecutionID{}
+	tID.OnGetGeneratedName().Return("abc")
+
+	tMeta := &mocks2.TaskExecutionMetadata{}
+	tMeta.OnGetTaskExecutionID().Return(tID)
+
+	rm := &mocks2.ResourceManager{}
+	rm.OnAllocateResourceMatch(ctx, core.ResourceNamespace("ns"), "abc", mock.Anything).Return(core.AllocationStatusGranted, nil)
+	rm.OnAllocateResourceMatch(ctx, core.ResourceNamespace("ns"), "abc2", mock.Anything).Return(core.AllocationStatusExhausted, nil)
+	rm.OnReleaseResource(ctx, core.ResourceNamespace("ns"), "abc").Return(nil)
+
+	tCtx := &mocks2.TaskExecutionContext{}
+	tCtx.OnTaskExecutionMetadata().Return(tMeta)
+	tCtx.OnResourceManager().Return(rm)
+
+	p := newPluginWithProperties(remote.PluginProperties{
+		ResourceQuotas: map[core.ResourceNamespace]int{
+			"ns": 1,
+		},
+	})
+	p.OnResourceRequirements(ctx, tCtx).Return("ns", core.ResourceConstraintsSpec{}, nil)
+
+	assert.NoError(t, releaseToken(ctx, p, tCtx, metrics))
+}