Cherry pick of #2350 #2355 #2357 #2344 #2358 #2360 into release-0.17 branch

.github/workflows/test-go.yaml

@@ -28,6 +28,7 @@ jobs:
         uses: actions/setup-go@v5
         with:
           go-version-file: ${{ env.GOPATH }}/src/github.com/kubeflow/katib/go.mod
+          cache-dependency-path: ${{ env.GOPATH }}/src/github.com/kubeflow/katib/go.sum
 
       - name: Check Go Modules, Generated Go/Python codes, and Format
         run: make check
@@ -50,6 +51,7 @@ jobs:
         uses: actions/setup-go@v5
         with:
           go-version-file: ${{ env.GOPATH }}/src/github.com/kubeflow/katib/go.mod
+          cache-dependency-path: ${{ env.GOPATH }}/src/github.com/kubeflow/katib/go.sum
 
       - name: Run Go test
         run: go mod download && make test ENVTEST_K8S_VERSION=${{ matrix.kubernetes-version }}

Makefile

@@ -8,12 +8,10 @@ COMMIT := v1beta1-$(shell git rev-parse --short=7 HEAD)
 KATIB_REGISTRY := docker.io/kubeflowkatib
 CPU_ARCH ?= linux/amd64,linux/arm64
 ENVTEST_K8S_VERSION ?= 1.29
-MOCKGEN_VERSION ?= $(shell grep 'github.com/golang/mock' go.mod | cut -d ' ' -f 2)
+MOCKGEN_VERSION ?= $(shell grep 'go.uber.org/mock' go.mod | cut -d ' ' -f 2)
 GO_VERSION=$(shell grep '^go' go.mod | cut -d ' ' -f 2)
+GOPATH ?= $(shell go env GOPATH)
 
-# for pytest
-PYTHONPATH := $(PYTHONPATH):$(CURDIR)/pkg/apis/manager/v1beta1/python:$(CURDIR)/pkg/apis/manager/health/python
-PYTHONPATH := $(PYTHONPATH):$(CURDIR)/pkg/metricscollector/v1beta1/common:$(CURDIR)/pkg/metricscollector/v1beta1/tfevent-metricscollector
 TEST_TENSORFLOW_EVENT_FILE_PATH ?= $(CURDIR)/test/unit/v1beta1/metricscollector/testdata/tfevent-metricscollector/logs
 
 # Run tests
@@ -93,17 +91,13 @@ controller-gen:
 # 4. Generate gRPC manager APIs (pkg/apis/manager/v1beta1/build.sh and pkg/apis/manager/health/build.sh)
 # 5. Generate Go mock codes
 generate: controller-gen
-ifndef GOPATH
-	$(error GOPATH not defined, please define GOPATH. Run "go help gopath" to learn more about GOPATH)
-endif
 ifndef HAS_MOCKGEN
-	go install github.com/golang/mock/mockgen@$(MOCKGEN_VERSION)
+	go install go.uber.org/mock/mockgen@$(MOCKGEN_VERSION)
 	$(info "mockgen has been installed")
 endif
 	go generate ./pkg/... ./cmd/...
 	hack/gen-python-sdk/gen-sdk.sh
-	pkg/apis/manager/v1beta1/build.sh
-	pkg/apis/manager/health/build.sh
+	hack/update-proto.sh
 	hack/update-mockgen.sh
 
 # Build images for the Katib v1beta1 components.
@@ -175,9 +169,9 @@ ifeq ("$(wildcard $(TEST_TENSORFLOW_EVENT_FILE_PATH))", "")
 endif
 
 pytest: prepare-pytest prepare-pytest-testdata
-	PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/suggestion --ignore=./test/unit/v1beta1/suggestion/test_skopt_service.py
-	PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/earlystopping
-	PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/metricscollector
+	pytest ./test/unit/v1beta1/suggestion --ignore=./test/unit/v1beta1/suggestion/test_skopt_service.py
+	pytest ./test/unit/v1beta1/earlystopping
+	pytest ./test/unit/v1beta1/metricscollector
 
 # The skopt service doesn't work appropriately with Python 3.11.
 # So, we need to run the test with Python 3.9.
@@ -187,4 +181,4 @@ pytest-skopt:
 	pip install six
 	pip install --prefer-binary -r test/unit/v1beta1/requirements.txt
 	pip install --prefer-binary -r cmd/suggestion/skopt/v1beta1/requirements.txt
