Repository: dragonflyoss/Dragonfly2
Pull request #1480 (2022-07-24 09:35:17)

feat: scheduler add training config
<!--- Provide a general summary of your changes in the Title above -->

## Description

<!--- Describe your changes in detail -->
add config for machine learning
## Related Issue

<!--- This project only accepts pull requests related to open issues -->
<!--- If suggesting a new feature or change, please discuss it in an issue first -->
<!--- If fixing a bug, there should be an issue describing it with steps to reproduce -->
<!--- Please link to the issue here: -->

## Motivation and Context

<!--- Why is this change required? What problem does it solve? -->

## Screenshots (if appropriate)

## Types of changes

<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation Update (if none of the other choices apply)

## Checklist

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->

- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the **CONTRIBUTING** document.
- [ ] I have added tests to cover my changes.



Patch:
diff --git a/pkg/dag/dag.go b/pkg/dag/dag.go
index 1464c187..6969545b 100644
--- a/pkg/dag/dag.go
+++ b/pkg/dag/dag.go
@@ -54,6 +54,12 @@ type DAG interface {
 	// GetVertices returns map of vertices.
 	GetVertices() map[string]*Vertex
 
+	// GetSourceVertices returns source vertices.
+	GetSourceVertices() map[string]*Vertex
+
+	// GetSinkVertices returns sink vertices.
+	GetSinkVertices() map[string]*Vertex
+
 	// VertexCount returns count of vertices.
 	VertexCount() int
 
@@ -252,6 +258,36 @@ func (d *dag) DeleteEdge(fromVertexID, toVertexID string) error {
 	return nil
 }
 
+// GetSourceVertices returns source vertices.
+func (d *dag) GetSourceVertices() map[string]*Vertex {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+
+	sourceVertices := make(map[string]*Vertex)
+	for k, v := range d.vertices {
+		if v.InDegree() == 0 {
+			sourceVertices[k] = v
+		}
+	}
+
+	return sourceVertices
+}
+
+// GetSinkVertices returns sink vertices.
+func (d *dag) GetSinkVertices() map[string]*Vertex {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+
+	sinkVertices := make(map[string]*Vertex)
+	for k, v := range d.vertices {
+		if v.OutDegree() == 0 {
+			sinkVertices[k] = v
+		}
+	}
+
+	return sinkVertices
+}
+
 // depthFirstSearch is a depth-first search of the directed acyclic graph.
 func (d *dag) depthFirstSearch(fromVertexID, toVertexID string) bool {
 	successors := make(map[string]struct{})
diff --git a/pkg/dag/mocks/dag_mock.go b/pkg/dag/mocks/dag_mock.go
index 43db04ca..8e63e929 100644
--- a/pkg/dag/mocks/dag_mock.go
+++ b/pkg/dag/mocks/dag_mock.go
@@ -102,6 +102,34 @@ func (mr *MockDAGMockRecorder) DeleteVertex(id interface{}) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteVertex", reflect.TypeOf((*MockDAG)(nil).DeleteVertex), id)
 }
 
+// GetSinkVertices mocks base method.
+func (m *MockDAG) GetSinkVertices() map[string]*dag.Vertex {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetSinkVertices")
+	ret0, _ := ret[0].(map[string]*dag.Vertex)
+	return ret0
+}
+
+// GetSinkVertices indicates an expected call of GetSinkVertices.
+func (mr *MockDAGMockRecorder) GetSinkVertices() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSinkVertices", reflect.TypeOf((*MockDAG)(nil).GetSinkVertices))
+}
+
+// GetSourceVertices mocks base method.
+func (m *MockDAG) GetSourceVertices() map[string]*dag.Vertex {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetSourceVertices")
+	ret0, _ := ret[0].(map[string]*dag.Vertex)
+	return ret0
+}
+
+// GetSourceVertices indicates an expected call of GetSourceVertices.
+func (mr *MockDAGMockRecorder) GetSourceVertices() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSourceVertices", reflect.TypeOf((*MockDAG)(nil).GetSourceVertices))
+}
+
 // GetVertex mocks base method.
 func (m *MockDAG) GetVertex(id string) (*dag.Vertex, error) {
 	m.ctrl.T.Helper()
diff --git a/scheduler/config/config.go b/scheduler/config/config.go
index 9c23cc3d..1e80e79a 100644
--- a/scheduler/config/config.go
+++ b/scheduler/config/config.go
@@ -81,6 +81,12 @@ func New() *Config {
 				HostGCInterval: DefaultSchedulerHostGCInterval,
 				HostTTL:        DefaultSchedulerHostTTL,
 			},
+			Training: &TrainingConfig{
+				Enable:               false,
+				EnableAutoRefresh:    false,
+				RefreshModelInterval: DefaultRefreshModelInterval,
+				CPU:                  DefaultCPU,
+			},
 		},
 		DynConfig: &DynConfig{
 			RefreshInterval: DefaultDynConfigRefreshInterval,
@@ -172,6 +178,16 @@ func (cfg *Config) Validate() error {
 		return errors.New("scheduler requires parameter taskTTL")
 	}
 
+	if cfg.Scheduler.Training != nil && cfg.Scheduler.Training.Enable {
+		if cfg.Scheduler.Training.CPU <= 0 {
+			return errors.New("training requires parameter cpu")
+		}
+
+		if cfg.Scheduler.Training.EnableAutoRefresh && cfg.Scheduler.Training.RefreshModelInterval <= 0 {
+			return errors.New("training requires parameter refreshModelInterval")
+		}
+	}
+
 	if cfg.DynConfig.RefreshInterval <= 0 {
 		return errors.New("dynconfig requires parameter refreshInterval")
 	}
@@ -284,6 +300,23 @@ type SchedulerConfig struct {
 
 	// Task and peer gc configuration.
 	GC *GCConfig `yaml:"gc" mapstructure:"gc"`
+
+	// Training configuration.
+	Training *TrainingConfig `yaml:"training" mapstructure:"training"`
+}
+
+type TrainingConfig struct {
+	// Enable training.
+	Enable bool `yaml:"enable" mapstructure:"enable"`
+
+	// Enable auto refresh model.
+	EnableAutoRefresh bool `yaml:"enableAutoRefresh" mapstructure:"enableAutoRefresh"`
+
+	// RefreshModelInterval is refresh interval for refreshing model.
+	RefreshModelInterval time.Duration `yaml:"refreshModelInterval" mapstructure:"refreshModelInterval"`
+
+	// CPU limit while training.
+	CPU int `yaml:"cpu" mapstructure:"cpu"`
 }
 
 type GCConfig struct {
diff --git a/scheduler/config/constants.go b/scheduler/config/constants.go
index d6d482a3..0a0d978b 100644
--- a/scheduler/config/constants.go
+++ b/scheduler/config/constants.go
@@ -76,6 +76,12 @@ const (
 
 	// DefaultSchedulerHostTTL is default ttl for host.
 	DefaultSchedulerHostTTL = 48 * time.Hour
+
+	// DefaultRefreshModelInterval is model refresh interval.
+	DefaultRefreshModelInterval = 168 * time.Hour
+
+	// DefaultCPU is default cpu usage.
+	DefaultCPU = 1
 )
 
 const (


Test patch:
diff --git a/pkg/dag/dag_test.go b/pkg/dag/dag_test.go
index 4e89a3df..3097fa7a 100644
--- a/pkg/dag/dag_test.go
+++ b/pkg/dag/dag_test.go
@@ -619,6 +619,100 @@ func TestDAGDeleteEdge(t *testing.T) {
 	}
 }
 
+func TestDAGSourceVertices(t *testing.T) {
+	tests := []struct {
+		name   string
+		expect func(t *testing.T, d DAG)
+	}{
+		{
+			name: "get source vertices",
+			expect: func(t *testing.T, d DAG) {
+				assert := assert.New(t)
+				var (
+					mockVertexEID = "bae"
+					mockVertexFID = "baf"
+				)
+				if err := d.AddVertex(mockVertexEID, mockVertexValue); err != nil {
+					assert.NoError(err)
+				}
+
+				if err := d.AddVertex(mockVertexFID, mockVertexValue); err != nil {
+					assert.NoError(err)
+				}
+
+				if err := d.AddEdge(mockVertexEID, mockVertexFID); err != nil {
+					assert.NoError(err)
+				}
+
+				sourceVertices := d.GetSourceVertices()
+				assert.Equal(len(sourceVertices), 1)
+				assert.Equal(sourceVertices[mockVertexEID].Value, mockVertexValue)
+			},
+		},
+		{
+			name: "source vertices not found",
+			expect: func(t *testing.T, d DAG) {
+				assert := assert.New(t)
+				sourceVertices := d.GetSourceVertices()
+				assert.Equal(len(sourceVertices), 0)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			d := NewDAG()
+			tc.expect(t, d)
+		})
+	}
+}
+
+func TestDAGSinkVertices(t *testing.T) {
+	tests := []struct {
+		name   string
+		expect func(t *testing.T, d DAG)
+	}{
+		{
+			name: "get sink vertices",
+			expect: func(t *testing.T, d DAG) {
+				assert := assert.New(t)
+				var (
+					mockVertexEID = "bae"
+					mockVertexFID = "baf"
+				)
+				if err := d.AddVertex(mockVertexEID, mockVertexValue); err != nil {
+					assert.NoError(err)
+				}
+
+				if err := d.AddVertex(mockVertexFID, mockVertexValue); err != nil {
+					assert.NoError(err)
+				}
+
+				if err := d.AddEdge(mockVertexEID, mockVertexFID); err != nil {
+					assert.NoError(err)
+				}
+
+				sinkVertices := d.GetSinkVertices()
+				assert.Equal(len(sinkVertices), 1)
+				assert.Equal(sinkVertices[mockVertexFID].Value, mockVertexValue)
+			},
+		},
+		{
+			name: "sink vertices not found",
+			expect: func(t *testing.T, d DAG) {
+				assert := assert.New(t)
+				sinkVertices := d.GetSinkVertices()
+				assert.Equal(len(sinkVertices), 0)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			d := NewDAG()
+			tc.expect(t, d)
+		})
+	}
+}
+
 func BenchmarkDAGAddVertex(b *testing.B) {
 	var ids []string
 	d := NewDAG()
diff --git a/pkg/dag/vertex_test.go b/pkg/dag/vertex_test.go
index bf72b5f3..5c063b85 100644
--- a/pkg/dag/vertex_test.go
+++ b/pkg/dag/vertex_test.go
@@ -24,7 +24,7 @@ import (
 
 const (
 	mockVertexID    = "foo"
-	mockVertexValue = "var"
+	mockVertexValue = "bar"
 )
 
 func TestNewVertex(t *testing.T) {
diff --git a/scheduler/config/config_test.go b/scheduler/config/config_test.go
index ef210b29..49ef5034 100644
--- a/scheduler/config/config_test.go
+++ b/scheduler/config/config_test.go
@@ -48,6 +48,12 @@ func TestConfig_Load(t *testing.T) {
 				HostGCInterval: 1 * time.Minute,
 				HostTTL:        10 * time.Minute,
 			},
+			Training: &TrainingConfig{
+				Enable:               true,
+				EnableAutoRefresh:    true,
+				RefreshModelInterval: 1 * time.Second,
+				CPU:                  2,
+			},
 		},
 		Server: &ServerConfig{
 			IP:       "127.0.0.1",
@@ -140,6 +146,12 @@ func TestConfig_New(t *testing.T) {
 				HostGCInterval: 30 * time.Minute,
 				HostTTL:        48 * time.Hour,
 			},
+			Training: &TrainingConfig{
+				Enable:               false,
+				EnableAutoRefresh:    false,
+				RefreshModelInterval: 168 * time.Hour,
+				CPU:                  1,
+			},
 		},
 		DynConfig: &DynConfig{
 			RefreshInterval: 10 * time.Second,
diff --git a/scheduler/config/testdata/scheduler.yaml b/scheduler/config/testdata/scheduler.yaml
index 4baea596..d8431a9b 100644
--- a/scheduler/config/testdata/scheduler.yaml
+++ b/scheduler/config/testdata/scheduler.yaml
@@ -20,6 +20,11 @@ scheduler:
     taskTTL: 600000000000
     hostGCInterval: 60000000000
     hostTTL: 600000000000
+  training:
+    enable: true
+    enableAutoRefresh: true
+    refreshModelInterval: 1000000000
+    cpu: 2
 
 dynconfig:
   refreshInterval: 300000000000