Repository: intelligent-machine-learning/dlrover (language: python)

Issue:
Directory design of EasyDL
1. How to organize directories of EasyDL?
```
|-brain # Automaticallly generates the resource plan of the job.
|-operator
  |-controllers
    |-elastic-job.   # Creates a k8s Job
    |-resource-scale  # Scale out or in the job resource according to the Custom Resource(CR)
|-elasticdl   # dispatches data shards to workers and monitors training nodes.
|-easydl  # APIs for the training loop of TensorFlow/Pytorch to use elastic training.
```

2. Which framework do we use to implement a trainer to support elastic training?
We need a trainer to catch the exception and rebuild the session if parameter servers change. Now, a trainer is implemented with `tf.estimator` framework in AntGroup. However, Keras is more common than `tf.estimator` and TF 2.x has supported training a Keras model using [ParameterServerStrategy](https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/ParameterServerStrategy). In another way, we can implement a trainer based on `tf.estimator` and convert a Keras model to an estimator model in TensorFlow.
 
 



Discussion / hints:

Fix patch:
diff --git a/dlrover/examples/iris_job.yaml b/dlrover/examples/iris_job.yaml
index a2ced95d7..450e9bf8d 100644
--- a/dlrover/examples/iris_job.yaml
+++ b/dlrover/examples/iris_job.yaml
@@ -2,6 +2,7 @@ apiVersion: elastic.iml.github.io/v1alpha1
 kind: ElasticJob
 metadata:
   name: elasticjob-sample
+  namespace: dlrover
 spec:
   distributionStrategy: ParameterServerStrategy
   optimizeMode: cluster
@@ -18,10 +19,12 @@ spec:
             - name: main
               image: easydl/tf-estimator:iris_dnn_v0
               imagePullPolicy: IfNotPresent
+              # yamllint disable
               command:
                 - /bin/bash
                 - -c
                 - "python -m model_zoo.tf_estimator.iris_dnn_elastic"
+              # yamllint enable
     worker:
       replicas: 0
       restartCount: 3
diff --git a/dlrover/go/brain/cmd/brain/main.go b/dlrover/go/brain/cmd/brain/main.go
index 93ced81de..45a1ee153 100644
--- a/dlrover/go/brain/cmd/brain/main.go
+++ b/dlrover/go/brain/cmd/brain/main.go
@@ -13,66 +13,74 @@
 
 package main
 
