Dragonfly2

1
/*
2
 *     Copyright 2023 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package types
18

19
import (
20
	"fmt"
21
)
22

23
const (
24
	// ModelFileName is model file name.
25
	ModelFileName = "model.graphdef"
26

27
	// ModelConfigFileName is model config file name.
28
	ModelConfigFileName = "config.pbtxt"
29

30
	// GNNModelNameSuffix is suffix of GNN model name.
31
	GNNModelNameSuffix = "gnn"
32

33
	// MLPModelNameSuffix is suffix of MLP model name.
34
	MLPModelNameSuffix = "mlp"
35

36
	// DefaultTritonPlatform is default triton backend configuration.
37
	DefaultTritonPlatform = "tensorrt_plan"
38
)
39

40
type ModelParams struct {
41
	ID uint `uri:"id" binding:"required"`
42
}
43

44
type UpdateModelRequest struct {
45
	BIO   string `json:"BIO" binding:"omitempty"`
46
	State string `json:"state" binding:"omitempty,oneof=active inactive"`
47
}
48

49
type GetModelsQuery struct {
50
	Name        string `json:"name" binding:"omitempty"`
51
	Type        string `json:"type" binding:"omitempty"`
52
	Version     string `json:"version"  binding:"omitempty"`
53
	SchedulerID uint   `json:"scheduler_id" binding:"omitempty"`
54
	Page        int    `form:"page" binding:"omitempty,gte=1"`
55
	PerPage     int    `form:"per_page" binding:"omitempty,gte=1,lte=10000000"`
56
}
57

58
type ModelEvaluation struct {
59
	Recall    float64 `json:"recall" binding:"omitempty,gte=0,lte=1"`
60
	Precision float64 `json:"precision" binding:"omitempty,gte=0,lte=1"`
61
	F1Score   float64 `json:"f1_score" binding:"omitempty,gte=0,lte=1"`
62
	MSE       float64 `json:"mse" binding:"omitempty,gte=0"`
63
	MAE       float64 `json:"mae" binding:"omitempty,gte=0"`
64
}
65

66
// MakeObjectKeyOfModelFile returns object key of model file.
67
func MakeObjectKeyOfModelFile(id string, version int) string {
68
	return fmt.Sprintf("%s/%s/%s", id, fmt.Sprint(version), ModelFileName)
69
}
70

71
// MakeObjectKeyOfModelConfigFile returns object key of model config file.
72
func MakeObjectKeyOfModelConfigFile(id string) string {
73
	return fmt.Sprintf("%s/%s", id, ModelConfigFileName)
74
}
75