-	PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/suggestion/test_skopt_service.py
+	pytest ./test/unit/v1beta1/suggestion/test_skopt_service.py

cmd/db-manager/v1beta1/main_test.go

@@ -20,7 +20,7 @@ import (
 	"context"
 	"testing"
 
-	"github.com/golang/mock/gomock"
+	"go.uber.org/mock/gomock"
 
 	health_pb "github.com/kubeflow/katib/pkg/apis/manager/health"
 	api_pb "github.com/kubeflow/katib/pkg/apis/manager/v1beta1"

cmd/earlystopping/medianstop/v1beta1/requirements.txt

@@ -1,5 +1,5 @@
 grpcio>=1.41.1
-protobuf>=3.19.5, <=3.20.3
+protobuf>=4.21.12,<5
 googleapis-common-protos==1.6.0
 kubernetes==22.6.0
 cython>=0.29.24

cmd/metricscollector/v1beta1/tfevent-metricscollector/main.py

@@ -15,6 +15,7 @@
 import grpc
 import argparse
 import api_pb2
+import api_pb2_grpc
 from pns import WaitMainProcesses
 import const
 from tfevent_loader import MetricsCollector
@@ -55,25 +56,28 @@ def parse_options():
     wait_all_processes = opt.wait_all_processes.lower() == "true"
     db_manager_server = opt.db_manager_server_addr.split(':')
     if len(db_manager_server) != 2:
-        raise Exception("Invalid Katib DB manager service address: %s" %
-                        opt.db_manager_server_addr)
+        raise Exception(
+            f"Invalid Katib DB manager service address: {opt.db_manager_server_addr}"
+        )
 
     WaitMainProcesses(
         pool_interval=opt.poll_interval,
         timout=opt.timeout,
         wait_all=wait_all_processes,
-        completed_marked_dir=opt.metrics_file_dir)
+        completed_marked_dir=opt.metrics_file_dir,
+    )
 
-    mc = MetricsCollector(opt.metric_names.split(';'))
+    mc = MetricsCollector(opt.metric_names.split(";"))
     observation_log = mc.parse_file(opt.metrics_file_dir)
 
-    channel = grpc.beta.implementations.insecure_channel(
-        db_manager_server[0], int(db_manager_server[1]))
-
-    with api_pb2.beta_create_DBManager_stub(channel) as client:
-        logger.info("In " + opt.trial_name + " " +
-                    str(len(observation_log.metric_logs)) + " metrics will be reported.")
-        client.ReportObservationLog(api_pb2.ReportObservationLogRequest(
-            trial_name=opt.trial_name,
-            observation_log=observation_log
-        ), timeout=timeout_in_seconds)
+    with grpc.insecure_channel(opt.db_manager_server_addr) as channel:
+        stub = api_pb2_grpc.DBManagerStub(channel)
+        logger.info(
+            f"In {opt.trial_name} {str(len(observation_log.metric_logs))} metrics will be reported."
+        )
+        stub.ReportObservationLog(
+            api_pb2.ReportObservationLogRequest(
+                trial_name=opt.trial_name, observation_log=observation_log
+            ),
+            timeout=timeout_in_seconds,
+        )

cmd/metricscollector/v1beta1/tfevent-metricscollector/requirements.txt

@@ -3,4 +3,4 @@ rfc3339>=6.2
 grpcio>=1.41.1
 googleapis-common-protos==1.6.0
 tensorflow==2.16.1
-protobuf<=3.20.3
+protobuf>=4.21.12,<5

cmd/suggestion/hyperband/v1beta1/requirements.txt

@@ -4,6 +4,6 @@ numpy>=1.25.2
 scikit-learn>=0.24.0
 scipy>=1.5.4
 forestci==0.3
-protobuf>=3.19.5, <=3.20.3
+protobuf>=4.21.12,<5
 googleapis-common-protos==1.6.0
 cython>=0.29.24

cmd/suggestion/hyperopt/v1beta1/requirements.txt

@@ -4,7 +4,7 @@ numpy>=1.25.2
 scikit-learn>=0.24.0
 scipy>=1.5.4
 forestci==0.3
-protobuf>=3.19.5, <=3.20.3
+protobuf>=4.21.12,<5
 googleapis-common-protos==1.6.0
 hyperopt==0.2.5
 cython>=0.29.24

cmd/suggestion/nas/darts/v1beta1/requirements.txt

@@ -1,4 +1,4 @@
 grpcio>=1.41.1