+import (
+	"context"
+	"flag"
+	log "github.com/golang/glog"
+	"github.com/intelligent-machine-learning/easydl/brain/pkg/common"
+	"github.com/intelligent-machine-learning/easydl/brain/pkg/config"
+	pb "github.com/intelligent-machine-learning/easydl/brain/pkg/proto"
+	"github.com/intelligent-machine-learning/easydl/brain/pkg/server"
+	"google.golang.org/grpc"
+	kubeclientset "k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/tools/clientcmd"
+	"net"
+)
+
 func main() {
-	//log.Info("Start DLRover Brain")
-	//flag.Parse()
-	//mConfig := config.CommandConfig
-	//
-	//ctx, cancel := context.WithCancel(context.Background())
-	//defer cancel()
-	//
-	//cfg, err := clientcmd.BuildConfigFromFlags(mConfig.APIServer, mConfig.KubeConfig)
-	//if err != nil {
-	//	log.Fatalf("Error building kube config: %v", err)
-	//}
+	log.Info("Start DLRover Brain")
+	flag.Parse()
+	mConfig := config.CommandConfig
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-	//kubeClient, err := kubeclientset.NewForConfig(cfg)
-	//if err != nil {
-	//	log.Fatalf("Error building kube client: %v", err.Error())
-	//}
+	cfg, err := clientcmd.BuildConfigFromFlags(mConfig.APIServer, mConfig.KubeConfig)
+	if err != nil {
+		log.Fatalf("Error building kube config: %v", err)
+	}
+
+	kubeClient, err := kubeclientset.NewForConfig(cfg)
+	if err != nil {
+		log.Fatalf("Error building kube client: %v", err.Error())
+	}
 
 	// Start error handler
-	//log.Info("Start error handler")
-	//errHandler, err := common.NewStopAllErrorHandler(cancel)
-	//if err != nil {
-	//	log.Fatalf("Create ErrorHandler error: %v", err)
-	//} else {
-	//	go errHandler.HandleError(ctx)
-	//}
-	//
-	//log.Infof("serviceConfigMapName=%s, serviceConfigMapKey=%s", mConfig.ServiceConfigMapName, mConfig.ServiceConfigMapKey)
-	//configManager := config.NewManager(mConfig.Namespace, mConfig.ServiceConfigMapName, mConfig.ServiceConfigMapKey, kubeClient)
-	//err = configManager.Run(ctx, errHandler)
-	//if err != nil {
-	//	log.Fatalf("Fail to run the config manager: %v", err)
-	//}
-	//
-	//conf, err := configManager.GetConfig()
-	//if err != nil {
-	//	log.Fatalf("Fail to get config: %v", err)
-	//}
-	//conf.Set(config.KubeClientInterface, kubeClient)
-	//conf.Set(config.Namespace, mConfig.Namespace)
-	//
-	//lis, err := net.Listen("tcp", mConfig.Port)
-	//if err != nil {
-	//	log.Fatalf("failed to listen: %v", err)
-	//}
-	//s := grpc.NewServer()
-	//
-	//brainServer, err := server.NewBrainServer(conf)
-	//if err != nil {
-	//	log.Fatalf("fail to create Brain server: %v", err)
-	//}
-	//
-	//err = brainServer.Run(ctx, errHandler)
-	//if err != nil {
-	//	log.Fatalf("Fail to run EasyDL server: %v", err)
-	//}
-	//
-	//pb.RegisterBrainServer(s, brainServer)
-	//log.Infof("server listening at %v", lis.Addr())
-	//if err = s.Serve(lis); err != nil {
-	//	log.Fatalf("failed to serve: %v", err)
-	//}
+	log.Info("Start error handler")
+	errHandler, err := common.NewStopAllErrorHandler(cancel)
+	if err != nil {
+		log.Fatalf("Create ErrorHandler error: %v", err)
+	} else {
+		go errHandler.HandleError(ctx)
+	}
+
+	log.Infof("namespace=%s, serviceConfigMapName=%s, serviceConfigMapKey=%s", mConfig.Namespace,
+		mConfig.ServiceConfigMapName, mConfig.ServiceConfigMapKey)
+	conf := config.NewEmptyConfig()
+	conf.Set(config.KubeClientInterface, kubeClient)
+	conf.Set(config.BrainServerConfigMapName, mConfig.ServiceConfigMapName)
+	conf.Set(config.BrainServerConfigMapKey, mConfig.ServiceConfigMapKey)
+	conf.Set(config.Namespace, mConfig.Namespace)
+
+	lis, err := net.Listen("tcp", mConfig.Port)
+	if err != nil {
+		log.Fatalf("failed to listen: %v", err)
+	}
+	s := grpc.NewServer()
+
+	brainServer, err := server.NewBrainServer(conf)
+	if err != nil {
+		log.Fatalf("fail to create Brain server: %v", err)
+	}
+
+	err = brainServer.Run(ctx, errHandler)
+	if err != nil {
+		log.Fatalf("Fail to run EasyDL server: %v", err)
+	}
+
+	pb.RegisterBrainServer(s, brainServer)
+	log.Infof("server listening at %v", lis.Addr())
+	if err = s.Serve(lis); err != nil {
+		log.Fatalf("failed to serve: %v", err)
+	}
 }
