Skip to content

Commit dec5030

Browse files
authored
Cast resample_probability before range check in PBT validation (#2685)
PBT's ValidateAlgorithmSettings compares resample_probability against 0 and 1 without casting it, but algorithm setting values arrive as strings. The comparison raises TypeError and the validation crashes with an internal error instead of returning a clear message, unlike the adjacent n_population and truncation_threshold checks which cast first. Cast the value with float, matching those checks, and add a PBT service test covering the validation (the service had no unit test before). Signed-off-by: Saivedant Hava <saivedant169@gmail.com>
1 parent fc55ea8 commit dec5030

2 files changed

Lines changed: 73 additions & 1 deletion

File tree

pkg/suggestion/v1beta1/pbt/service.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def ValidateAlgorithmSettings(self, request, context):
6666
)
6767
if (
6868
"resample_probability" in settings
69-
and not 0 <= settings["resample_probability"] <= 1
69+
and not 0 <= float(settings["resample_probability"]) <= 1
7070
):
7171
return self._set_validate_context_error(
7272
context,
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# Copyright 2026 The Kubeflow Authors.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import unittest
16+
17+
import grpc
18+
import grpc_testing
19+
import utils
20+
21+
from pkg.apis.manager.v1beta1.python import api_pb2
22+
from pkg.suggestion.v1beta1.pbt.service import PbtService
23+
24+
25+
class TestPbt(unittest.TestCase):
26+
def setUp(self):
27+
servicers = {
28+
api_pb2.DESCRIPTOR.services_by_name["Suggestion"]: PbtService()
29+
}
30+
self.test_server = grpc_testing.server_from_dictionary(
31+
servicers, grpc_testing.strict_real_time()
32+
)
33+
34+
def _spec(self, resample_probability=None):
35+
settings = [
36+
api_pb2.AlgorithmSetting(name="suggestion_trial_dir", value="/tmp"),
37+
api_pb2.AlgorithmSetting(name="n_population", value="10"),
38+
api_pb2.AlgorithmSetting(name="truncation_threshold", value="0.2"),
39+
]
40+
if resample_probability is not None:
41+
settings.append(
42+
api_pb2.AlgorithmSetting(
43+
name="resample_probability", value=resample_probability
44+
)
45+
)
46+
return api_pb2.ExperimentSpec(
47+
algorithm=api_pb2.AlgorithmSpec(
48+
algorithm_name="pbt", algorithm_settings=settings
49+
)
50+
)
51+
52+
def test_validate_algorithm_settings(self):
53+
# Valid, without the optional resample_probability.
54+
_, _, code, _ = utils.call_validate(self.test_server, self._spec())
55+
self.assertEqual(code, grpc.StatusCode.OK)
56+
57+
# Valid, with a resample_probability inside [0, 1]. Its value arrives as a
58+
# string, so the range check must cast it before comparing.
59+
_, _, code, _ = utils.call_validate(
60+
self.test_server, self._spec(resample_probability="0.5")
61+
)
62+
self.assertEqual(code, grpc.StatusCode.OK)
63+
64+
# Invalid, resample_probability outside [0, 1].
65+
_, _, code, _ = utils.call_validate(
66+
self.test_server, self._spec(resample_probability="1.5")
67+
)
68+
self.assertEqual(code, grpc.StatusCode.INVALID_ARGUMENT)
69+
70+
71+
if __name__ == "__main__":
72+
unittest.main()

0 commit comments

Comments
 (0)