Skip to content

Commit

Permalink
Make default cluster label configurable (flyteorg#86)
Browse files Browse the repository at this point in the history
* Make default cluster label configurable

* Unit tests
  • Loading branch information
EngHabu authored May 11, 2020
1 parent 5b8724a commit 49b9c77
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 8 deletions.
2 changes: 2 additions & 0 deletions flyteplugins/go/tasks/plugins/hive/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ var (
TokenKey: "FLYTE_QUBOLE_CLIENT_TOKEN",
LruCacheSize: 2000,
Workers: 15,
DefaultClusterLabel: "default",
ClusterConfigs: []ClusterConfig{{PrimaryLabel: "default", Labels: []string{"default"}, Limit: 100, ProjectScopeQuotaProportionCap: 0.7, NamespaceScopeQuotaProportionCap: 0.7}},
DestinationClusterConfigs: []DestinationClusterConfig{},
}
Expand All @@ -62,6 +63,7 @@ type Config struct {
TokenKey string `json:"quboleTokenKey" pflag:",Name of the key where to find Qubole token in the secret manager."`
LruCacheSize int `json:"lruCacheSize" pflag:",Size of the AutoRefreshCache"`
Workers int `json:"workers" pflag:",Number of parallel workers to refresh the cache"`
DefaultClusterLabel string `json:"defaultClusterConfig" pflag:",The default cluster label. This will be used if label is not specified on the hive job."`
ClusterConfigs []ClusterConfig `json:"clusterConfigs" pflag:"-,A list of cluster configs. Each of the configs corresponds to a service cluster"`
DestinationClusterConfigs []DestinationClusterConfig `json:"destinationClusterConfigs" pflag:"-,A list configs specifying the destination service cluster for (project, domain)"`
}
Expand Down
19 changes: 11 additions & 8 deletions flyteplugins/go/tasks/plugins/hive/execution_state.go
Original file line number Diff line number Diff line change
Expand Up @@ -259,13 +259,13 @@ func GetQueryInfo(ctx context.Context, tCtx core.TaskExecutionContext) (
return
}

func mapLabelToPrimaryLabel(ctx context.Context, quboleCfg *config.Config, label string) (string, bool) {
primaryLabel := DefaultClusterPrimaryLabel
found := false
func mapLabelToPrimaryLabel(ctx context.Context, quboleCfg *config.Config, label string) (primaryLabel string, found bool) {
primaryLabel = quboleCfg.DefaultClusterLabel
found = false

if label == "" {
logger.Debugf(ctx, "Input cluster label is an empty string; falling back to using the default primary label [%v]", label, DefaultClusterPrimaryLabel)
return primaryLabel, found
logger.Debugf(ctx, "Input cluster label is an empty string; falling back to using the default primary label [%v]", label, primaryLabel)
return
}

// Using a linear search because N is small and because of ClusterConfig's struct definition
Expand All @@ -280,8 +280,11 @@ func mapLabelToPrimaryLabel(ctx context.Context, quboleCfg *config.Config, label
}
}

logger.Debugf(ctx, "Cannot find the primary cluster label for label [%v] in configmap; "+
"falling back to using the default primary label [%v]", label, DefaultClusterPrimaryLabel)
if !found {
logger.Debugf(ctx, "Cannot find the primary cluster label for label [%v] in configmap; "+
"falling back to using the default primary label [%v]", label, primaryLabel)
}

return primaryLabel, found
}

Expand Down Expand Up @@ -319,7 +322,7 @@ func getClusterPrimaryLabel(ctx context.Context, tCtx core.TaskExecutionContext,
}

// Else we return the default primary label
return DefaultClusterPrimaryLabel
return cfg.DefaultClusterLabel
}

func KickOffQuery(ctx context.Context, tCtx core.TaskExecutionContext, currentState ExecutionState, quboleClient client.QuboleClient,
Expand Down
1 change: 1 addition & 0 deletions flyteplugins/go/tasks/plugins/hive/execution_state_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -348,6 +348,7 @@ func TestKickOffQuery(t *testing.T) {

func createMockQuboleCfg() *config.Config {
return &config.Config{
DefaultClusterLabel: "default",
ClusterConfigs: []config.ClusterConfig{
{PrimaryLabel: "primary A", Labels: []string{"primary A", "A", "label A", "A-prod"}, Limit: 10},
{PrimaryLabel: "primary B", Labels: []string{"B"}, Limit: 10},
Expand Down

0 comments on commit 49b9c77

Please sign in to comment.