diff --git a/dlrover/go/brain/cmd/k8smonitor/main.go b/dlrover/go/brain/cmd/k8smonitor/main.go
index 5396b0d76..92a43a828 100644
--- a/dlrover/go/brain/cmd/k8smonitor/main.go
+++ b/dlrover/go/brain/cmd/k8smonitor/main.go
@@ -53,16 +53,6 @@ func main() {
 
 	log.Infof("namespace=%s, serviceConfigMapName=%s, serviceConfigMapKey=%s", mConfig.Namespace,
 		mConfig.ServiceConfigMapName, mConfig.ServiceConfigMapKey)
-	//configManager := config.NewManager(mConfig.Namespace, mConfig.ServiceConfigMapName, mConfig.ServiceConfigMapKey, kubeClient)
-	//err = configManager.Run(ctx, errHandler)
-	//if err != nil {
-	//	log.Fatalf("Fail to run the config manager: %v", err)
-	//}
-	//
-	//conf, err := configManager.GetConfig()
-	//if err != nil {
-	//	log.Fatalf("Fail to get config: %v", err)
-	//}
 
 	conf := config.NewEmptyConfig()
 	conf.Set(config.KubeClientInterface, kubeClient)
diff --git a/dlrover/go/brain/manifests/k8s/brain-service-configmap.yaml b/dlrover/go/brain/manifests/k8s/brain-service-configmap.yaml
new file mode 100644
index 000000000..d75864249
--- /dev/null
+++ b/dlrover/go/brain/manifests/k8s/brain-service-configmap.yaml
@@ -0,0 +1,11 @@
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: dlrover-bain-service
+  namespace: dlrover
+data:
+  config: |
+    namespace: "dlrover"
+    data-store.config-map.name: "dlrover-data-store"
+    data-store.config-map.key: "datastore"
diff --git a/dlrover/go/brain/pkg/config/command_config.go b/dlrover/go/brain/pkg/config/command_config.go
index 8ac864d85..87ac3b25d 100644
--- a/dlrover/go/brain/pkg/config/command_config.go
+++ b/dlrover/go/brain/pkg/config/command_config.go
@@ -26,6 +26,8 @@ const (
 	SpecServiceConfigMapKey = "serviceConfigMapKey"
 	// SpecNamespaceName is the spec name of namespace
 	SpecNamespaceName = "namespace"
+	// SpecServerPort is the port of the EasyDL server gRPC service
+	SpecServerPort = "port"
 )
 
 // Spec is the struct of configure specifications
@@ -40,6 +42,8 @@ type Spec struct {
 	ServiceConfigMapKey string
 	// Namespace is the namespace
 	Namespace string
+	// Port is the port of the EasyDL server gRPC service
+	Port string
 }
 
 // CommandConfig is the variable of type ConfigSpec
@@ -66,4 +70,9 @@ func init() {
 		SpecNamespaceName,
 		"",
 		"The namespace name used for k8s ConfigMap & LeaderElection")
+	flag.StringVar(
+		&CommandConfig.Port,
+		SpecServerPort,
+		":50001",
+		"Port of the EasyDL server gRPC service")
 }