-protobuf>=3.19.5, <=3.20.3
+protobuf>=4.21.12,<5
 googleapis-common-protos==1.6.0
 cython>=0.29.24

cmd/suggestion/nas/enas/v1beta1/requirements.txt

@@ -2,4 +2,4 @@ grpcio>=1.41.1
 googleapis-common-protos==1.6.0
 cython>=0.29.24
 tensorflow==2.16.1
-protobuf<=3.20.3
+protobuf>=4.21.12,<5

cmd/suggestion/optuna/v1beta1/requirements.txt

@@ -1,4 +1,4 @@
 grpcio>=1.41.1
-protobuf>=3.19.5, <=3.20.3
+protobuf>=4.21.12,<5
 googleapis-common-protos==1.53.0
 optuna==3.3.0

cmd/suggestion/pbt/v1beta1/requirements.txt

@@ -1,4 +1,4 @@
 grpcio>=1.41.1
-protobuf>=3.19.5, <=3.20.3
+protobuf>=4.21.12,<5
 googleapis-common-protos==1.53.0
 numpy==1.25.2

cmd/suggestion/skopt/v1beta1/requirements.txt

@@ -7,7 +7,7 @@ numpy==1.23.5
 scikit-learn>=0.24.0, <=1.3.0
 scipy>=1.5.4
 forestci==0.3
-protobuf>=3.19.5, <=3.20.3
+protobuf>=4.21.12,<5
 googleapis-common-protos==1.6.0
 scikit-optimize>=0.9.0
 cython>=0.29.24

examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.cpu

@@ -4,10 +4,12 @@ ENV TARGET_DIR /opt/darts-cnn-cifar10
 
 ADD examples/v1beta1/trial-images/darts-cnn-cifar10 ${TARGET_DIR}
 
-WORKDIR  ${TARGET_DIR}
+WORKDIR ${TARGET_DIR}
 
+# TODO (andreyvelich): This is required since torchvision==0.17.1 is incompatible with numpy 2.0
+RUN pip install numpy==1.26.0
 RUN pip install --prefer-binary --no-cache-dir torch==2.2.1 torchvision==0.17.1
-RUN pip install --prefer-binary --no-cache-dir -r requirements.txt
+
 RUN chgrp -R 0 ${TARGET_DIR} \
   && chmod -R g+rwX ${TARGET_DIR}
 

examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu

@@ -6,6 +6,9 @@ WORKDIR /opt/pytorch-mnist
 
 # Add folder for the logs.
 RUN mkdir /katib
+
+# TODO (andreyvelich): This is required since torchvision==0.17.1 is incompatible with numpy 2.0
+RUN pip install numpy==1.26.0
 RUN pip install --prefer-binary --no-cache-dir torch==2.2.1 torchvision==0.17.1
 RUN pip install --prefer-binary --no-cache-dir -r requirements.txt
 

examples/v1beta1/trial-images/pytorch-mnist/requirements.txt

@@ -1,2 +1 @@
 cloudml-hypertune==0.1.0.dev6
-Pillow>=9.1.1

go.mod

