wandb

1
package server
2

3
import (
4
	"errors"
5
	"path/filepath"
6

7
	"github.com/wandb/wandb/core/internal/corelib"
8
	"github.com/wandb/wandb/core/pkg/service"
9
	"google.golang.org/protobuf/proto"
10
)
11

12
type MetricHandler struct {
13
	definedMetrics map[string]*service.MetricRecord
14
	globMetrics    map[string]*service.MetricRecord
15
}
16

17
func NewMetricHandler() *MetricHandler {
18
	return &MetricHandler{
19
		definedMetrics: make(map[string]*service.MetricRecord),
20
		globMetrics:    make(map[string]*service.MetricRecord),
21
	}
22
}
23

24
// addMetric adds a metric to the target map. If the metric already exists, it will be merged
25
// with the existing metric. If the overwrite flag is set, the metric will be overwritten.
26
func addMetric(arg interface{}, key string, target *map[string]*service.MetricRecord) (*service.MetricRecord, error) {
27
	var metric *service.MetricRecord
28

29
	switch v := arg.(type) {
30
	case string:
31
		metric = &service.MetricRecord{
32
			Name: v,
33
		}
34
	case *service.MetricRecord:
35
		metric = v
36
	default:
37
		// Handle invalid input
38
		return nil, errors.New("invalid input")
39
	}
40

41
	if metric.GetXControl().GetOverwrite() {
42
		(*target)[key] = metric
43
	} else {
44
		if existingMetric, ok := (*target)[key]; ok {
45
			proto.Merge(existingMetric, metric)
46
		} else {
47
			(*target)[key] = metric
48
		}
49
	}
50
	return metric, nil
51
}
52

53
// createMatchingGlobMetric check if a key matches a glob pattern, if it does create a new defined metric
54
// based on the glob metric and return it.
55
func (mh *MetricHandler) createMatchingGlobMetric(key string) *service.MetricRecord {
56
	for pattern, globMetric := range mh.globMetrics {
57
		if match, err := filepath.Match(pattern, key); err != nil {
58
			// h.logger.CaptureError("error matching metric", err)
59
			continue
60
		} else if match {
61
			metric := proto.Clone(globMetric).(*service.MetricRecord)
62
			metric.Name = key
63
			metric.Options.Defined = false
64
			metric.GlobName = ""
65
			return metric
66
		}
67
	}
68
	return nil
69
}
70

71
// handleStepMetric handles the step metric for a given metric key. If the step metric is not
72
// defined, it will be added to the defined metrics map.
73
func (h *Handler) handleStepMetric(key string) {
74
	if key == "" {
75
		return
76
	}
77

78
	// already exists no need to add
79
	if _, defined := h.metricHandler.definedMetrics[key]; defined {
80
		return
81
	}
82

83
	metric, err := addMetric(key, key, &h.metricHandler.definedMetrics)
84

85
	if err != nil {
86
		h.logger.CaptureError("error adding metric to map", err)
87
		return
88
	}
89

90
	stepRecord := &service.Record{
91
		RecordType: &service.Record_Metric{
92
			Metric: metric,
93
		},
94
		Control: &service.Control{
95
			Local: true,
96
		},
97
	}
98
	h.sendRecord(stepRecord)
99
}
100

101
func (h *Handler) handleMetric(record *service.Record, metric *service.MetricRecord) {
102
	// metric can have a glob name or a name
103
	// TODO: replace glob-name/name with one-of field
104
	switch {
105
	case metric.GetGlobName() != "":
106
		if _, err := addMetric(metric, metric.GetGlobName(), &h.metricHandler.globMetrics); err != nil {
107
			h.logger.CaptureError("error adding metric to map", err)
108
			return
109
		}
110
		h.sendRecord(record)
111
	case metric.GetName() != "":
112
		if _, err := addMetric(metric, metric.GetName(), &h.metricHandler.definedMetrics); err != nil {
113
			h.logger.CaptureError("error adding metric to map", err)
114
			return
115
		}
116
		h.handleStepMetric(metric.GetStepMetric())
117
		h.sendRecord(record)
118
	default:
119
		h.logger.CaptureError("invalid metric", errors.New("invalid metric"))
120
	}
121
}
122

123
type MetricSender struct {
124
	definedMetrics map[string]*service.MetricRecord
125
	metricIndex    map[string]int32
126
	configMetrics  []map[int]interface{}
127
}
128

129
func NewMetricSender() *MetricSender {
130
	return &MetricSender{
131
		definedMetrics: make(map[string]*service.MetricRecord),
132
		metricIndex:    make(map[string]int32),
133
		configMetrics:  make([]map[int]interface{}, 0),
134
	}
135
}
136

137
// encodeMetricHints encodes the metric hints for the given metric record. The metric hints
138
// are used to configure the plots in the UI.
139
func (s *Sender) encodeMetricHints(_ *service.Record, metric *service.MetricRecord) {
140

141
	_, err := addMetric(metric, metric.GetName(), &s.metricSender.definedMetrics)
142
	if err != nil {
143
		return
144
	}
145

146
	if metric.GetStepMetric() != "" {
147
		index, ok := s.metricSender.metricIndex[metric.GetStepMetric()]
148
		if ok {
149
			metric = proto.Clone(metric).(*service.MetricRecord)
150
			metric.StepMetric = ""
151
			metric.StepMetricIndex = index + 1
152
		}
153
	}
154

155
	encodeMetric := corelib.ProtoEncodeToDict(metric)
156
	if index, ok := s.metricSender.metricIndex[metric.GetName()]; ok {
157
		s.metricSender.configMetrics[index] = encodeMetric
158
	} else {
159
		nextIndex := len(s.metricSender.configMetrics)
160
		s.metricSender.configMetrics = append(s.metricSender.configMetrics, encodeMetric)
161
		s.metricSender.metricIndex[metric.GetName()] = int32(nextIndex)
162
	}
163
}
164