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")
 }