katalyst-core/cmd/katalyst-agent/app/enableagents.go at bdd283d2e500129a6a2b7f99841046ab2f848394 · JustinChengLZ/katalyst-core · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
/*
Copyright 2022 The Katalyst Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"sync"

	"k8s.io/apimachinery/pkg/util/sets"

	"github.com/kubewharf/katalyst-core/cmd/katalyst-agent/app/agent"
	"github.com/kubewharf/katalyst-core/cmd/katalyst-agent/app/agent/qrm"
	_ "github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/cpu"
	_ "github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/gpu"
	_ "github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/io"
	_ "github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/memory"
	_ "github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/network"
	"github.com/kubewharf/katalyst-core/pkg/agent/utilcomponent/periodicalhandler"
	phconsts "github.com/kubewharf/katalyst-core/pkg/agent/utilcomponent/periodicalhandler/consts"
)

// AgentStarter is used to start katalyst agents
type AgentStarter struct {
	Init      agent.InitFunc
	ExtraConf interface{}
}

// AgentsDisabledByDefault is the set of controllers which is disabled by default
var AgentsDisabledByDefault = sets.NewString()

// agentInitializers is used to store the initializing function for each agent
var agentInitializers sync.Map

func init() {
	agentInitializers.Store(agent.ReporterManagerAgent, AgentStarter{Init: agent.InitReporterManager})
	agentInitializers.Store(agent.EvictionManagerAgent, AgentStarter{Init: agent.InitEvictionManager})
	agentInitializers.Store(agent.QoSSysAdvisor, AgentStarter{Init: agent.InitSysAdvisor})
	agentInitializers.Store(phconsts.PeriodicalHandlerManagerName, AgentStarter{Init: periodicalhandler.NewPeriodicalHandlerManager})

	agentInitializers.Store(agent.ORMAgent, AgentStarter{Init: agent.InitORM})
	agentInitializers.Store(agent.AuditManager, AgentStarter{Init: agent.InitAuditManager})

	// qrm plugins are registered at top level of agent
	agentInitializers.Store(qrm.QRMPluginNameCPU, AgentStarter{Init: qrm.InitQRMCPUPlugins})
	agentInitializers.Store(qrm.QRMPluginNameMemory, AgentStarter{Init: qrm.InitQRMMemoryPlugins})
	agentInitializers.Store(qrm.QRMPluginNameNetwork, AgentStarter{Init: qrm.InitQRMNetworkPlugins})
	agentInitializers.Store(qrm.QRMPluginNameIO, AgentStarter{Init: qrm.InitQRMIOPlugins})
	agentInitializers.Store(qrm.QRMPluginNameGPU, AgentStarter{Init: qrm.InitQRMGPUPlugins})
}

// RegisterAgentInitializer is used to register user-defined agents
func RegisterAgentInitializer(name string, starter AgentStarter) {
	agentInitializers.Store(name, starter)
}

// GetAgentInitializers returns those agents with initialized functions
func GetAgentInitializers() map[string]AgentStarter {
	agents := make(map[string]AgentStarter)
	agentInitializers.Range(func(key, value interface{}) bool {
		agents[key.(string)] = value.(AgentStarter)
		return true
	})
	return agents
}