feat: openmodelz support deployment event (#189)

Signed-off-by: xieydd <[email protected]>

agent/Makefile

@@ -126,9 +126,6 @@ mockgen-install:
 addlicense-install:
 	go install github.com/google/addlicense@latest
 
-sqlc-install:
-	go install github.com/sqlc-dev/sqlc/cmd/sqlc@latest
-
 # https://github.com/swaggo/swag/pull/1322, we should use master instead of latest for now.
 swag-install:
 	go install github.com/swaggo/swag/cmd/[email protected]
@@ -195,7 +192,5 @@ release:
 		goreleaser/goreleaser-cross:${GOLANG_CROSS_VERSION} \
 		release --rm-dist
 
-generate: mockgen-install sqlc-install swag
-	@mockgen -source pkg/query/querier.go -destination pkg/query/mock/mock.go -package mock
+generate: mockgen-install swag
 	@mockgen -source pkg/runtime/runtime.go -destination pkg/runtime/mock/mock.go -package mock
-	# @sqlc generate

agent/api/types/event.go

@@ -1,5 +1,7 @@
 package types
 
+import "time"
+
 const (
 	DeploymentCreateEvent     = "deployment-create"
 	DeploymentUpdateEvent     = "deployment-update"
@@ -8,3 +10,12 @@ const (
 	DeploymentScaleDownEvent  = "deployment-scale-down"
 	DeploymentScaleBlockEvent = "deployment-scale-block"
 )
+
+type DeploymentEvent struct {
+	ID           string    `json:"id"`
+	CreatedAt    time.Time `json:"created_at"`
+	UserID       string    `json:"user_id"`
+	DeploymentID string    `json:"deployment_id"`
+	EventType    string    `json:"event_type"`
+	Message      string    `json:"message"`
+}

agent/client/const.go

@@ -13,22 +13,23 @@ const defaultAddr = "0.0.0.0:8080"
 const apiBasePath = ""
 
 const (
-	gatewayInferControlPlanePath                 = "/system/inferences"
-	gatewayInferScaleControlPath                 = "/system/scale-inference"
-	gatewayInferInstanceControlPlanePath         = "/system/inference/%s/instances"
-	gatewayInferInstanceExecControlPlanePath     = "/system/inference/%s/instance/%s/exec"
-	gatewayServerControlPlanePath                = "/system/servers"
-	gatewayServerLabelCreateControlPlanePath     = "/system/server/%s/labels"
-	gatewayServerNodeDeleteControlPlanePath      = "/system/server/%s/delete"
-	gatewayNamespaceControlPlanePath             = "/system/namespaces"
-	gatewayBuildControlPlanePath                 = "/system/build"
-	gatewayBuildInstanceControlPlanePath         = "/system/build/%s"
-	gatewayImageCacheControlPlanePath            = "/system/image-cache"
-	modelzCloudClusterControlPlanePath           = "/api/v1/users/%s/clusters/%s"
-	modelzCloudClusterWithUserControlPlanePath   = "/api/v1/users/%s/clusters"
-	modelzCloudClusterAPIKeyControlPlanePath     = "/api/v1/users/%s/clusters/%s/api_keys"
-	modelzCloudClusterNamespaceControlPlanePath  = "/api/v1/users/%s/clusters/%s/namespaces"
-	modelzCloudClusterDeploymentControlPlanePath = "/api/v1/users/%s/clusters/%s/deployments/%s/agent"
+	gatewayInferControlPlanePath                      = "/system/inferences"
+	gatewayInferScaleControlPath                      = "/system/scale-inference"
+	gatewayInferInstanceControlPlanePath              = "/system/inference/%s/instances"
+	gatewayInferInstanceExecControlPlanePath          = "/system/inference/%s/instance/%s/exec"
+	gatewayServerControlPlanePath                     = "/system/servers"
+	gatewayServerLabelCreateControlPlanePath          = "/system/server/%s/labels"
+	gatewayServerNodeDeleteControlPlanePath           = "/system/server/%s/delete"
+	gatewayNamespaceControlPlanePath                  = "/system/namespaces"
+	gatewayBuildControlPlanePath                      = "/system/build"
+	gatewayBuildInstanceControlPlanePath              = "/system/build/%s"
+	gatewayImageCacheControlPlanePath                 = "/system/image-cache"
+	modelzCloudClusterControlPlanePath                = "/api/v1/users/%s/clusters/%s"
+	modelzCloudClusterWithUserControlPlanePath        = "/api/v1/users/%s/clusters"
+	modelzCloudClusterAPIKeyControlPlanePath          = "/api/v1/users/%s/clusters/%s/api_keys"
+	modelzCloudClusterNamespaceControlPlanePath       = "/api/v1/users/%s/clusters/%s/namespaces"
+	modelzCloudClusterDeploymentControlPlanePath      = "/api/v1/users/%s/clusters/%s/deployments/%s/agent"
+	modelzCloudClusterDeploymentEventControlPlanePath = "/api/v1/users/%s/clusters/%s/deployments/%s/event"
 )
 
 const (

agent/client/modelz_cloud.go

@@ -165,3 +165,25 @@ func (cli *Client) GetUIDFromDeploymentID(ctx context.Context, token string, clu
 	}
 	return "", fmt.Errorf("failed to get uid from deployment id, status code: %d", resp.statusCode)
 }
+
+func (cli *Client) CreateDeploymentEvent(ctx context.Context, token string, event types.DeploymentEvent) error {
+	urlValues := url.Values{}
+	agentToken, err := ParseAgentToken(token)
+	if err != nil {
+		return err
+	}
+	headers := make(map[string][]string)
+	headers["Authorization"] = []string{"Bearer " + agentToken.Token}
+	urlPath := fmt.Sprintf(modelzCloudClusterDeploymentEventControlPlanePath, agentToken.UID, agentToken.ClusterName, event.DeploymentID)
+
+	resp, err := cli.post(ctx, urlPath, urlValues, event, headers)
+	if err != nil {
+		return err
+	}
+	defer ensureReaderClosed(resp)
+
+	if resp.statusCode == http.StatusCreated {
+		return nil
+	}
+	return fmt.Errorf("failed to create deployment event, status code: %d", resp.statusCode)
+}

agent/pkg/app/config.go

@@ -54,10 +54,6 @@ func configFromCLI(c *cli.Context) config.Config {
 	cfg.Metrics.PrometheusHost = c.String(flagMetricsPrometheusHost)
 	cfg.Metrics.PrometheusPort = c.Int(flagMetricsPrometheusPort)
 
-	// postgres database
-	cfg.DB.EventEnabled = c.Bool(flagEventEnabled)
-	cfg.DB.URL = c.String(flagDBURL)
-
 	// modelz cloud
 	cfg.ModelZCloud.Enabled = c.Bool(flagModelZCloudEnabled)
 	cfg.ModelZCloud.URL = c.String(flagModelZCloudURL)
@@ -68,5 +64,6 @@ func configFromCLI(c *cli.Context) config.Config {
 	cfg.ModelZCloud.UpstreamTimeout = c.Duration(flagModelZCloudUpstreamTimeout)
 	cfg.ModelZCloud.MaxIdleConnections = c.Int(flagModelZCloudMaxIdleConnections)
 	cfg.ModelZCloud.MaxIdleConnectionsPerHost = c.Int(flagModelZCloudMaxIdleConnectionsPerHost)
+	cfg.ModelZCloud.EventEnabled = c.Bool(flagModelZCloudEventEnabled)
 	return cfg
 }

agent/pkg/app/root.go

@@ -64,10 +64,6 @@ const (
 	flaglogsLokiUser  = "logs-loki-user"
 	flagLogsLokiToken = "logs-loki-token"
 
-	// db
-	flagEventEnabled = "event-enabled"
-	flagDBURL        = "db-url"
-
 	// modelz cloud
 	flagModelZCloudEnabled                   = "modelz-cloud-enabled"
 	flagModelZCloudURL                       = "modelz-cloud-url"
@@ -78,6 +74,7 @@ const (
 	flagModelZCloudUpstreamTimeout           = "modelz-cloud-upstream-timeout"
 	flagModelZCloudMaxIdleConnections        = "modelz-cloud-max-idle-connections"
 	flagModelZCloudMaxIdleConnectionsPerHost = "modelz-cloud-max-idle-connections-per-host"
+	flagModelZCloudEventEnabled              = "modelz-cloud-event-enabled"
 )
 
 type App struct {
@@ -317,21 +314,6 @@ func New() App {
 			Usage:   "Loki service auth token",
 			EnvVars: []string{"MODELZ_AGENT_LOGS_LOKI_TOKEN"},
 		},
-		&cli.BoolFlag{
-			Name:    flagEventEnabled,
-			Hidden:  true,
-			Usage:   "Enable event logging",
-			Value:   false,
-			EnvVars: []string{"MODELZ_AGENT_EVENT_ENABLED"},
-			Aliases: []string{"ee"},
-		},
-		&cli.StringFlag{
-			Name:    flagDBURL,
-			Usage:   "Postgres database URL",
-			Hidden:  true,
-			Aliases: []string{"du"},
-			EnvVars: []string{"MODELZ_AGENT_DB_URL"},
-		},
 		&cli.BoolFlag{
 			Name:    flagModelZCloudEnabled,
 			Usage:   "Enable modelz cloud, agent as modelz cloud agent",
@@ -393,6 +375,13 @@ func New() App {
 			Aliases: []string{"mich"},
 			Value:   1024,
 		},
+		&cli.BoolFlag{
+			Name:    flagModelZCloudEventEnabled,
+			Usage:   "Enable event logging for modelz cloud.",
+			Value:   false,
+			EnvVars: []string{"MODELZ_AGENT_MODELZ_CLOUD_EVENT_ENABLED"},
+			Aliases: []string{"mze"},
+		},
 	}
 	internalApp.Action = runServer
 

agent/pkg/config/config.go

@@ -14,7 +14,6 @@ type Config struct {
 	Build       BuildConfig       `json:"build,omitempty"`
 	Metrics     MetricsConfig     `json:"metrics,omitempty"`
 	Logs        LogsConfig        `json:"logs,omitempty"`
-	DB          PostgresConfig    `json:"db,omitempty"`
 	ModelZCloud ModelZCloudConfig `json:"modelz_cloud,omitempty"`
 }
 
@@ -34,6 +33,7 @@ type ModelZCloudConfig struct {
 	UpstreamTimeout           time.Duration     `json:"upstream_timeout,omitempty"`
 	MaxIdleConnections        int               `json:"max_idle_connections,omitempty"`
 	MaxIdleConnectionsPerHost int               `json:"max_idle_connections_per_host,omitempty"`
+	EventEnabled              bool              `json:"event_enabled,omitempty"`
 }
 
 type LogsConfig struct {
@@ -88,11 +88,6 @@ type KubeConfig struct {
 	ResyncPeriod time.Duration `json:"resync_period,omitempty"`
 }
 
-type PostgresConfig struct {
-	EventEnabled bool   `json:"event_enabled,omitempty"`
-	URL          string `json:"url,omitempty"`
-}
-
 func New() Config {
 	return Config{
 		KubeConfig: KubeConfig{},
@@ -138,12 +133,6 @@ func (c Config) Validate() error {
 		return errors.New("metrics config is required")
 	}
 
-	if c.DB.EventEnabled {
-		if c.DB.URL == "" {
-			return errors.New("db config is required")
-		}
-	}
-
 	if c.Ingress.IngressEnabled {
 		if c.Ingress.Namespace == "" {
 			return errors.New("ingress namespace is required")

agent/pkg/event/event.go

@@ -4,23 +4,24 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/google/uuid"
-
+	"github.com/sirupsen/logrus"
+	"github.com/tensorchord/openmodelz/agent/api/types"
 	"github.com/tensorchord/openmodelz/agent/client"
-	"github.com/tensorchord/openmodelz/agent/pkg/query"
 )
 
 type Interface interface {
 	CreateDeploymentEvent(namespace, deployment, event, message string) error
 }
 
 type EventRecorder struct {
-	DB query.Querier
+	Client     *client.Client
+	AgentToken string
 }
 
-func NewEventRecorder(q query.Querier) Interface {
+func NewEventRecorder(client *client.Client, token string) Interface {
 	return &EventRecorder{
-		DB: q,
+		Client:     client,
+		AgentToken: token,
 	}
 }
 
@@ -31,24 +32,18 @@ func (e *EventRecorder) CreateDeploymentEvent(namespace, deployment, event, mess
 	} else if user == "" {
 		return fmt.Errorf("user id is empty")
 	}
-	userId, err := uuid.Parse(user)
-	if err != nil {
-		return err
-	}
 
-	deploymentId, err := uuid.Parse(deployment)
+	deploymentEvent := types.DeploymentEvent{
+		UserID:       user,
+		DeploymentID: deployment,
+		EventType:    event,
+		Message:      message,
+	}
+	err = e.Client.CreateDeploymentEvent(context.TODO(), e.AgentToken, deploymentEvent)
 	if err != nil {
+		logrus.Errorf("failed to create deployment event: %v", err)
 		return err
 	}
 
-	params := query.CreateDeploymentEventParams{
-		UserID:       uuid.NullUUID{UUID: userId, Valid: true},
-		DeploymentID: uuid.NullUUID{UUID: deploymentId, Valid: true},
-		EventType:    NullStringBuilder(event, true),
-		Message:      NullStringBuilder(message, true),
-	}
-	if _, err := e.DB.CreateDeploymentEvent(context.TODO(), params); err != nil {
-		return err
-	}
 	return nil
 }