From ca903e25e598d3232026380098465bf84c352380 Mon Sep 17 00:00:00 2001 From: Yuki Iwai Date: Mon, 22 Aug 2022 20:28:30 +0900 Subject: [PATCH] Add --connect-timeout flag to katib-db-manager (#1937) --- cmd/db-manager/v1beta1/Dockerfile | 1 - cmd/db-manager/v1beta1/main.go | 9 +++++++-- docs/developer-guide.md | 10 +++++++++- pkg/db/v1beta1/common/const.go | 1 - pkg/db/v1beta1/db.go | 7 ++++--- pkg/db/v1beta1/mysql/mysql.go | 4 ++-- pkg/db/v1beta1/postgres/postgres.go | 6 +++--- 7 files changed, 25 insertions(+), 13 deletions(-) diff --git a/cmd/db-manager/v1beta1/Dockerfile b/cmd/db-manager/v1beta1/Dockerfile index 7dae3cb5de2..d614722de78 100644 --- a/cmd/db-manager/v1beta1/Dockerfile +++ b/cmd/db-manager/v1beta1/Dockerfile @@ -28,4 +28,3 @@ WORKDIR /app COPY --from=build-env /bin/grpc_health_probe /bin/ COPY --from=build-env /go/src/github.com/kubeflow/katib/katib-db-manager /app/ ENTRYPOINT ["./katib-db-manager"] -CMD ["-w", "kubernetes"] diff --git a/cmd/db-manager/v1beta1/main.go b/cmd/db-manager/v1beta1/main.go index 114a609de3d..17f949f0690 100644 --- a/cmd/db-manager/v1beta1/main.go +++ b/cmd/db-manager/v1beta1/main.go @@ -22,6 +22,7 @@ import ( "fmt" "net" "os" + "time" health_pb "github.com/kubeflow/katib/pkg/apis/manager/health" api_pb "github.com/kubeflow/katib/pkg/apis/manager/v1beta1" @@ -34,7 +35,8 @@ import ( ) const ( - port = "0.0.0.0:6789" + port = "0.0.0.0:6789" + defaultConnectTimeout = time.Second * 60 ) var dbIf common.KatibDBInterface @@ -87,14 +89,17 @@ func (s *server) Check(ctx context.Context, in *health_pb.HealthCheckRequest) (* } func main() { + var connectTimeout time.Duration + flag.DurationVar(&connectTimeout, "connect-timeout", defaultConnectTimeout, "Timeout before calling error during database connection. (e.g. 120s)") flag.Parse() + var err error dbNameEnvName := common.DBNameEnvName dbName := os.Getenv(dbNameEnvName) if dbName == "" { klog.Fatal("DB_NAME env is not set. Exiting") } - dbIf, err = db.NewKatibDBInterface(dbName) + dbIf, err = db.NewKatibDBInterface(dbName, connectTimeout) if err != nil { klog.Fatalf("Failed to open db connection: %v", err) } diff --git a/docs/developer-guide.md b/docs/developer-guide.md index 93bed7c9963..d8a8c605c2b 100644 --- a/docs/developer-guide.md +++ b/docs/developer-guide.md @@ -58,7 +58,7 @@ make generate Below is a list of command-line flags accepted by Katib controller: | Name | Type | Default | Description | -| ------------------------------- | ------------------------- | ----------------------------- | ---------------------------------------------------------------------------------------------------------------------- | +|---------------------------------|---------------------------|-------------------------------|------------------------------------------------------------------------------------------------------------------------| | enable-grpc-probe-in-suggestion | bool | true | Enable grpc probe in suggestions | | experiment-suggestion-name | string | "default" | The implementation of suggestion interface in experiment controller | | metrics-addr | string | ":8080" | The address the metric endpoint binds to | @@ -68,6 +68,14 @@ Below is a list of command-line flags accepted by Katib controller: | enable-leader-election | bool | false | Enable leader election for katib-controller. Enabling this will ensure there is only one active katib-controller. | | leader-election-id | string | "3fbc96e9.katib.kubeflow.org" | The ID for leader election. | +## DB Manager Flags + +Below is a list of command-line flags accepted by Katib DB Manager: + +| Name | Type | Default | Description | +|-----------------|---------------|---------|---------------------------------------------------------| +| connect-timeout | time.Duration | 60s | Timeout before calling error during database connection | + ## Workflow design Please see [workflow-design.md](./workflow-design.md). diff --git a/pkg/db/v1beta1/common/const.go b/pkg/db/v1beta1/common/const.go index 97927c07731..3bdb09da66e 100644 --- a/pkg/db/v1beta1/common/const.go +++ b/pkg/db/v1beta1/common/const.go @@ -20,7 +20,6 @@ import "time" const ( ConnectInterval = 5 * time.Second - ConnectTimeout = 60 * time.Second DBUserEnvName = "DB_USER" DBNameEnvName = "DB_NAME" diff --git a/pkg/db/v1beta1/db.go b/pkg/db/v1beta1/db.go index 7d48069b589..d5834711023 100644 --- a/pkg/db/v1beta1/db.go +++ b/pkg/db/v1beta1/db.go @@ -18,6 +18,7 @@ package db import ( "errors" + "time" "github.com/kubeflow/katib/pkg/db/v1beta1/common" "github.com/kubeflow/katib/pkg/db/v1beta1/mysql" @@ -25,14 +26,14 @@ import ( "k8s.io/klog" ) -func NewKatibDBInterface(dbName string) (common.KatibDBInterface, error) { +func NewKatibDBInterface(dbName string, connectTimeout time.Duration) (common.KatibDBInterface, error) { if dbName == common.MySqlDBNameEnvValue { klog.Info("Using MySQL") - return mysql.NewDBInterface() + return mysql.NewDBInterface(connectTimeout) } else if dbName == common.PostgresSQLDBNameEnvValue { klog.Info("Using Postgres") - return postgres.NewDBInterface() + return postgres.NewDBInterface(connectTimeout) } return nil, errors.New("Invalid DB Name") } diff --git a/pkg/db/v1beta1/mysql/mysql.go b/pkg/db/v1beta1/mysql/mysql.go index 4e7b4427b95..4fe863e8b47 100644 --- a/pkg/db/v1beta1/mysql/mysql.go +++ b/pkg/db/v1beta1/mysql/mysql.go @@ -73,8 +73,8 @@ func NewWithSQLConn(db *sql.DB) (common.KatibDBInterface, error) { return d, nil } -func NewDBInterface() (common.KatibDBInterface, error) { - db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, common.ConnectTimeout) +func NewDBInterface(connectTimeout time.Duration) (common.KatibDBInterface, error) { + db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, connectTimeout) if err != nil { return nil, fmt.Errorf("DB open failed: %v", err) } diff --git a/pkg/db/v1beta1/postgres/postgres.go b/pkg/db/v1beta1/postgres/postgres.go index 05781a4dfe2..587f973c301 100644 --- a/pkg/db/v1beta1/postgres/postgres.go +++ b/pkg/db/v1beta1/postgres/postgres.go @@ -49,7 +49,7 @@ func getDbName() string { common.PostgreSQLDBHostEnvName, common.DefaultPostgreSQLHost) dbPort := env.GetEnvOrDefault( common.PostgreSQLDBPortEnvName, common.DefaultPostgreSQLPort) - dbName := env.GetEnvOrDefault(common.DefaultPostgreSQLDatabase, + dbName := env.GetEnvOrDefault(common.PostgreSQLDatabase, common.DefaultPostgreSQLDatabase) psqlInfo := fmt.Sprintf("host=%s port=%s user=%s "+ @@ -59,8 +59,8 @@ func getDbName() string { return psqlInfo } -func NewDBInterface() (common.KatibDBInterface, error) { - db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, common.ConnectTimeout) +func NewDBInterface(connectTimeout time.Duration) (common.KatibDBInterface, error) { + db, err := common.OpenSQLConn(dbDriver, getDbName(), common.ConnectInterval, connectTimeout) if err != nil { return nil, fmt.Errorf("DB open failed: %v", err) }