diff --git a/dlrover/go/brain/pkg/config/common.go b/dlrover/go/brain/pkg/config/common.go
index 39300df2d..975f295ca 100644
--- a/dlrover/go/brain/pkg/config/common.go
+++ b/dlrover/go/brain/pkg/config/common.go
@@ -55,6 +55,11 @@ const (
 	// KubeWatcherHandlerConfigMapKey is the key of kube watcher handler config map
 	KubeWatcherHandlerConfigMapKey = "kube-watcher.handler.config-map.key"
 
+	// BrainServerConfigMapName is the name of brain server config map
+	BrainServerConfigMapName = "brain.server.config-map.name"
+	// BrainServerConfigMapKey is the name of brain server config key
+	BrainServerConfigMapKey = "brain.server.config-map.key"
+
 	// QueryBackwardTimePeriodInHour is the config key for query backward time period in hour
 	QueryBackwardTimePeriodInHour = "query.backward.time-period.hour"
 
diff --git a/dlrover/go/brain/pkg/datastore/dbbase/recorder.go b/dlrover/go/brain/pkg/datastore/dbbase/recorder.go
index 1d42c672b..fe988687f 100644
--- a/dlrover/go/brain/pkg/datastore/dbbase/recorder.go
+++ b/dlrover/go/brain/pkg/datastore/dbbase/recorder.go
@@ -83,7 +83,7 @@ func (r *DatabaseRecorder) Get(row interface{}, condition Condition) error {
 	}
 	session := r.Table(r.TableName)
 	session = condition.Apply(session)
-	found, err := session.Desc("id").Get(row)
+	found, err := session.Get(row)
 	if err != nil {
 		log.Errorf("Failed to get %v of %+v: %v", r.TableName, condition, err)
 		return err
diff --git a/dlrover/go/brain/pkg/datastore/implementation/elasticjob_datastore.go b/dlrover/go/brain/pkg/datastore/implementation/elasticjob_datastore.go
index d27696f18..23f1723d3 100644
--- a/dlrover/go/brain/pkg/datastore/implementation/elasticjob_datastore.go
+++ b/dlrover/go/brain/pkg/datastore/implementation/elasticjob_datastore.go
@@ -18,6 +18,7 @@ import (
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/common"
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/config"
 	datastoreapi "github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/api"
+	dsimplutils "github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/implementation/utils"
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/recorder/mysql"
 )
 
@@ -65,5 +66,5 @@ func (store *ElasticJobDataStore) PersistData(condition *datastoreapi.Condition,
 
 // GetData returns data for a given condition
 func (store *ElasticJobDataStore) GetData(condition *datastoreapi.Condition, data interface{}) error {
-	return nil
+	return dsimplutils.GetData(store.client, condition, data)
 }
diff --git a/dlrover/go/brain/pkg/datastore/manager.go b/dlrover/go/brain/pkg/datastore/manager.go
index e5154a936..747504a84 100644
--- a/dlrover/go/brain/pkg/datastore/manager.go
+++ b/dlrover/go/brain/pkg/datastore/manager.go
@@ -21,6 +21,7 @@ import (
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/config"
 	datastoreapi "github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/api"
 	dsimpl "github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/implementation"
+	"sync"
 )
 
 // Manager is the struct of data store manager
@@ -28,6 +29,7 @@ type Manager struct {
 	conf          *config.Config
 	configManager *config.Manager
 	dataStores    map[string]datastoreapi.DataStore
+	mutex         *sync.RWMutex
 }
 
 // NewManager creates a data store manager
@@ -41,6 +43,8 @@ func NewManager(conf *config.Config) *Manager {
 
 	return &Manager{
 		configManager: configManager,
+		dataStores:    make(map[string]datastoreapi.DataStore),
+		mutex:         &sync.RWMutex{},
 	}
 }
 
@@ -63,11 +67,24 @@ func (m *Manager) Run(ctx context.Context, errReporter common.ErrorReporter) err
 
 // CreateDataStore creates a data store for a given name
 func (m *Manager) CreateDataStore(name string) (datastoreapi.DataStore, error) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if dataStore, found := m.dataStores[name]; found {
+		return dataStore, nil
+	}
+
 	conf := m.conf.GetConfig(name)
 	if conf == nil {
 		err := fmt.Errorf("There is no config for data store %s", name)
 		return nil, err
 	}
 
-	return dsimpl.CreateDataStore(name, conf)
+	dataStore, err := dsimpl.CreateDataStore(name, conf)
+	if err != nil {
+		log.Errorf("fail to create data store %s: %v", name, err)
+		return nil, err
+	}
+	m.dataStores[name] = dataStore
+	return dataStore, nil
 }
diff --git a/dlrover/go/brain/pkg/platform/k8s/common/elastic_job.go b/dlrover/go/brain/pkg/platform/k8s/common/elastic_job.go
new file mode 100644
index 000000000..912d407b3
--- /dev/null
+++ b/dlrover/go/brain/pkg/platform/k8s/common/elastic_job.go
@@ -0,0 +1,23 @@
+// Copyright 2023 The DLRover Authors. All rights reserved.
+// 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 common
+
+const (
+	// JobPhaseRunning indicates job is running
+	JobPhaseRunning = "Running"
+	// JobPhaseError indicates job runs into errors
+	JobPhaseError = "Error"
+	// JobPhaseSucceeded indicates job completed successfully
+	JobPhaseSucceeded = "Succeeded"
+)
diff --git a/dlrover/go/brain/pkg/platform/k8s/implementation/watchhandler/elasticjob_handler.go b/dlrover/go/brain/pkg/platform/k8s/implementation/watchhandler/elasticjob_handler.go
index 0fb477bb1..4992334c1 100644
--- a/dlrover/go/brain/pkg/platform/k8s/implementation/watchhandler/elasticjob_handler.go
+++ b/dlrover/go/brain/pkg/platform/k8s/implementation/watchhandler/elasticjob_handler.go
@@ -19,6 +19,7 @@ import (
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/config"
 	datastoreapi "github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/api"
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/recorder/mysql"
+	k8scommon "github.com/intelligent-machine-learning/easydl/brain/pkg/platform/k8s/common"
 	handlerutils "github.com/intelligent-machine-learning/easydl/brain/pkg/platform/k8s/implementation/watchhandler/utils"
 	watchercommon "github.com/intelligent-machine-learning/easydl/brain/pkg/platform/k8s/watcher/common"
 	elasticv1alpha1 "github.com/intelligent-machine-learning/easydl/dlrover/go/operator/api/v1alpha1"
@@ -85,7 +86,43 @@ func (handler *ElasticJobHandler) HandleCreateEvent(object runtime.Object, event
 
 // HandleUpdateEvent handles update events
 func (handler *ElasticJobHandler) HandleUpdateEvent(object runtime.Object, oldObject runtime.Object, event watchercommon.Event) error {
-	return nil
+	newJob := &elasticv1alpha1.ElasticJob{}
+	oldJob := &elasticv1alpha1.ElasticJob{}
+
+	unstructNewObj := object.(*unstructured.Unstructured)
+	unstructOldObj := oldObject.(*unstructured.Unstructured)
+
+	runtime.DefaultUnstructuredConverter.FromUnstructured(unstructOldObj.Object, oldJob)
+	runtime.DefaultUnstructuredConverter.FromUnstructured(unstructNewObj.Object, newJob)
+
+	jobUID := string(newJob.UID)
+
+	cond := &datastoreapi.Condition{
+		Type: common.TypeGetDataGetJob,
+		Extra: &mysql.JobCondition{
+			UID: jobUID,
+		},
+	}
+	record := &mysql.Job{}
+
+	err := handler.dataStore.GetData(cond, record)
+	if err != nil {
+		log.Errorf("fail to get job record for %s: %v", jobUID, err)
+		return err
+	}
+
+	if newJob.Status.Phase == k8scommon.JobPhaseRunning && record.StartedAt.IsZero() {
+		record.StartedAt = newJob.Status.StartTime.Time
+	}
+
+	if newJob.Status.Phase == k8scommon.JobPhaseSucceeded && record.FinishedAt.IsZero() {
+		record.FinishedAt = newJob.Status.CompletionTime.Time
+	}
+
+	upsertCond := &datastoreapi.Condition{
+		Type: common.TypeUpsertJob,
+	}
+	return handler.dataStore.PersistData(upsertCond, record, nil)
 }
 
 // HandleDeleteEvent handles delete events
diff --git a/dlrover/go/brain/pkg/server/server.go b/dlrover/go/brain/pkg/server/server.go
index 7f8044d02..41aab45a3 100644
--- a/dlrover/go/brain/pkg/server/server.go
+++ b/dlrover/go/brain/pkg/server/server.go
@@ -16,6 +16,7 @@ package server
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	log "github.com/golang/glog"
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/common"
@@ -25,36 +26,72 @@ import (
 	dsimpl "github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/implementation"
 	"github.com/intelligent-machine-learning/easydl/brain/pkg/datastore/recorder/mysql"
 	pb "github.com/intelligent-machine-learning/easydl/brain/pkg/proto"
+	"k8s.io/client-go/kubernetes"
 )
 
 const (
 	defaultDataStoreName = dsimpl.BaseDataStoreName
+	logName              = "brain-server"
 )
 
 // BrainServer is the interface of DLRover Brain
 type BrainServer struct {
 	pb.UnimplementedBrainServer
-	dsManager  *datastore.Manager
-	dataStores map[string]datastoreapi.DataStore
+
+	kubeClientSet kubernetes.Interface
+
+	conf          *config.Config
+	configManager *config.Manager
+
+	dsManager *datastore.Manager
 }
 
 // NewBrainServer creates an EasyDLServer instance
 func NewBrainServer(conf *config.Config) (*BrainServer, error) {
-	dsManager := datastore.NewManager(conf)
+	namespace := conf.GetString(config.Namespace)
+	configMapName := conf.GetString(config.BrainServerConfigMapName)
+	configMapKey := conf.GetString(config.BrainServerConfigMapKey)
+	kubeClientSet := conf.GetKubeClientInterface()
+
 	return &BrainServer{
-		dsManager:  dsManager,
-		dataStores: make(map[string]datastoreapi.DataStore),
+		configManager: config.NewManager(namespace, configMapName, configMapKey, kubeClientSet),
+		kubeClientSet: kubeClientSet,
 	}, nil
 }
 
 // Run starts the server
 func (s *BrainServer) Run(ctx context.Context, errReporter common.ErrorReporter) error {
+	err := s.configManager.Run(ctx, errReporter)
+	if err != nil {
+		err = fmt.Errorf("[%s] failed to initialize config manager: %v", logName, err)
+		log.Error(err)
+		return err
+	}
+	s.conf, err = s.configManager.GetConfig()
+	if err != nil {
+		log.Errorf("[%s] fail to get brain server config: %v", logName, err)
+		return err
+	}
+	log.Infof("[%s] brain server config: %v", logName, s.conf)
+
+	dsConf := config.NewEmptyConfig()
+	dsConf.Set(config.KubeClientInterface, s.kubeClientSet)
+	dsConf.Set(config.DataStoreConfigMapName, s.conf.GetString(config.DataStoreConfigMapName))
+	dsConf.Set(config.DataStoreConfigMapKey, s.conf.GetString(config.DataStoreConfigMapKey))
+	dsConf.Set(config.Namespace, s.conf.GetString(config.Namespace))
+
+	s.dsManager = datastore.NewManager(dsConf)
+	err = s.dsManager.Run(ctx, errReporter)
+	if err != nil {
+		log.Errorf("[%s] fail to run the data store manager: %v", logName, err)
+		return err
+	}
 	return s.dsManager.Run(ctx, errReporter)
 }
 
 // PersistMetrics persists job metrics to data store
 func (s *BrainServer) PersistMetrics(ctx context.Context, in *pb.JobMetrics) (*empty.Empty, error) {
-	dataStore, err := s.getDataStore(in.DataStore)
+	dataStore, err := s.dsManager.CreateDataStore(in.DataStore)
 	if err != nil {
 		return nil, err
 	}
@@ -65,18 +102,6 @@ func (s *BrainServer) PersistMetrics(ctx context.Context, in *pb.JobMetrics) (*e
 	return &empty.Empty{}, nil
 }
 
-func (s *BrainServer) getDataStore(name string) (datastoreapi.DataStore, error) {
-	_, found := s.dataStores[name]
-	if !found {
-		dataStore, err := s.dsManager.CreateDataStore(name)
-		if err != nil {
-			return nil, err
-		}
-		s.dataStores[name] = dataStore
-	}
-	return s.dataStores[name], nil
-}
-
 // Optimize returns the initial resource of a job.
 func (s *BrainServer) Optimize(ctx context.Context, in *pb.OptimizeRequest) (*pb.OptimizeResponse, error) {
 	return nil, nil
@@ -84,7 +109,7 @@ func (s *BrainServer) Optimize(ctx context.Context, in *pb.OptimizeRequest) (*pb
 
 // GetJobMetrics returns a job metrics
 func (s *BrainServer) GetJobMetrics(ctx context.Context, in *pb.JobMetricsRequest) (*pb.JobMetricsResponse, error) {
-	dataStore, err := s.getDataStore(defaultDataStoreName)
+	dataStore, err := s.dsManager.CreateDataStore(defaultDataStoreName)
 	if err != nil {
 		log.Errorf("fail to get data store %s", defaultDataStoreName)
 		return &pb.JobMetricsResponse{


Test patch:
diff --git a/dlrover/go/brain/pkg/datastore/manager_test.go b/dlrover/go/brain/pkg/datastore/manager_test.go
new file mode 100644
index 000000000..e730792fc
--- /dev/null
+++ b/dlrover/go/brain/pkg/datastore/manager_test.go
@@ -0,0 +1,14 @@
+// Copyright 2023 The DLRover Authors. All rights reserved.
+// 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 datastore