Replace gRPC code generation tool from Znly/protoc to Buf

Makefile

@@ -10,10 +10,8 @@ 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)
 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)
 	$(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

buf.gen.yaml

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

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,
+        )

go.mod

@@ -8,7 +8,6 @@ require (
 	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 +20,8 @@ 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
 	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

hack/update-codegen.sh

@@ -40,7 +40,7 @@ echo ">> Using ${CODEGEN_PKG} for the code generator"
 chmod +x "${CODEGEN_PKG}/generate-groups.sh"
 chmod +x "${CODEGEN_PKG}/generate-internal-groups.sh"
 
-PROJECT_ROOT=${GOPATH}/src/github.com/kubeflow/katib
+PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd -P)"
 
 modules=(experiments suggestions trials common)
 versions=(v1beta1)

hack/update-openapigen.sh

@@ -39,7 +39,7 @@ echo ">> Using ${CODEGEN_PKG} for the code generator"
 # Ensure we can execute.
 chmod +x "${CODEGEN_PKG}/generate-groups.sh"
 
-PROJECT_ROOT=${GOPATH}/src/github.com/kubeflow/katib
+PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd -P)"
 VERSION_LIST=(v1beta1)
 SWAGGER_VERSION="0.1"
 

hack/update-proto.sh

@@ -0,0 +1,44 @@
+#!/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}"
+  "${GOBIN}/buf" generate
+popd
+
+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,13 @@
+version: v2
+plugins:
+- 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: .