@@ -7,8 +7,6 @@ require (
 	github.com/awalterschulze/gographviz v2.0.3+incompatible
 	github.com/c-bata/goptuna v0.8.0
 	github.com/go-sql-driver/mysql v1.5.0
-	github.com/golang/mock v1.6.0
-	github.com/golang/protobuf v1.5.4
 	github.com/google/go-cmp v0.6.0
 	github.com/google/go-containerregistry v0.15.2
 	github.com/google/go-containerregistry/pkg/authn/k8schain v0.0.0-20230517160804-b7ad3f13a62c
@@ -21,8 +19,9 @@ require (
 	github.com/shirou/gopsutil/v3 v3.22.5
 	github.com/spf13/viper v1.9.0
 	github.com/tidwall/gjson v1.14.1
-	golang.org/x/net v0.23.0
+	go.uber.org/mock v0.4.0
 	google.golang.org/grpc v1.58.3
+	google.golang.org/protobuf v1.33.0
 	k8s.io/api v0.29.3
 	k8s.io/apimachinery v0.29.3
 	k8s.io/client-go v0.29.3
@@ -84,6 +83,7 @@ require (
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang-jwt/jwt/v4 v4.5.0 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/gnostic-models v0.6.8 // indirect
 	github.com/google/go-containerregistry/pkg/authn/kubernetes v0.0.0-20230516205744-dbecb1de8cfa // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
@@ -128,6 +128,7 @@ require (
 	golang.org/x/crypto v0.21.0 // indirect
 	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e // indirect
 	golang.org/x/mod v0.14.0 // indirect
+	golang.org/x/net v0.23.0 // indirect
 	golang.org/x/oauth2 v0.12.0 // indirect
 	golang.org/x/sync v0.5.0 // indirect
 	golang.org/x/sys v0.18.0 // indirect
@@ -139,7 +140,6 @@ require (
 	gonum.org/v1/gonum v0.8.2 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
-	google.golang.org/protobuf v1.33.0 // indirect
 	gopkg.in/fsnotify/fsnotify.v1 v1.4.7 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/ini.v1 v1.63.2 // indirect

go.sum

@@ -263,7 +263,6 @@ github.com/golang/mock v1.4.1/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt
 github.com/golang/mock v1.4.3/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw=
 github.com/golang/mock v1.4.4/go.mod h1:l3mdAwkq5BuhzHwde/uurv3sEJeZMXNpwsxVWU71h+4=
 github.com/golang/mock v1.5.0/go.mod h1:CWnOUgYIOo4TcNZ0wHX3YZCqsaM1I1Jvs6v3mP3KVu8=
-github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
 github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.0/go.mod h1:Qd/q+1AKNOZr9uGQzbzCmRO6sUih6GTPZv6a1/R87v0=
@@ -689,6 +688,8 @@ go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
+go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
+go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=

hack/update-mockgen.sh

@@ -26,11 +26,11 @@ SCRIPT_ROOT="$(dirname "${BASH_SOURCE[0]}")/.."
 cd "${SCRIPT_ROOT}"
 
 # Grab mockgen version from go.mod
-MOCKGEN_VERSION=$(grep 'github.com/golang/mock' go.mod | awk '{print $2}')
+MOCKGEN_VERSION=$(grep 'go.uber.org/mock' go.mod | awk '{print $2}')
 
 if [[ ! $(mockgen -version) == "${MOCKGEN_VERSION}" ]]; then
   echo "You must use ${MOCKGEN_VERSION} mockgen version to run this script"
-  echo "To install mockgen follow this doc: https://github.com/golang/mock/tree/main#installation"
+  echo "To install mockgen follow this doc: https://github.com/uber-go/mock#installation"
   echo "Run 'mockgen -version' to check the installed version"
   exit 1
 fi

hack/update-proto.sh

@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+
+# Copyright 2024 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.
+
+set -e
+
+SOURCE="${BASH_SOURCE[0]}"
+while [ -h "$SOURCE" ]; do # resolve $SOURCE until the file is no longer a symlink
+  DIR="$( cd -P "$( dirname "$SOURCE" )" && pwd )"
+  SOURCE="$(readlink "$SOURCE")"
+  [[ $SOURCE != /* ]] && SOURCE="$DIR/$SOURCE" # if $SOURCE was a relative symlink, we need to resolve it relative to the path where the symlink file was located
+done
+ROOT_DIR="$( cd -P "$( dirname "$SOURCE" )/.." && pwd )"
+
+mkdir -p "${ROOT_DIR}/bin"
+export GOBIN=$ROOT_DIR/bin
+
+if [ ! -f "${GOBIN}/buf" ]; then
+  go install github.com/bufbuild/buf/cmd/[email protected]
+fi
+
+pushd "${ROOT_DIR}/pkg/apis/manager/health"
+  "${GOBIN}/buf" generate
+popd
+
+pushd "${ROOT_DIR}/pkg/apis/manager/v1beta1"
+  "${GOBIN}/buf" generate
+popd

pkg/apis/manager/health/buf.gen.yaml

@@ -0,0 +1,21 @@
+version: v2
+plugins:
+- remote: buf.build/protocolbuffers/go:v1.33.0
+  out: .
+  opt: module=github.com/kubeflow/katib/pkg/apis/manager/health
+
+- remote: buf.build/grpc/go:v1.3.0
+  out: .
+  opt: module=github.com/kubeflow/katib/pkg/apis/manager/health,require_unimplemented_servers=false
+
+- remote: buf.build/protocolbuffers/python:v26.1
+  out: python
+
+- remote: buf.build/protocolbuffers/pyi:v26.1
+  out: python
+
+- remote: buf.build/grpc/python:v1.64.1
+  out: python
+
+inputs:
+- directory: .