Make metric keys configurable (flyteorg#446)

* making metric keys configurable Signed-off-by: Daniel Rammer <daniel@union.ai> * fixed lint issue Signed-off-by: Daniel Rammer <daniel@union.ai> * go generated to update config flags Signed-off-by: Daniel Rammer <daniel@union.ai>
eapolinario · Jun 1, 2022 · bd4f6e2 · bd4f6e2
1 parent 774c457
commit bd4f6e2
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 28 deletions.
diff --git a/flytepropeller/cmd/controller/cmd/root.go b/flytepropeller/cmd/controller/cmd/root.go
@@ -7,27 +7,28 @@ import (
 	"os"
 	"runtime"
 
+	"github.com/flyteorg/flytepropeller/pkg/controller"
+	config2 "github.com/flyteorg/flytepropeller/pkg/controller/config"
 	"github.com/flyteorg/flytepropeller/pkg/controller/executors"
-	"sigs.k8s.io/controller-runtime/pkg/manager"
+	"github.com/flyteorg/flytepropeller/pkg/signals"
 
+	"github.com/flyteorg/flytestdlib/config"
+	"github.com/flyteorg/flytestdlib/config/viper"
+	"github.com/flyteorg/flytestdlib/contextutils"
+	"github.com/flyteorg/flytestdlib/logger"
 	"github.com/flyteorg/flytestdlib/profutils"
 	"github.com/flyteorg/flytestdlib/promutils"
-	"golang.org/x/sync/errgroup"
-	"k8s.io/klog"
-
-	config2 "github.com/flyteorg/flytepropeller/pkg/controller/config"
-
-	"github.com/flyteorg/flytestdlib/config/viper"
+	"github.com/flyteorg/flytestdlib/promutils/labeled"
 	"github.com/flyteorg/flytestdlib/version"
 
-	"github.com/flyteorg/flytestdlib/config"
-	"github.com/flyteorg/flytestdlib/logger"
+	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 
-	"github.com/spf13/cobra"
+	"golang.org/x/sync/errgroup"
 
-	"github.com/flyteorg/flytepropeller/pkg/controller"
-	"github.com/flyteorg/flytepropeller/pkg/signals"
+	"k8s.io/klog"
+
+	"sigs.k8s.io/controller-runtime/pkg/manager"
 )
 
 const (
@@ -108,6 +109,13 @@ func executeRootCmd(baseCtx context.Context, cfg *config2.Config) error {
 	// set up signals so we handle the first shutdown signal gracefully
 	ctx := signals.SetupSignalHandler(baseCtx)
 
+	// set metric keys
+	keys := contextutils.MetricKeysFromStrings(cfg.MetricKeys)
+	logger.Infof(context.TODO(), "setting metrics keys to %+v", keys)
+	if len(keys) > 0 {
+		labeled.SetMetricKeys(keys...)
+	}
+
 	// Add the propeller subscope because the MetricsPrefix only has "flyte:" to get uniform collection of metrics.
 	propellerScope := promutils.NewScope(cfg.MetricsPrefix).NewSubScope("propeller").NewSubScope(cfg.LimitNamespace)
 	limitNamespace := ""

diff --git a/flytepropeller/cmd/controller/cmd/webhook.go b/flytepropeller/cmd/controller/cmd/webhook.go
@@ -3,21 +3,24 @@ package cmd
 import (
 	"context"
 
-	"github.com/flyteorg/flytepropeller/pkg/controller/executors"
-	"sigs.k8s.io/controller-runtime/pkg/manager"
-
 	"github.com/flyteorg/flytepropeller/pkg/controller"
-	"github.com/flyteorg/flytestdlib/promutils"
-	"golang.org/x/sync/errgroup"
-
-	webhookConfig "github.com/flyteorg/flytepropeller/pkg/webhook/config"
-	"github.com/flyteorg/flytestdlib/profutils"
-
 	"github.com/flyteorg/flytepropeller/pkg/controller/config"
+	"github.com/flyteorg/flytepropeller/pkg/controller/executors"
 	"github.com/flyteorg/flytepropeller/pkg/signals"
 	"github.com/flyteorg/flytepropeller/pkg/webhook"
+	webhookConfig "github.com/flyteorg/flytepropeller/pkg/webhook/config"
+
+	"github.com/flyteorg/flytestdlib/contextutils"
 	"github.com/flyteorg/flytestdlib/logger"
+	"github.com/flyteorg/flytestdlib/profutils"
+	"github.com/flyteorg/flytestdlib/promutils"
+	"github.com/flyteorg/flytestdlib/promutils/labeled"
+
 	"github.com/spf13/cobra"
+
+	"golang.org/x/sync/errgroup"
+
+	"sigs.k8s.io/controller-runtime/pkg/manager"
 )
 
 var webhookCmd = &cobra.Command{
@@ -85,6 +88,13 @@ func runWebhook(origContext context.Context, propellerCfg *config.Config, cfg *w
 	// set up signals so we handle the first shutdown signal gracefully
 	ctx := signals.SetupSignalHandler(origContext)
 
+	// set metric keys
+	keys := contextutils.MetricKeysFromStrings(propellerCfg.MetricKeys)
+	logger.Infof(context.TODO(), "setting metrics keys to %+v", keys)
+	if len(keys) > 0 {
+		labeled.SetMetricKeys(keys...)
+	}
+
 	webhookScope := promutils.NewScope(cfg.MetricsPrefix).NewSubScope("webhook")
 	limitNamespace := ""
 	if propellerCfg.LimitNamespace != defaultNamespace {

diff --git a/flytepropeller/cmd/controller/main.go b/flytepropeller/cmd/controller/main.go
@@ -2,17 +2,10 @@ package main
 
 import (
 	_ "github.com/flyteorg/flytepropeller/plugins"
-	"github.com/flyteorg/flytestdlib/contextutils"
-	"github.com/flyteorg/flytestdlib/promutils/labeled"
 
 	"github.com/flyteorg/flytepropeller/cmd/controller/cmd"
 )
 
-func init() {
-	labeled.SetMetricKeys(contextutils.ProjectKey, contextutils.DomainKey, contextutils.WorkflowIDKey,
-		contextutils.TaskIDKey)
-}
-
 func main() {
 	cmd.Execute()
 }
diff --git a/flytepropeller/pkg/controller/config/config.go b/flytepropeller/pkg/controller/config/config.go
@@ -34,6 +34,7 @@ import (
 	"time"
 
 	"github.com/flyteorg/flytestdlib/config"
+	"github.com/flyteorg/flytestdlib/contextutils"
 	"k8s.io/apimachinery/pkg/types"
 )
 
@@ -105,6 +106,8 @@ var (
 		MetadataPrefix:      "metadata/propeller",
 		EnableAdminLauncher: true,
 		MetricsPrefix:       "flyte",
+		MetricKeys: []string{contextutils.ProjectKey.String(), contextutils.DomainKey.String(),
+			contextutils.WorkflowIDKey.String(), contextutils.TaskIDKey.String()},
 		EventConfig: EventConfig{
 			RawOutputPolicy: RawOutputPolicyReference,
 		},
@@ -128,6 +131,7 @@ type Config struct {
 	DefaultRawOutputPrefix string               `json:"rawoutput-prefix" pflag:",a fully qualified storage path of the form s3://flyte/abc/..., where all data sandboxes should be stored."`
 	Queue                  CompositeQueueConfig `json:"queue,omitempty" pflag:",Workflow workqueue configuration, affects the way the work is consumed from the queue."`
 	MetricsPrefix          string               `json:"metrics-prefix" pflag:",An optional prefix for all published metrics."`
+	MetricKeys             []string             `json:"metrics-keys" pflag:",Metrics labels applied to prometheus metrics emitted by the service."`
 	EnableAdminLauncher    bool                 `json:"enable-admin-launcher" pflag:"Enable remote Workflow launcher to Admin"`
 	MaxWorkflowRetries     int                  `json:"max-workflow-retries" pflag:"Maximum number of retries per workflow"`
 	MaxTTLInHours          int                  `json:"max-ttl-hours" pflag:"Maximum number of hours a completed workflow should be retained. Number between 1-23 hours"`

diff --git a/flytepropeller/pkg/controller/config/config_flags.go b/flytepropeller/pkg/controller/config/config_flags.go
diff --git a/flytepropeller/pkg/controller/config/config_flags_test.go b/flytepropeller/pkg/controller/config/config_flags_test.go