Support postgres as a katib db

.github/workflows/darts-cifar10-e2e-test.yaml

@@ -24,6 +24,7 @@ jobs:
           experiments: ${{ matrix.experiments }}
           # Comma Delimited
           trial-images: darts-cnn-cifar10-cpu
+          database-type: ${{ matrix.database-type }}
 
     strategy:
       fail-fast: false
@@ -33,3 +34,4 @@ jobs:
         kubernetes-version: ["v1.21.13", "v1.22.10", "v1.23.7"]
         # Comma Delimited
         experiments: ["darts-cpu"]
+        database-type: ["mysql", "postgres"]

.github/workflows/enas-cifar10-e2e-test.yaml

@@ -24,6 +24,7 @@ jobs:
           experiments: ${{ matrix.experiments }}
           # Comma Delimited
           trial-images: enas-cnn-cifar10-cpu
+          database-type: ${{ matrix.database-type }}
 
     strategy:
       fail-fast: false
@@ -33,3 +34,4 @@ jobs:
         kubernetes-version: ["v1.21.13", "v1.22.10", "v1.23.7"]
         # Comma Delimited
         experiments: ["enas-cpu"]
+        database-type: ["mysql", "postgres"]

.github/workflows/mxnet-mnist-e2e-test.yaml

@@ -24,6 +24,7 @@ jobs:
           experiments: ${{ matrix.experiments }}
           # Comma Delimited
           trial-images: mxnet-mnist
+          database-type: ${{ matrix.database-type }}
 
     strategy:
       fail-fast: false
@@ -39,3 +40,4 @@ jobs:
           # others
           - "grid,bayesian-optimization,tpe"
           - "multivariate-tpe,cma-es,hyperband"
+        database-type: ["mysql", "postgres"]

.github/workflows/pytorch-mnist-e2e-test.yaml

@@ -25,6 +25,7 @@ jobs:
           training-operator: true
           # Comma Delimited
           trial-images: pytorch-mnist-cpu
+          database-type: ${{ matrix.database-type }}
 
     strategy:
       fail-fast: false
@@ -36,3 +37,4 @@ jobs:
         experiments:
           - "file-metrics-collector,pytorchjob-mnist"
           - "median-stop-with-json-format,file-metrics-collector-with-json-format"
+        database-type: ["mysql", "postgres"]

.github/workflows/simple-pbt-e2e-test.yaml

@@ -24,6 +24,7 @@ jobs:
           experiments: ${{ matrix.experiments }}
           # Comma Delimited
           trial-images: simple-pbt
+          database-type: ${{ matrix.database-type }}
 
     strategy:
       fail-fast: false
@@ -34,3 +35,4 @@ jobs:
         kubernetes-version: ["v1.21.12", "v1.22.9", "v1.23.6"]
         # Comma Delimited
         experiments: ["simple-pbt"]
+        database-type: ["mysql", "postgres"]

.github/workflows/template-e2e-test/action.yaml

@@ -7,13 +7,18 @@ inputs:
   training-operator:
     required: false
     type: boolean
+    default: false
   trial-images:
     required: true
     type: string
   katib-ui:
     required: true
     type: boolean
     default: false
+  database-type:
+    required: false
+    type: string
+    default: mysql
 
 runs:
   using: composite
@@ -24,7 +29,7 @@ runs:
 
     - name: Set Up Katib
       shell: bash
-      run: ./test/e2e/v1beta1/scripts/gh-actions/setup-katib.sh ${{ inputs.katib-ui }} ${{ inputs.training-operator }}
+      run: ./test/e2e/v1beta1/scripts/gh-actions/setup-katib.sh ${{ inputs.katib-ui }} ${{ inputs.training-operator }} ${{ inputs.database-type }}
 
     - name: Run E2E Experiment
       shell: bash

.github/workflows/tf-mnist-with-summaries-e2e-test.yaml

@@ -25,6 +25,7 @@ jobs:
           training-operator: true
           # Comma Delimited
           trial-images: tf-mnist-with-summaries
+          database-type: ${{ matrix.database-type }}
 
     strategy:
       fail-fast: false
@@ -34,3 +35,4 @@ jobs:
         kubernetes-version: ["v1.21.13", "v1.22.10", "v1.23.7"]
         # Comma Delimited
         experiments: ["tfjob-mnist-with-summaries"]
