refactor(core): clean up define metric and add missing functionality

.circleci/config.yml

@@ -468,7 +468,7 @@ jobs:
  name: "Run nox tests: <<parameters.session_base>>"
  no_output_timeout: 10m
  command: "\
- nox -db uv -s '<<parameters.session_base>>-\
+ nox -s '<<parameters.session_base>>-\
  <<parameters.python_version_major>>.\
  <<parameters.python_version_minor>>\
  (core=<<parameters.core>>)'"
@@ -1581,12 +1581,12 @@ workflows:
  session_base: unit_tests
  codecov_flags: unit
 
- filters:
- branches:
- only:
- # - main
- - /^release-.*/
- - /^.*-ci-win$/
+ # filters:
+ #  branches:
+ #  only:
+ #  # - main
+ #  - /^release-.*/
+ #  - /^.*-ci-win$/
 
  matrix:
  parameters:

core/internal/runmetric/runmetric.go

@@ -0,0 +1,24 @@
+package runmetric
+
+import (
+ "github.com/wandb/wandb/core/pkg/service"
+ "google.golang.org/protobuf/proto"
+)
+
+type RunMetricDict = map[string]*service.MetricRecord
+
+func addMetric(
+ metric *service.MetricRecord,
+ key string,
+ target *RunMetricDict,
+) {
+ if metric.GetXControl().GetOverwrite() {
+ (*target)[key] = metric
+ return
+ }
+ if existingMetric, ok := (*target)[key]; ok {
+ proto.Merge(existingMetric, metric)
+ } else {
+ (*target)[key] = metric
+ }
+}

core/internal/runmetric/runmetric_handler.go

@@ -0,0 +1,103 @@
+package runmetric
+
+import (
+ "fmt"
+ "path/filepath"
+ "strings"
+
+ "github.com/wandb/wandb/core/pkg/service"
+ "google.golang.org/protobuf/proto"
+)
+
+// RunMetricHanlder is used to track the run metrics
+type RunMetricHanlder struct {
+ definedMetrics RunMetricDict
+ globMetrics RunMetricDict
+}
+
+// NewMetricHandler creates a new RunMetric
+func NewMetricHandler() *RunMetricHanlder {
+ return &RunMetricHanlder{
+ definedMetrics: make(RunMetricDict),
+ globMetrics: make(RunMetricDict),
+ }
+}
+
+func (runMetric *RunMetricHanlder) MatchMetric(key string) (*service.MetricRecord, bool) {
+ if key == "" {
+ return nil, false
+ }
+
+ // Skip internal metrics
+ if strings.HasPrefix(key, "_") {
+ return nil, false
+ }
+
+ if metric, ok := runMetric.definedMetrics[key]; ok {
+ return metric, true
+ }
+
+ if metric, ok := runMetric.findGlobMatch(key); ok {
+ metric.Name = key
+ metric.GlobName = ""
+ metric.Options.Defined = false
+ return metric, false
+ }
+
+ return nil, false
+}
+
+func (runMetric *RunMetricHanlder) findGlobMatch(key string) (*service.MetricRecord, bool) {
+ if key == "" {
+ return nil, false
+ }
+
+ for pattern, metric := range runMetric.globMetrics {
+ if match, err := filepath.Match(pattern, key); err != nil {
+ // h.logger.CaptureError("error matching metric", err)
+ continue
+ } else if match {
+ return proto.Clone(metric).(*service.MetricRecord), true
+ }
+ }
+
+ return nil, false
+}
+
+func (runMetric *RunMetricHanlder) AddStepMetric(metric *service.MetricRecord) *service.MetricRecord {
+ if metric.GetName() == "" {
+ return nil
+ }
+
+ if metric.GetStepMetric() == "" {
+ return nil
+ }
+
+ if _, ok := runMetric.definedMetrics[metric.GetStepMetric()]; ok {
+ return nil
+ }
+
+ stepMetric := &service.MetricRecord{
+ Name: metric.GetStepMetric(),
+ }
+
+ addMetric(stepMetric, stepMetric.GetName(), &runMetric.definedMetrics)
+
+ return stepMetric
+}
+
+func (runMetric *RunMetricHanlder) AddMetric(metric *service.MetricRecord) error {
+ // metric can have a glob name or a name
+ // TODO: replace glob-name/name with one-of field
+ switch {
+ case metric.GetGlobName() != "":
+ addMetric(metric, metric.GetGlobName(), &runMetric.globMetrics)
+ return nil
+ case metric.GetName() != "":
+ addMetric(metric, metric.GetName(), &runMetric.definedMetrics)
+ return nil
+ default:
+ err := fmt.Errorf("metric must have a name or glob name")
+ return err
+ }
+}

core/internal/runmetric/runmetric_sender.go

@@ -0,0 +1,59 @@
+package runmetric
+
+import (
+ "fmt"
+
+ "github.com/wandb/wandb/core/internal/corelib"
+ "github.com/wandb/wandb/core/pkg/service"
+ "google.golang.org/protobuf/proto"
+)
+
+type RunMetricSender struct {
+ metrics RunMetricDict
+ indeces map[string]int32
+ Config []map[int]interface{}
+}
+
+func NewRunMetricSender() *RunMetricSender {
+ return &RunMetricSender{
+ metrics: make(RunMetricDict),
+ indeces: make(map[string]int32),
+ Config: make([]map[int]interface{}, 0),
+ }
+}
+
+func (runMetric *RunMetricSender) EncodeConfigHints(metric *service.MetricRecord) error {
+ if metric.GetName() == "" {
+ err := fmt.Errorf("metric name is required")
+ return err
+ }
+
+ if metric.GetGlobName() != "" {
+ err := fmt.Errorf("glob name is not allowed")
+ return err
+ }
+
+ addMetric(metric, metric.GetName(), &runMetric.metrics)
+
+ // Check if the metric has a step metric
+ // If it does, we need to increment the step metric index
+ // and remove the step metric from the metric
+ if metric.GetStepMetric() != "" {
+ index, ok := runMetric.indeces[metric.GetStepMetric()]
+ if ok {
+ metric = proto.Clone(metric).(*service.MetricRecord)
+ metric.StepMetric = ""
+ metric.StepMetricIndex = index + 1
+ }
+ }
+
+ encoded := corelib.ProtoEncodeToDict(metric)
+ if index, ok := runMetric.indeces[metric.GetName()]; ok {
+ runMetric.Config[index] = encoded
+ } else {
+ next := int32(len(runMetric.Config))
+ runMetric.Config = append(runMetric.Config, encoded)
+ runMetric.indeces[metric.GetName()] = next
+ }
+ return nil
+}