add common package (#491)

kubeflow · May 9, 2019 · 41143e8 · 41143e8
1 parent a38806c
commit 41143e8
Show file tree

Hide file tree

Showing 10 changed files with 49 additions and 18 deletions.
diff --git a/pkg/types.go → pkg/common/v1alpha1/common.go b/pkg/types.go → pkg/common/v1alpha1/common.go
@@ -1,4 +1,4 @@
-package pkg
+package v1alpha1
 
 import (
 	"os"

diff --git a/pkg/common/v1alpha2/common.go b/pkg/common/v1alpha2/common.go
@@ -0,0 +1,31 @@
+package v1alpha2
+
+import (
+	"os"
+
+	experimentsv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/experiment/v1alpha2"
+)
+
+const (
+	KatibManagerServiceIPEnvName        = "KATIB_MANAGER_PORT_6789_TCP_ADDR"
+	KatibManagerServicePortEnvName      = "KATIB_MANAGER_PORT_6789_TCP_PORT"
+	KatibManagerServiceNamespaceEnvName = "KATIB_MANAGER_NAMESPACE"
+	KatibManagerService                 = "katib-manager"
+	KatibManagerPort                    = "6789"
+	ManagerAddr                   = KatibManagerService + ":" + KatibManagerPort
+)
+
+func GetManagerAddr() string {
+	ns := os.Getenv(experimentsv1alpha2.DefaultKatibNamespaceEnvName)
+	if len(ns) == 0 {
+		addr := os.Getenv(KatibManagerServiceIPEnvName)
+		port := os.Getenv(KatibManagerServicePortEnvName)
+		if len(addr) > 0 && len(port) > 0 {
+			return addr + ":" + port
+		} else {
+			return ManagerAddr
+		}
+	} else {
+		return KatibManagerService + "." + ns + ":" + KatibManagerPort
+	}
+}
diff --git a/pkg/controller/v1alpha1/studyjob/katib_api_util.go b/pkg/controller/v1alpha1/studyjob/katib_api_util.go
@@ -21,7 +21,7 @@ import (
 	"log"
 	"reflect"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha1"
 	katibv1alpha1 "github.com/kubeflow/katib/pkg/api/operators/apis/studyjob/v1alpha1"
 	katibapi "github.com/kubeflow/katib/pkg/api/v1alpha1"
 	"google.golang.org/grpc"
@@ -37,7 +37,7 @@ func initializeStudy(instance *katibv1alpha1.StudyJob) error {
 
 	instance.Status.Condition = katibv1alpha1.ConditionRunning
 
-	conn, err := grpc.Dial(pkg.ManagerAddr, grpc.WithInsecure())
+	conn, err := grpc.Dial(common.ManagerAddr, grpc.WithInsecure())
 	if err != nil {
 		log.Printf("Connect katib manager error %v", err)
 		instance.Status.Condition = katibv1alpha1.ConditionFailed
@@ -248,7 +248,7 @@ func populateConfigForNAS(instance *katibv1alpha1.StudyJob) (*katibapi.StudyConf
 }
 
 func deleteStudy(instance *katibv1alpha1.StudyJob) error {
-	conn, err := grpc.Dial(pkg.ManagerAddr, grpc.WithInsecure())
+	conn, err := grpc.Dial(common.ManagerAddr, grpc.WithInsecure())
 	if err != nil {
 		log.Printf("Connect katib manager error %v", err)
 		return err

diff --git a/pkg/controller/v1alpha1/studyjob/manifest_parser.go b/pkg/controller/v1alpha1/studyjob/manifest_parser.go
@@ -20,7 +20,7 @@ import (
 	"fmt"
 	"text/template"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha1"
 	katibv1alpha1 "github.com/kubeflow/katib/pkg/api/operators/apis/studyjob/v1alpha1"
 	katibapi "github.com/kubeflow/katib/pkg/api/v1alpha1"
 	"github.com/kubeflow/katib/pkg/manager/v1alpha1/studyjobclient"
@@ -100,7 +100,7 @@ func getMetricsCollectorManifest(studyID string, trialID string, workerID string
 		"WorkerID":       workerID,
 		"WorkerKind":     workerKind,
 		"NameSpace":      namespace,
-		"ManagerSerivce": pkg.GetManagerAddr(),
+		"ManagerSerivce": common.GetManagerAddr(),
 	}
 	if mcs != nil && mcs.GoTemplate.RawTemplate != "" {
 		mtp, err = template.New("MetricsCollector").Parse(mcs.GoTemplate.RawTemplate)

diff --git a/pkg/controller/v1alpha1/studyjob/studyjob_controller.go b/pkg/controller/v1alpha1/studyjob/studyjob_controller.go
@@ -20,7 +20,7 @@ import (
 	"log"
 	"strconv"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha1"
 	katibv1alpha1 "github.com/kubeflow/katib/pkg/api/operators/apis/studyjob/v1alpha1"
 	katibapi "github.com/kubeflow/katib/pkg/api/v1alpha1"
 	pytorchjobv1beta1 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1beta1"
@@ -484,7 +484,7 @@ func (r *ReconcileStudyJobController) checkStatus(instance *katibv1alpha1.StudyJ
 		grpc.WithInsecure(),
 		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(maxMsgSize)),
 	}
-	conn, err := grpc.Dial(pkg.ManagerAddr, opts...)
+	conn, err := grpc.Dial(common.ManagerAddr, opts...)
 	if err != nil {
 		log.Printf("Connect katib manager error %v", err)
 		instance.Status.Condition = katibv1alpha1.ConditionFailed

diff --git a/pkg/controller/v1alpha2/experiment/util/manifest_util.go b/pkg/controller/v1alpha2/experiment/util/manifest_util.go
@@ -20,7 +20,7 @@ import (
 	"strings"
 	"text/template"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha2"
 	katibv1alpha2 "github.com/kubeflow/katib/pkg/api/operators/apis/experiment/v1alpha2"
 	"github.com/kubeflow/katib/pkg/util/v1alpha2/katibclient"
 
@@ -35,7 +35,7 @@ func getMetricsCollectorManifest(experimentName string, trialName string, jobKin
 		"Trial":          trialName,
 		"JobKind":        jobKind,
 		"NameSpace":      namespace,
-		"ManagerService": pkg.GetManagerAddr(),
+		"ManagerService": common.GetManagerAddr(),
 		"MetricNames":    strings.Join(metricNames, ";"),
 	}
 	if mcs != nil && mcs.GoTemplate.RawTemplate != "" {

diff --git a/pkg/suggestion/v1alpha1/grid_service.go b/pkg/suggestion/v1alpha1/grid_service.go
@@ -6,7 +6,7 @@ import (
 	"log"
 	"strconv"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha1"
 	"github.com/kubeflow/katib/pkg/api/v1alpha1"
 
 	"google.golang.org/grpc"
@@ -138,7 +138,7 @@ func (s *GridSuggestService) genGrids(studyID string, pcs []*api.ParameterConfig
 }
 
 func (s *GridSuggestService) GetSuggestions(ctx context.Context, in *api.GetSuggestionsRequest) (*api.GetSuggestionsReply, error) {
-	conn, err := grpc.Dial(pkg.ManagerAddr, grpc.WithInsecure())
+	conn, err := grpc.Dial(common.ManagerAddr, grpc.WithInsecure())
 	if err != nil {
 		log.Printf("could not connect: %v", err)
 		return &api.GetSuggestionsReply{}, err

diff --git a/pkg/suggestion/v1alpha1/hyperband_service.go b/pkg/suggestion/v1alpha1/hyperband_service.go
@@ -9,7 +9,7 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha1"
 	"github.com/kubeflow/katib/pkg/api/v1alpha1"
 
 	"google.golang.org/grpc"
@@ -403,7 +403,7 @@ func (h *HyperBandSuggestService) shLoopParamUpdate(studyID string, hbparam *Hyp
 }
 
 func (h *HyperBandSuggestService) GetSuggestions(ctx context.Context, in *api.GetSuggestionsRequest) (*api.GetSuggestionsReply, error) {
-	conn, err := grpc.Dial(pkg.ManagerAddr, grpc.WithInsecure())
+	conn, err := grpc.Dial(common.ManagerAddr, grpc.WithInsecure())
 	if err != nil {
 		log.Fatalf("could not connect: %v", err)
 		return &api.GetSuggestionsReply{}, err

diff --git a/pkg/suggestion/v1alpha1/random_service.go b/pkg/suggestion/v1alpha1/random_service.go
@@ -7,7 +7,7 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha1"
 	"github.com/kubeflow/katib/pkg/api/v1alpha1"
 	"google.golang.org/grpc"
 )
@@ -33,7 +33,7 @@ func (s *RandomSuggestService) IntRandom(min, max int) int {
 }
 
 func (s *RandomSuggestService) GetSuggestions(ctx context.Context, in *api.GetSuggestionsRequest) (*api.GetSuggestionsReply, error) {
-	conn, err := grpc.Dial(pkg.ManagerAddr, grpc.WithInsecure())
+	conn, err := grpc.Dial(common.ManagerAddr, grpc.WithInsecure())
 	if err != nil {
 		log.Fatalf("could not connect: %v", err)
 		return &api.GetSuggestionsReply{}, err

diff --git a/pkg/ui/ui.go b/pkg/ui/ui.go
@@ -9,7 +9,7 @@ import (
 	"strings"
 	"time"
 
-	"github.com/kubeflow/katib/pkg"
+	common "github.com/kubeflow/katib/pkg/common/v1alpha1"
 	studyjobv1alpha1 "github.com/kubeflow/katib/pkg/api/operators/apis/studyjob/v1alpha1"
 	"github.com/kubeflow/katib/pkg/api/v1alpha1"
 	"github.com/kubeflow/katib/pkg/manager/v1alpha1/studyjobclient"
@@ -54,7 +54,7 @@ func (k *KatibUIHandler) connectManager() (*grpc.ClientConn, api.ManagerClient,
 		grpc.WithInsecure(),
 		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(maxMsgSize)),
 	}
-	conn, err := grpc.Dial(pkg.ManagerAddr, opts...)
+	conn, err := grpc.Dial(common.ManagerAddr, opts...)
 	if err != nil {
 		log.Printf("Connect katib manager error %v", err)
 		return nil, nil, err