feat(credit): INS-5058 added check and collect for model trigger

cmd/main/main.go

@@ -32,6 +32,9 @@ import (
 	grpcrecovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
 	openfga "github.com/openfga/api/proto/openfga/v1"
 
+	"github.com/instill-ai/x/temporal"
+	"github.com/instill-ai/x/zapadapter"
+
 	"github.com/instill-ai/model-backend/config"
 	"github.com/instill-ai/model-backend/pkg/acl"
 	"github.com/instill-ai/model-backend/pkg/constant"
@@ -43,14 +46,13 @@ import (
 	"github.com/instill-ai/model-backend/pkg/service"
 	"github.com/instill-ai/model-backend/pkg/usage"
 	"github.com/instill-ai/model-backend/pkg/utils"
-	"github.com/instill-ai/x/temporal"
-	"github.com/instill-ai/x/zapadapter"
+
+	mgmtpb "github.com/instill-ai/protogen-go/core/mgmt/v1beta"
+	modelpb "github.com/instill-ai/protogen-go/model/model/v1alpha"
 
 	database "github.com/instill-ai/model-backend/pkg/db"
 	customlogger "github.com/instill-ai/model-backend/pkg/logger"
 	customotel "github.com/instill-ai/model-backend/pkg/logger/otel"
-	mgmtpb "github.com/instill-ai/protogen-go/core/mgmt/v1beta"
-	modelpb "github.com/instill-ai/protogen-go/model/model/v1alpha"
 )
 
 var propagator = b3.New(b3.WithInjectEncoding(b3.B3MultipleHeader))
@@ -231,7 +233,7 @@ func main() {
 
 	modelpb.RegisterModelPublicServiceServer(
 		publicGrpcS,
-		handler.NewPublicHandler(ctx, serv, rayService))
+		handler.NewPublicHandler(ctx, serv, rayService, nil))
 
 	modelpb.RegisterModelPrivateServiceServer(
 		privateGrpcS,

cmd/worker/main.go

@@ -18,11 +18,12 @@ import (
 
 	temporalclient "go.temporal.io/sdk/client"
 
+	"github.com/instill-ai/x/temporal"
+	"github.com/instill-ai/x/zapadapter"
+
 	"github.com/instill-ai/model-backend/config"
 	"github.com/instill-ai/model-backend/pkg/ray"
 	"github.com/instill-ai/model-backend/pkg/repository"
-	"github.com/instill-ai/x/temporal"
-	"github.com/instill-ai/x/zapadapter"
 
 	database "github.com/instill-ai/model-backend/pkg/db"
 	customlogger "github.com/instill-ai/model-backend/pkg/logger"
@@ -150,7 +151,7 @@ func main() {
 		initTemporalNamespace(ctx, tempClient)
 	}
 
-	cw := modelWorker.NewWorker(repository.NewRepository(db, redisClient), redisClient, rayService)
+	cw := modelWorker.NewWorker(repository.NewRepository(db, redisClient), redisClient, rayService, nil)
 
 	w := worker.New(tempClient, modelWorker.TaskQueue, worker.Options{})
 

pkg/constant/constant.go

@@ -13,6 +13,7 @@ const DefaultUserID string = "admin"
 const HeaderUserUIDKey = "Instill-User-Uid"
 const HeaderVisitorUIDKey = "Instill-Visitor-Uid"
 const HeaderAuthTypeKey = "Instill-Auth-Type"
+const HeaderRequesterUID = "Instill-Requester-Uid"
 
 // Ray proto path
 const RayProtoPath string = "assets/ray/proto"

pkg/handler/handler.go

@@ -8,6 +8,7 @@ import (
 	"github.com/instill-ai/model-backend/pkg/datamodel"
 	"github.com/instill-ai/model-backend/pkg/ray"
 	"github.com/instill-ai/model-backend/pkg/service"
+	"github.com/instill-ai/model-backend/pkg/usage"
 
 	healthcheckPB "github.com/instill-ai/protogen-go/common/healthcheck/v1beta"
 	modelpb "github.com/instill-ai/protogen-go/model/model/v1alpha"
@@ -17,15 +18,20 @@ var tracer = otel.Tracer("model-backend.public-handler.tracer")
 
 type PublicHandler struct {
 	modelpb.UnimplementedModelPublicServiceServer
-	service service.Service
-	ray     ray.Ray
+	service           service.Service
+	ray               ray.Ray
+	modelUsageHandler usage.ModelUsageHandler
 }
 
-func NewPublicHandler(ctx context.Context, s service.Service, r ray.Ray) modelpb.ModelPublicServiceServer {
+func NewPublicHandler(ctx context.Context, s service.Service, r ray.Ray, h usage.ModelUsageHandler) modelpb.ModelPublicServiceServer {
 	datamodel.InitJSONSchema(ctx)
+	if h == nil {
+		h = usage.NewNoopModelUsageHandler()
+	}
 	return &PublicHandler{
-		service: s,
-		ray:     r,
+		service:           s,
+		ray:               r,
+		modelUsageHandler: h,
 	}
 }
 

pkg/handler/trigger.go

@@ -18,20 +18,23 @@ import (
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/reflect/protoreflect"
 
+	"github.com/instill-ai/x/sterr"
+
 	"github.com/instill-ai/model-backend/config"
 	"github.com/instill-ai/model-backend/pkg/constant"
 	"github.com/instill-ai/model-backend/pkg/datamodel"
 	"github.com/instill-ai/model-backend/pkg/repository"
 	"github.com/instill-ai/model-backend/pkg/resource"
 	"github.com/instill-ai/model-backend/pkg/service"
+	"github.com/instill-ai/model-backend/pkg/usage"
 	"github.com/instill-ai/model-backend/pkg/utils"
-	"github.com/instill-ai/x/sterr"
 
-	custom_logger "github.com/instill-ai/model-backend/pkg/logger"
-	custom_otel "github.com/instill-ai/model-backend/pkg/logger/otel"
 	commonpb "github.com/instill-ai/protogen-go/common/task/v1alpha"
 	mgmtpb "github.com/instill-ai/protogen-go/core/mgmt/v1beta"
 	modelpb "github.com/instill-ai/protogen-go/model/model/v1alpha"
+
+	custom_logger "github.com/instill-ai/model-backend/pkg/logger"
+	custom_otel "github.com/instill-ai/model-backend/pkg/logger/otel"
 )
 
 type TriggerNamespaceModelRequestInterface interface {
@@ -98,11 +101,20 @@ func (h *PublicHandler) triggerNamespaceModel(ctx context.Context, req TriggerNa
 		span.SetStatus(1, err.Error())
 		return commonpb.Task_TASK_UNSPECIFIED, nil, err
 	}
-	if err := authenticateUser(ctx, false); err != nil {
+	if err = authenticateUser(ctx, false); err != nil {
 		span.SetStatus(1, err.Error())
 		return commonpb.Task_TASK_UNSPECIFIED, nil, err
 	}
 
+	userUID := uuid.FromStringOrNil(resource.GetRequestSingleHeader(ctx, constant.HeaderUserUIDKey))
+	if err = h.modelUsageHandler.Check(ctx, &usage.ModelUsageHandlerParams{
+		UserUID:      userUID,
+		OwnerUID:     ns.NsUID,
+		RequesterUID: uuid.FromStringOrNil(resource.GetRequestSingleHeader(ctx, constant.HeaderRequesterUID)),
+	}); err != nil {
+		return commonpb.Task_TASK_UNSPECIFIED, nil, status.Errorf(codes.FailedPrecondition, "performing usage check: %s", err.Error())
+	}
+
 	pbModel, err := h.service.GetNamespaceModelByID(ctx, ns, modelID, modelpb.View_VIEW_FULL)
 	if err != nil {
 		span.SetStatus(1, err.Error())
@@ -123,24 +135,24 @@ func (h *PublicHandler) triggerNamespaceModel(ctx context.Context, req TriggerNa
 
 	var version *datamodel.ModelVersion
 	versionID := req.GetVersion()
+	modelUID := uuid.FromStringOrNil(pbModel.Uid)
+
 	if versionID == "" {
-		version, err = h.service.GetRepository().GetLatestModelVersionByModelUID(ctx, uuid.FromStringOrNil(pbModel.Uid))
+		version, err = h.service.GetRepository().GetLatestModelVersionByModelUID(ctx, modelUID)
 		if err != nil {
 			return commonpb.Task_TASK_UNSPECIFIED, nil, status.Error(codes.NotFound, err.Error())
 		}
 	} else {
-		version, err = h.service.GetModelVersionAdmin(ctx, uuid.FromStringOrNil(pbModel.Uid), versionID)
+		version, err = h.service.GetModelVersionAdmin(ctx, modelUID, versionID)
 		if err != nil {
 			return commonpb.Task_TASK_UNSPECIFIED, nil, status.Error(codes.NotFound, err.Error())
 		}
 	}
 
-	userUID := resource.GetRequestSingleHeader(ctx, constant.HeaderUserUIDKey)
-
 	usageData := &utils.UsageMetricData{
 		OwnerUID:           ns.NsUID.String(),
 		OwnerType:          mgmtpb.OwnerType_OWNER_TYPE_USER,
-		UserUID:            userUID,
+		UserUID:            userUID.String(),
 		UserType:           mgmtpb.OwnerType_OWNER_TYPE_USER,
 		ModelUID:           pbModel.Uid,
 		Mode:               mgmtpb.Mode_MODE_SYNC,
@@ -150,13 +162,8 @@ func (h *PublicHandler) triggerNamespaceModel(ctx context.Context, req TriggerNa
 		ModelTask:          pbModel.Task,
 	}
 
-	// write usage/metric datapoint and prediction record
+	// write usage/metric datapoint
 	defer func(u *utils.UsageMetricData, startTime time.Time) {
-		// TODO: prediction feature not ready
-		// pred.ComputeTimeDuration = time.Since(startTime).Seconds()
-		// if err := h.service.CreateModelPrediction(ctx, pred); err != nil {
-		// 	logger.Warn("model prediction write failed")
-		// }
 		u.ComputeTimeDuration = time.Since(startTime).Seconds()
 		if err := h.service.WriteNewDataPoint(ctx, usageData); err != nil {
 			logger.Warn("usage/metric write failed")
@@ -360,6 +367,15 @@ func (h *PublicHandler) triggerAsyncNamespaceModel(ctx context.Context, req Trig
 		return nil, err
 	}
 
+	userUID := uuid.FromStringOrNil(resource.GetRequestSingleHeader(ctx, constant.HeaderUserUIDKey))
+	if err = h.modelUsageHandler.Check(ctx, &usage.ModelUsageHandlerParams{
+		UserUID:      userUID,
+		OwnerUID:     ns.NsUID,
+		RequesterUID: uuid.FromStringOrNil(resource.GetRequestSingleHeader(ctx, constant.HeaderRequesterUID)),
+	}); err != nil {
+		return nil, status.Errorf(codes.FailedPrecondition, "performing usage check: %s", err.Error())
+	}
+
 	pbModel, err := h.service.GetNamespaceModelByID(ctx, ns, modelID, modelpb.View_VIEW_FULL)
 	if err != nil {
 		span.SetStatus(1, err.Error())
@@ -397,24 +413,22 @@ func (h *PublicHandler) triggerAsyncNamespaceModel(ctx context.Context, req Trig
 		return nil, status.Error(codes.InvalidArgument, err.Error())
 	}
 
-	userUID := resource.GetRequestSingleHeader(ctx, constant.HeaderUserUIDKey)
-
 	// TODO: temporary solution to store input json for latest operation
 	inputRequestJSON, err := protojson.Marshal(req)
 	if err != nil {
 		return nil, err
 	}
 	h.service.GetRedisClient().Set(
 		ctx,
-		fmt.Sprintf("model_trigger_input:%s:%s", userUID, pbModel.Uid),
+		fmt.Sprintf("model_trigger_input:%s:%s", userUID.String(), pbModel.Uid),
 		inputRequestJSON,
 		time.Duration(config.Config.Server.Workflow.MaxWorkflowTimeout)*time.Second,
 	)
 
 	usageData := &utils.UsageMetricData{
 		OwnerUID:           ns.NsUID.String(),
 		OwnerType:          mgmtpb.OwnerType_OWNER_TYPE_USER,
-		UserUID:            userUID,
+		UserUID:            userUID.String(),
 		UserType:           mgmtpb.OwnerType_OWNER_TYPE_USER,
 		ModelUID:           pbModel.Uid,
 		Mode:               mgmtpb.Mode_MODE_ASYNC,
@@ -424,14 +438,9 @@ func (h *PublicHandler) triggerAsyncNamespaceModel(ctx context.Context, req Trig
 		ModelTask:          pbModel.Task,
 	}
 
-	// write usage/metric datapoint and prediction record
+	// write usage/metric datapoint
 	defer func(u *utils.UsageMetricData, startTime time.Time) {
 		if u.Status == mgmtpb.Status_STATUS_ERRORED {
-			// TODO: prediction feature not ready
-			// pred.ComputeTimeDuration = time.Since(startTime).Seconds()
-			// if err := h.service.CreateModelPrediction(ctx, pred); err != nil {
-			// 	logger.Warn("model prediction write failed")
-			// }
 			u.ComputeTimeDuration = time.Since(startTime).Seconds()
 			if err := h.service.WriteNewDataPoint(ctx, usageData); err != nil {
 				logger.Warn("usage/metric write failed")

pkg/ray/ray.go

@@ -23,9 +23,10 @@ import (
 	"github.com/instill-ai/model-backend/pkg/constant"
 	"github.com/instill-ai/model-backend/pkg/ray/rayserver"
 
-	custom_logger "github.com/instill-ai/model-backend/pkg/logger"
 	commonpb "github.com/instill-ai/protogen-go/common/task/v1alpha"
 	modelpb "github.com/instill-ai/protogen-go/model/model/v1alpha"
+
+	custom_logger "github.com/instill-ai/model-backend/pkg/logger"
 )
 
 type Ray interface {
@@ -103,10 +104,10 @@ func (r *ray) Init() {
 		}
 		initConfigData, err := yaml.Marshal(&initDeployConfig)
 		if err != nil {
-			fmt.Printf("error while Marshaling deployment config: %v", err)
+			fmt.Printf("error while Marshaling deployment config: %v\n", err)
 		}
 		if err := os.WriteFile(r.configFilePath, initConfigData, 0666); err != nil {
-			fmt.Printf("error creating deployment config: %v", err)
+			fmt.Printf("error creating deployment config: %v\n", err)
 		}
 	}
 

pkg/service/service.go

@@ -22,6 +22,8 @@ import (
 
 	custom_logger "github.com/instill-ai/model-backend/pkg/logger"
 
+	"github.com/instill-ai/x/errmsg"
+
 	"github.com/instill-ai/model-backend/config"
 	"github.com/instill-ai/model-backend/pkg/acl"
 	"github.com/instill-ai/model-backend/pkg/constant"
@@ -31,7 +33,6 @@ import (
 	"github.com/instill-ai/model-backend/pkg/resource"
 	"github.com/instill-ai/model-backend/pkg/utils"
 	"github.com/instill-ai/model-backend/pkg/worker"
-	"github.com/instill-ai/x/errmsg"
 
 	artifactpb "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha"
 	commonpb "github.com/instill-ai/protogen-go/common/task/v1alpha"
@@ -461,6 +462,7 @@ func (s *service) TriggerNamespaceModelByID(ctx context.Context, ns resource.Nam
 		workflowOptions,
 		"TriggerModelWorkflow",
 		&worker.TriggerModelWorkflowRequest{
+			TriggerUID:         uuid.FromStringOrNil(triggerUID),
 			ModelID:            dbModel.ID,
 			ModelUID:           dbModel.UID,
 			ModelVersion:       *version,

pkg/service/service_test.go

@@ -6,23 +6,23 @@ import (
 	"context"
 	"testing"
 
+	"github.com/gofrs/uuid"
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 
-	uuid "github.com/gofrs/uuid"
-	gomock "github.com/golang/mock/gomock"
-
+	"github.com/instill-ai/model-backend/pkg/datamodel"
 	"github.com/instill-ai/model-backend/pkg/service"
 
-	datamodel "github.com/instill-ai/model-backend/pkg/datamodel"
 	modelPB "github.com/instill-ai/protogen-go/model/model/v1alpha"
 )
 
 const ID = "modelID"
+
 const OWNER = "users/909c3278-f7d1-461c-9352-87741bef1ds1"
 
 var ModelDefinition, _ = uuid.FromString("909c3278-f7d1-461c-9352-87741bef11d3")
 
-//TODO: async method, need to figure out how to test this
+// TODO: async method, need to figure out how to test this
 // func TestCreateModel(t *testing.T) {
 // 	t.Run("CreateModel", func(t *testing.T) {
 // 		ctrl := gomock.NewController(t)
@@ -456,6 +456,21 @@ func TestGetModelDefinition(t *testing.T) {
 		_, err := s.GetModelDefinition(context.Background(), "github")
 		assert.NoError(t, err)
 	})
+
+	t.Run("GetModelDefinitionByUID", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
+		mockRepository := NewMockRepository(ctrl)
+		mockRepository.
+			EXPECT().
+			GetModelDefinitionByUID(ModelDefinition).
+			Return(&datamodel.ModelDefinition{}, nil).
+			Times(1)
+		s := service.NewService(mockRepository, nil, nil, nil, nil, nil, nil, nil, "")
+
+		_, err := s.GetModelDefinitionByUID(context.Background(), ModelDefinition)
+		assert.NoError(t, err)
+	})
 }
 
 func TestListModelDefinitions(t *testing.T) {
@@ -474,3 +489,6 @@ func TestListModelDefinitions(t *testing.T) {
 		assert.NoError(t, err)
 	})
 }
+
+func TestService_ListNamespaceModelVersions(t *testing.T) {
+}