Extend Katib API for NAS jobs

cmd/manager/main.go

@@ -29,34 +29,59 @@ type server struct {
 }
 
 func (s *server) CreateStudy(ctx context.Context, in *api_pb.CreateStudyRequest) (*api_pb.CreateStudyReply, error) {
+	var studyID string
+	var err error
 	if in == nil || in.StudyConfig == nil {
 		return &api_pb.CreateStudyReply{}, errors.New("StudyConfig is missing.")
 	}
-	studyID, err := dbIf.CreateStudy(in.StudyConfig)
-	if err != nil {
-		return &api_pb.CreateStudyReply{}, err
+
+	if in.StudyConfig.JobType != "nas" {
+		//If it is a HP job
+		studyID, err = dbIf.CreateStudy(in.StudyConfig)
+		if err != nil {
+			return &api_pb.CreateStudyReply{}, err
+		}
+		s.SaveStudy(ctx, &api_pb.SaveStudyRequest{
+			StudyName:   in.StudyConfig.Name,
+			Owner:       in.StudyConfig.Owner,
+			Description: "StudyID: " + studyID,
+		})
+
+	} else {
+		//If it is a NAS job
+		studyID = "1"
+		log.Printf("NAS job is created, id= %v", studyID)
 	}
-	s.SaveStudy(ctx, &api_pb.SaveStudyRequest{
-		StudyName:   in.StudyConfig.Name,
-		Owner:       in.StudyConfig.Owner,
-		Description: "StudyID: " + studyID,
-	})
 	return &api_pb.CreateStudyReply{StudyId: studyID}, nil
 }
 
 func (s *server) DeleteStudy(ctx context.Context, in *api_pb.DeleteStudyRequest) (*api_pb.DeleteStudyReply, error) {
 	if in == nil || in.StudyId == "" {
 		return &api_pb.DeleteStudyReply{}, errors.New("StudyId is missing.")
 	}
-	err := dbIf.DeleteStudy(in.StudyId)
-	if err != nil {
-		return &api_pb.DeleteStudyReply{}, err
+	if in.JobType != "nas" {
+		//If it is a HP job
+		err := dbIf.DeleteStudy(in.StudyId)
+		if err != nil {
+			return &api_pb.DeleteStudyReply{}, err
+		}
+	} else {
+		//If it is a NAS job
+		log.Printf("NAS job is deleted, id= %v", in.StudyId)
 	}
 	return &api_pb.DeleteStudyReply{StudyId: in.StudyId}, nil
 }
 
 func (s *server) GetStudy(ctx context.Context, in *api_pb.GetStudyRequest) (*api_pb.GetStudyReply, error) {
-	sc, err := dbIf.GetStudyConfig(in.StudyId)
+	var sc *api_pb.StudyConfig
+	var err error
+	if in.JobType != "nas" {
+		//If it is a HP job
+		sc, err = dbIf.GetStudyConfig(in.StudyId)
+	} else {
+		//If it is a NAS job
+		log.Printf("Get NAS job, id = %v", in.StudyId)
+	}
 	return &api_pb.GetStudyReply{StudyConfig: sc}, err
 }
 
@@ -91,8 +116,8 @@ func (s *server) GetTrials(ctx context.Context, in *api_pb.GetTrialsRequest) (*a
 }
 
 func (s *server) GetTrial(ctx context.Context, in *api_pb.GetTrialRequest) (*api_pb.GetTrialReply, error) {
-        t, err := dbIf.GetTrial(in.TrialId)
-        return &api_pb.GetTrialReply{Trial: t}, err
+	t, err := dbIf.GetTrial(in.TrialId)
+	return &api_pb.GetTrialReply{Trial: t}, err
 }
 
 func (s *server) GetSuggestions(ctx context.Context, in *api_pb.GetSuggestionsRequest) (*api_pb.GetSuggestionsReply, error) {
@@ -217,11 +242,17 @@ func (s *server) GetWorkerFullInfo(ctx context.Context, in *api_pb.GetWorkerFull
 func (s *server) SetSuggestionParameters(ctx context.Context, in *api_pb.SetSuggestionParametersRequest) (*api_pb.SetSuggestionParametersReply, error) {
 	var err error
 	var id string
-	if in.ParamId == "" {
-		id, err = dbIf.SetSuggestionParam(in.SuggestionAlgorithm, in.StudyId, in.SuggestionParameters)
+	if in.JobType != "nas" {
+		//If it is a HP job
+		if in.ParamId == "" {
+			id, err = dbIf.SetSuggestionParam(in.SuggestionAlgorithm, in.StudyId, in.SuggestionParameters)
+		} else {
+			id = in.ParamId
+			err = dbIf.UpdateSuggestionParam(in.ParamId, in.SuggestionParameters)
+		}
 	} else {
-		id = in.ParamId
-		err = dbIf.UpdateSuggestionParam(in.ParamId, in.SuggestionParameters)
+		//If it is a NAS job
+		log.Printf("Set suggestion Params for NAS job id= %v", in.StudyId)
 	}
 	return &api_pb.SetSuggestionParametersReply{ParamId: id}, err
 }

cmd/suggestion/nasrl/Dockerfile

@@ -0,0 +1,8 @@
+FROM python:3
+
+ADD . /usr/src/app/github.com/kubeflow/katib
+WORKDIR /usr/src/app/github.com/kubeflow/katib/cmd/suggestion/nasrl
+RUN pip install --no-cache-dir -r requirements.txt
+ENV PYTHONPATH /usr/src/app/github.com/kubeflow/katib:/usr/src/app/github.com/kubeflow/katib/pkg/api/python
+
+ENTRYPOINT ["python", "-u", "main.py"]

cmd/suggestion/nasrl/main.py

@@ -0,0 +1,29 @@
+import grpc
+from concurrent import futures
+
+import time
+
+from pkg.api.python import api_pb2_grpc
+from pkg.suggestion.nasrl_service import NasrlService
+from pkg.suggestion.types import DEFAULT_PORT
+from logging import getLogger, StreamHandler, INFO, DEBUG
+
+
+_ONE_DAY_IN_SECONDS = 60 * 60 * 24
+
+
+def serve():
+    print("NAS RL Suggestion Service")
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+    api_pb2_grpc.add_SuggestionServicer_to_server(NasrlService(), server)
+    server.add_insecure_port(DEFAULT_PORT)
+    print("Listening...")
+    server.start()
+    try:
+        while True:
+            time.sleep(_ONE_DAY_IN_SECONDS)
+    except KeyboardInterrupt:
+        server.stop(0)
+
+if __name__ == "__main__":
+    serve()

cmd/suggestion/nasrl/requirements.txt

@@ -0,0 +1,9 @@
+grpcio
+duecredit
+cloudpickle==0.5.6
+numpy>=1.13.3
+scikit-learn>=0.19.0
+scipy>=0.19.1
+forestci
+protobuf
+googleapis-common-protos

examples/nasjob-example-RL.yaml

@@ -0,0 +1,142 @@
+apiVersion: "kubeflow.org/v1alpha1"
+kind: StudyJob
+metadata:
+  namespace: kubeflow
+  labels:
+    controller-tools.k8s.io: "1.0"
+  name: nas-rl-example
+spec:
+  studyName: nas-rl-example
+  owner: crd
+  optimizationtype: maximize
+  objectivevaluename: Validation-accuracy
+  optimizationgoal: 0.99
+  requestcount: 4
+  metricsnames:
+    - accuracy
+  jobType: nas
+  graphConfig:
+    numLayers: 8
+    inputSize:
+      - 32
+      - 32
+      - 3
+    outputSize:
+      - 10
+  operations:
+    - operationType: convolution
+      parameterconfigs:
+        - name: filter_size
+          parametertype: categorical
+          feasible:
+            list:
+            - "3"
+            - "5"
+            - "7"
+        - name: num_filter
+          parametertype: categorical
+          feasible:
+            list:
+            - "32"
+            - "48"
+            - "64"
+            - "96"
+            - "128"
+        - name: stride
+          parametertype: categorical
+          feasible:
+            list:
+            - "1"
+            - "2"
+        - name: activation
+          parametertype: categorical
+          feasible:
+            list:
+            - relu
+            - sigmoid
+            - tanh
+            - identity
+    - operationType: reduction
+      parameterconfigs:
+        - name: reduction_type
+          parametertype: categorical
+          feasible:
+            list:
+            - maxpooling
+            - avgpooling
+        - name: pool_size
+          parametertype: int
+          feasible:
+            min: "1"
+            max: "3"
+        - name: stride
+          parametertype: categorical
+          feasible:
+            list:
+            - "1"
+            - "2"
+    - operationType: dense
+      parameterconfigs:
+        - name: num_units
+          parametertype: categorical
+          feasible:
+            list:
+              - "128"
+              - "256"
+              - "512"
+              - "1024"
+        - name: activation
+          parametertype: categorical
+          feasible:
+            list:
+            - relu
+            - sigmoid
+            - tanh
+            - identity
+        - name: dropout_rate
+          parametertype: range
+          feasible:
+            min: "0.1"
+            max: "0.6"
+            step: "0.1"
+  workerSpec:
+    goTemplate:
+        rawTemplate: |-
+          apiVersion: batch/v1
+          kind: Job
+          metadata:
+            name: {{.WorkerID}}
+            namespace: kubeflow
+          spec:
+            template:
+              spec:
+                imagePullSecrets:
+                  - name: gitlabregcred
+                containers:
+                - name: {{.WorkerID}}
+                  image: katib/mxnet-mnist-example
+                  command:
+                  - "python"
+                  - "/mxnet/example/image-classification/train_mnist.py"
+                  - "--batch-size=32"
+                  {{- with .HyperParameters}}
+                  {{- range .}}
+                  - "{{.Name}}={{.Value}}"
+                  {{- end}}
+                  {{- end}}
+                restartPolicy: Never
+  suggestionSpec:
+    suggestionAlgorithm: "nasrl"
+    suggestionParameters:
+      - name: "num_control_cells"
+        value: "32"
+      - name: "embedding_dim"
+        value: "20"
+      - name: "accuracy_beta"
+        value: "0.8"
+      - name: "clip_rewards"
+        value: "0"
+      - name: "embedding_dim"
+        value: "20"
+      - name: "discount_factor"
+        value: "0.95"