+        database-type: ["mysql", "postgres"]

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/google/go-containerregistry/pkg/authn/k8schain v0.0.0-20211222182933-7c19fa370dbd
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/hpcloud/tail v1.0.1-0.20180514194441-a1dbeea552b7
+	github.com/lib/pq v1.10.6
 	github.com/mattbaird/jsonpatch v0.0.0-20171005235357-81af80346b1a
 	github.com/onsi/gomega v1.17.0
 	github.com/prometheus/client_golang v1.11.0

go.sum

@@ -869,6 +869,8 @@ github.com/lib/pq v1.3.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.8.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lib/pq v1.9.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lib/pq v1.10.3/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
+github.com/lib/pq v1.10.6 h1:jbk+ZieJ0D7EVGJYpL9QTz7/YW6UHbmdnZWYyK5cdBs=
+github.com/lib/pq v1.10.6/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/logrusorgru/aurora v0.0.0-20181002194514-a7b3b318ed4e/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
 github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4=
 github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=

manifests/v1beta1/components/postgres/kustomization.yaml

@@ -0,0 +1,9 @@
+---
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization
+
+resources:
+  - postgres.yaml
+  - pvc.yaml
+  - secret.yaml
+  - service.yaml

manifests/v1beta1/components/postgres/postgres.yaml

@@ -0,0 +1,42 @@
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: katib-postgres
+  namespace: kubeflow
+  labels:
+    katib.kubeflow.org/component: postgres
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      katib.kubeflow.org/component: postgres
+  strategy:
+    type: Recreate
+  template:
+    metadata:
+      labels:
+        katib.kubeflow.org/component: postgres
+      annotations:
+        sidecar.istio.io/inject: "false"
+    spec:
+      containers:
+        - name: katib-postgres
+          image: postgres:14.3-alpine
+          envFrom:
+            - secretRef:
+                name: katib-postgres-secrets
+          env:
+            - name: PGDATA
+              value: /var/lib/postgresql/data/pgdata
+          ports:
+            - name: postgres
+              containerPort: 5432
+              protocol: TCP
+          volumeMounts:
+            - name: katib-postgres
+              mountPath: /var/lib/postgresql/data
+      volumes:
+        - name: katib-postgres
+          persistentVolumeClaim:
+            claimName: katib-postgres

manifests/v1beta1/components/postgres/pvc.yaml

@@ -0,0 +1,12 @@
+---
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: katib-postgres
+  namespace: kubeflow
+spec:
+  accessModes:
+    - ReadWriteOnce
+  resources:
+    requests:
+      storage: 3Gi

manifests/v1beta1/components/postgres/secret.yaml

@@ -0,0 +1,10 @@
+---
+apiVersion: v1
+kind: Secret
+type: Opaque
+metadata:
+  name: katib-postgres-secrets
+data:
+  POSTGRES_USER: a2F0aWI=  # katib
+  POSTGRES_PASSWORD: a2F0aWI=  # katib
+  POSTGRES_DB: a2F0aWI=  # katib

manifests/v1beta1/components/postgres/service.yaml

@@ -0,0 +1,16 @@
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: katib-postgres
+  namespace: kubeflow
+  labels:
+    katib.kubeflow.org/component: postgres
+spec:
+  type: ClusterIP
+  ports:
+    - port: 5432
+      protocol: TCP
+      name: dbapi
+  selector:
+    katib.kubeflow.org/component: postgres

manifests/v1beta1/installs/katib-standalone-postgres/kustomization.yaml

