Skip to content
This repository has been archived by the owner on Jul 17, 2024. It is now read-only.

Commit

Permalink
feat: add telemetry (#103)
Browse files Browse the repository at this point in the history
Signed-off-by: Keming <[email protected]>
  • Loading branch information
kemingy authored Aug 3, 2023
1 parent 23b18f2 commit 700fb86
Show file tree
Hide file tree
Showing 12 changed files with 215 additions and 11 deletions.
4 changes: 2 additions & 2 deletions mdz/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -109,12 +109,12 @@ export GOFLAGS ?= -count=1

.DEFAULT_GOAL:=build

build: build-local ## Build the release version of envd
build: build-release ## Build the release version

help: ## Display this help
@awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n make \033[36m<target>\033[0m\n"} /^[a-zA-Z0-9_-]+:.*?##/ { printf " \033[36m%-15s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST)

debug: debug-local ## Build the debug version of envd
debug: debug-local ## Build the debug version

# more info about `GOGC` env: https://github.com/golangci/golangci-lint#memory-usage-of-golangci-lint
lint: $(GOLANGCI_LINT) ## Lint GO code
Expand Down
2 changes: 2 additions & 0 deletions mdz/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,8 @@ require (
github.com/rivo/uniseg v0.2.0 // indirect
github.com/rogpeppe/go-internal v1.10.0 // indirect
github.com/russross/blackfriday/v2 v2.1.0 // indirect
github.com/segmentio/analytics-go/v3 v3.2.1 // indirect
github.com/segmentio/backo-go v1.0.0 // indirect
github.com/spf13/pflag v1.0.5 // indirect
github.com/swaggo/swag v1.8.12 // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
Expand Down
4 changes: 4 additions & 0 deletions mdz/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,10 @@ github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjR
github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
github.com/segmentio/analytics-go/v3 v3.2.1 h1:G+f90zxtc1p9G+WigVyTR0xNfOghOGs/PYAlljLOyeg=
github.com/segmentio/analytics-go/v3 v3.2.1/go.mod h1:p8owAF8X+5o27jmvUognuXxdtqvSGtD0ZrfY2kcS9bE=
github.com/segmentio/backo-go v1.0.0 h1:kbOAtGJY2DqOR0jfRkYEorx/b18RgtepGtY3+Cpe6qA=
github.com/segmentio/backo-go v1.0.0/go.mod h1:kJ9mm9YmoWSkk+oQ+5Cj8DEoRCX2JT6As4kEtIIOp1M=
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
Expand Down
7 changes: 7 additions & 0 deletions mdz/pkg/cmd/deploy.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
petname "github.com/dustinkirkland/golang-petname"
"github.com/spf13/cobra"
"github.com/tensorchord/openmodelz/agent/api/types"
"github.com/tensorchord/openmodelz/mdz/pkg/telemetry"
)

var (
Expand Down Expand Up @@ -110,6 +111,12 @@ func commandDeploy(cmd *cobra.Command, args []string) error {
}
}

telemetry.GetTelemetry().Record(
"deploy",
telemetry.AddField("GPU", deployGPU),
telemetry.AddField("FromZero", deployMinReplicas == 0),
)

if _, err := agentClient.InferenceCreate(
cmd.Context(), namespace, inf); err != nil {
cmd.PrintErrf("Failed to create the inference: %s\n", errors.Cause(err))
Expand Down
2 changes: 2 additions & 0 deletions mdz/pkg/cmd/list.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"github.com/jedib0t/go-pretty/v6/table"
"github.com/spf13/cobra"
"github.com/tensorchord/openmodelz/agent/api/types"
"github.com/tensorchord/openmodelz/mdz/pkg/telemetry"
)

const (
Expand Down Expand Up @@ -48,6 +49,7 @@ func init() {
}

func commandList(cmd *cobra.Command, args []string) error {
telemetry.GetTelemetry().Record("list")
infs, err := agentClient.InferenceList(cmd.Context(), namespace)
if err != nil {
cmd.PrintErrf("Failed to list inferences: %v\n", err)
Expand Down
14 changes: 11 additions & 3 deletions mdz/pkg/cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,15 @@ import (
"github.com/spf13/cobra/doc"

"github.com/tensorchord/openmodelz/agent/client"
"github.com/tensorchord/openmodelz/mdz/pkg/telemetry"
)

var (
// Used for flags.
mdzURL string
namespace string
debug bool
mdzURL string
namespace string
debug bool
disableTelemetry bool

agentClient *client.Client
)
Expand Down Expand Up @@ -62,6 +64,8 @@ func init() {

rootCmd.PersistentFlags().BoolVarP(&debug, "debug", "", false, "Enable debug logging")

rootCmd.PersistentFlags().BoolVarP(&disableTelemetry, "disable-telemetry", "", false, "Disable anonymous telemetry")

// Cobra also supports local flags, which will only run
// when this action is called directly.
rootCmd.AddGroup(&cobra.Group{ID: "basic", Title: "Basic Commands:"})
Expand Down Expand Up @@ -89,6 +93,10 @@ func commandInit(cmd *cobra.Command, args []string) error {
return err
}
}

if err := telemetry.Initialize(!disableTelemetry); err != nil {
logrus.WithError(err).Debug("Failed to initialize telemetry")
}
return nil
}

Expand Down
3 changes: 3 additions & 0 deletions mdz/pkg/cmd/scale.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package cmd

import (
"github.com/spf13/cobra"
"github.com/tensorchord/openmodelz/mdz/pkg/telemetry"
)

var (
Expand Down Expand Up @@ -73,6 +74,8 @@ func commandScale(cmd *cobra.Command, args []string) error {
deployment.Spec.Scaling.TargetLoad = int32Ptr(targetInflightRequests)
}

telemetry.GetTelemetry().Record("scale")

if _, err := agentClient.DeploymentUpdate(cmd.Context(), namespace, deployment); err != nil {
cmd.PrintErrf("Failed to update deployment: %s\n", err)
return err
Expand Down
5 changes: 4 additions & 1 deletion mdz/pkg/cmd/server_join.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"github.com/spf13/cobra"

"github.com/tensorchord/openmodelz/mdz/pkg/server"
"github.com/tensorchord/openmodelz/mdz/pkg/telemetry"
)

// serverJoinCmd represents the server join command
Expand Down Expand Up @@ -46,10 +47,12 @@ func commandServerJoin(cmd *cobra.Command, args []string) error {
},
})
if err != nil {
cmd.PrintErrf("Failed to join the cluster: %s\n", errors.Cause(err))
cmd.PrintErrf("Failed to configure before join: %s\n", errors.Cause(err))
return err
}

telemetry.GetTelemetry().Record("server join")

_, err = engine.Run()
if err != nil {
cmd.PrintErrf("Failed to join the cluster: %s\n", errors.Cause(err))
Expand Down
2 changes: 2 additions & 0 deletions mdz/pkg/cmd/server_list.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"github.com/jedib0t/go-pretty/v6/table"
"github.com/spf13/cobra"
"github.com/tensorchord/openmodelz/agent/api/types"
"github.com/tensorchord/openmodelz/mdz/pkg/telemetry"
)

var (
Expand Down Expand Up @@ -40,6 +41,7 @@ func init() {
}

func commandServerList(cmd *cobra.Command, args []string) error {
telemetry.GetTelemetry().Record("server list")
servers, err := agentClient.ServerList(cmd.Context())
if err != nil {
cmd.PrintErrf("Failed to list servers: %s\n", errors.Cause(err))
Expand Down
7 changes: 7 additions & 0 deletions mdz/pkg/cmd/server_start.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (

"github.com/tensorchord/openmodelz/agent/pkg/consts"
"github.com/tensorchord/openmodelz/mdz/pkg/server"
"github.com/tensorchord/openmodelz/mdz/pkg/telemetry"
"github.com/tensorchord/openmodelz/mdz/pkg/version"
)

Expand Down Expand Up @@ -60,6 +61,12 @@ func commandServerStart(cmd *cobra.Command, args []string) error {
domainWithSuffix := fmt.Sprintf("%s.%s", args[0], serverStartDomain)
domain = &domainWithSuffix
}
defer func(start time.Time) {
telemetry.GetTelemetry().Record(
"server start",
telemetry.AddField("duration", time.Since(start).Seconds()),
)
}(time.Now())
engine, err := server.NewStart(server.Options{
Verbose: serverVerbose,
Runtime: server.Runtime(serverStartRuntime),
Expand Down
166 changes: 166 additions & 0 deletions mdz/pkg/telemetry/telemetry.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
package telemetry

import (
"io"
"os"
"path/filepath"
"runtime"
"sync"
"time"

"github.com/cockroachdb/errors"
"github.com/google/uuid"
segmentio "github.com/segmentio/analytics-go/v3"
"github.com/sirupsen/logrus"

"github.com/tensorchord/openmodelz/mdz/pkg/version"
)

type TelemetryField func(*segmentio.Properties)

type Telemetry interface {
Record(command string, args ...TelemetryField)
}

type defaultTelemetry struct {
client segmentio.Client
uid string
enabled bool
}

const telemetryToken = "65WHA9bxCNX74K3HjgplMOmsio9LkYSI"

var (
once sync.Once
telemetry *defaultTelemetry
telemetryConfigFile string
)

func init() {
home, err := os.UserHomeDir()
if err != nil {
panic(err)
}
telemetryConfigFile = filepath.Join(home, ".config", "openmodelz", "telemetry")
}

func GetTelemetry() Telemetry {
return telemetry
}

func Initialize(enabled bool) error {
once.Do(func() {
client, err := segmentio.NewWithConfig(telemetryToken, segmentio.Config{
BatchSize: 1,
})
if err != nil {
panic(err)
}
telemetry = &defaultTelemetry{
client: client,
enabled: enabled,
}
})
return telemetry.init()
}

func (t *defaultTelemetry) init() error {
if !t.enabled {
return nil
}
// detect if the config file already exists
_, err := os.Stat(telemetryConfigFile)
if err != nil {
if !os.IsNotExist(err) {
return errors.Wrap(err, "failed to stat telemetry config file")
}
t.uid = uuid.New().String()
return t.dumpConfig()
}
if err = t.loadConfig(); err != nil {
return errors.Wrap(err, "failed to load telemetry config")
}
t.Idnetify()
return nil
}

func (t *defaultTelemetry) dumpConfig() error {
if err := os.MkdirAll(filepath.Dir(telemetryConfigFile), os.ModeDir|0700); err != nil {
return errors.Wrap(err, "failed to create telemetry config directory")
}
file, err := os.Create(telemetryConfigFile)
if err != nil {
return errors.Wrap(err, "failed to create telemetry config file")
}
defer file.Close()
_, err = file.WriteString(t.uid)
if err != nil {
return errors.Wrap(err, "failed to write telemetry config file")
}
return nil
}

func (t *defaultTelemetry) loadConfig() error {
file, err := os.Open(telemetryConfigFile)
if err != nil {
return errors.Wrap(err, "failed to open telemetry config file")
}
defer file.Close()
uid, err := io.ReadAll(file)
if err != nil {
return errors.Wrap(err, "failed to read telemetry config file")
}
t.uid = string(uid)
return nil
}

func (t *defaultTelemetry) Idnetify() {
if !t.enabled {
return
}
v := version.GetOpenModelzVersion()
if err := t.client.Enqueue(segmentio.Identify{
UserId: t.uid,
Context: &segmentio.Context{
OS: segmentio.OSInfo{
Name: runtime.GOOS,
Version: runtime.GOARCH,
},
App: segmentio.AppInfo{
Name: "openmodelz",
Version: v,
},
},
Timestamp: time.Now(),
Traits: segmentio.NewTraits(),
}); err != nil {
logrus.WithError(err).Debug("failed to identify user")
return
}
}

func AddField(name string, value interface{}) TelemetryField {
return func(p *segmentio.Properties) {
p.Set(name, value)
}
}

func (t *defaultTelemetry) Record(command string, fields ...TelemetryField) {
if !t.enabled {
return
}
logrus.WithField("UID", t.uid).WithField("command", command).Debug("send telemetry")
track := segmentio.Track{
UserId: t.uid,
Event: command,
Properties: segmentio.NewProperties(),
}
for _, field := range fields {
field(&track.Properties)
}
if err := t.client.Enqueue(track); err != nil {
logrus.WithError(err).Debug("failed to send telemetry")
}
// make sure the msg can be sent out
t.client.Close()
}
10 changes: 5 additions & 5 deletions mdz/pkg/version/version.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ var (
developmentFlag = "false"
)

// Version contains envd version information
// Version contains OpenModelz version information
type Version struct {
Version string
BuildDate string
Expand All @@ -65,8 +65,8 @@ func SetGitTagForE2ETest(tag string) {
gitTag = tag
}

// GetEnvdVersion gets Envd version information
func GetEnvdVersion() string {
// GetOpenModelzVersion gets OpenModelz version information
func GetOpenModelzVersion() string {
var versionStr string

if gitCommit != "" && gitTag != "" &&
Expand Down Expand Up @@ -97,7 +97,7 @@ func GetEnvdVersion() string {
// GetVersion returns the version information
func GetVersion() Version {
return Version{
Version: GetEnvdVersion(),
Version: GetOpenModelzVersion(),
BuildDate: buildDate,
GitCommit: gitCommit,
GitTag: gitTag,
Expand Down Expand Up @@ -128,5 +128,5 @@ func UserAgent() string {
version = matches[0][1] + "-dev"
}

return "envd/" + version
return "modelz/" + version
}

0 comments on commit 700fb86

Please sign in to comment.