@@ -0,0 +1,46 @@
+---
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization
+namespace: kubeflow
+resources:
+  # Namespace.
+  - ../../components/namespace/
+  # Katib controller.
+  - ../../components/controller/
+  # Katib CRDs.
+  - ../../components/crd/
+  # Katib DB manager.
+  - ../../components/db-manager/
+  # Katib DB postgres.
+  - ../../components/postgres/
+  # Katib UI.
+  - ../../components/ui/
+  # Katib Cert Generator
+  - ../../components/cert-generator/
+  # Katib webhooks.
+  - ../../components/webhook/
+images:
+  - name: docker.io/kubeflowkatib/katib-controller
+    newName: docker.io/kubeflowkatib/katib-controller
+    newTag: latest
+  - name: docker.io/kubeflowkatib/katib-db-manager
+    newName: docker.io/kubeflowkatib/katib-db-manager
+    newTag: latest
+  - name: docker.io/kubeflowkatib/katib-ui
+    newName: docker.io/kubeflowkatib/katib-ui
+    newTag: latest
+  - name: docker.io/kubeflowkatib/cert-generator
+    newName: docker.io/kubeflowkatib/cert-generator
+    newTag: latest
+patches:
+  - patch: |-
+      - op: replace
+        path: /spec/template/spec/containers/0/env
+        value:
+          - name: DB_NAME
+            value: "postgres"
+          - name: DB_PASSWORD
+            value: "katib"
+    target:
+      kind: Deployment
+      labelSelector: "katib.kubeflow.org/component=db-manager"

pkg/db/v1beta1/common/connection.go

@@ -0,0 +1,49 @@
+/*
+Copyright 2022 The Kubeflow Authors.
+
+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
+
+import (
+	"database/sql"
+	"fmt"
+	"time"
+
+	"k8s.io/klog"
+)
+
+func OpenSQLConn(driverName string, dataSourceName string, interval time.Duration,
+	timeout time.Duration) (*sql.DB, error) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	timeoutC := time.After(timeout)
+	for {
+		select {
+		case <-ticker.C:
+			if db, err := sql.Open(driverName, dataSourceName); err == nil {
+				// if db, err := sql.Open(driverName, dataSourceName); err == nil {
+				if err = db.Ping(); err == nil {
+					return db, nil
+				}
+				klog.Errorf("Ping to Katib db failed: %v", err)
+			} else {
+				klog.Errorf("Open sql connection failed: %v", err)
+			}
+		case <-timeoutC:
+			return nil, fmt.Errorf("Timeout waiting for DB conn successfully opened.")
+		}
+	}
+}

pkg/db/v1beta1/common/const.go

@@ -17,14 +17,12 @@ limitations under the License.
 package common
 
 const (
-	DBUserEnvName = "DB_USER"
-
-	DBNameEnvName = "DB_NAME"
+	DBUserEnvName     = "DB_USER"
+	DBNameEnvName     = "DB_NAME"
+	DBPasswordEnvName = "DB_PASSWORD"
 
 	MySqlDBNameEnvValue = "mysql"
 
-	DBPasswordEnvName = "DB_PASSWORD"
-
 	MySQLDBHostEnvName = "KATIB_MYSQL_DB_HOST"
 	MySQLDBPortEnvName = "KATIB_MYSQL_DB_PORT"
 	MySQLDatabase      = "KATIB_MYSQL_DB_DATABASE"
@@ -33,4 +31,15 @@ const (
 	DefaultMySQLDatabase = "katib"
 	DefaultMySQLHost     = "katib-mysql"
 	DefaultMySQLPort     = "3306"
+
+	PostgresSQLDBNameEnvValue = "postgres"
+
+	PostgreSQLDBHostEnvName = "KATIB_POSTGRESQL_DB_HOST"
+	PostgreSQLDBPortEnvName = "KATIB_POSTGRESQL_DB_PORT"
+	PostgreSQLDatabase      = "KATIB_POSTGRESQL_DB_DATABASE"
+
+	DefaultPostgreSQLUser     = "katib"
+	DefaultPostgreSQLDatabase = "katib"
+	DefaultPostgreSQLHost     = "katib-postgres"
+	DefaultPostgreSQLPort     = "5432"
 )

pkg/db/v1beta1/db.go

@@ -21,12 +21,18 @@ import (
 
 	"github.com/kubeflow/katib/pkg/db/v1beta1/common"
 	"github.com/kubeflow/katib/pkg/db/v1beta1/mysql"
+	"github.com/kubeflow/katib/pkg/db/v1beta1/postgres"
+	"k8s.io/klog"
 )
 
 func NewKatibDBInterface(dbName string) (common.KatibDBInterface, error) {
 
 	if dbName == common.MySqlDBNameEnvValue {
+		klog.Info("Using MySQL")
 		return mysql.NewDBInterface()
+	} else if dbName == common.PostgresSQLDBNameEnvValue {
+		klog.Info("Using Postgres")
+		return postgres.NewDBInterface()
 	}
 	return nil, errors.New("Invalid DB Name")
 }