Merge 005df58 into 5b356c6

requirements/TUNE_REQUIREMENTS.txt

@@ -1,4 +1,4 @@
 mypy_extensions~=0.4; python_version < '3.8'
 optuna~=2.10
 torch>=1.8.1
-ax-platform~=0.2; python_version >= '3.7'
+ax-platform~=0.2.2; python_version >= '3.7'

tests/tune/base_asserts_test.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-from math import log10
+from spock.addons.tune.ax import AxInterface
 
 
 class AllTypes:
@@ -44,8 +44,12 @@ def test_hp_two(self, arg_builder):
 
 class SampleTypes:
     def test_sampling(self, arg_builder):
-        # Draw 100 random samples and make sure all fall within all of the bounds or sets
-        for _ in range(25):
+        # Draw random samples and make sure all fall within all of the bounds or sets
+        if isinstance(arg_builder._tuner_interface._lib_interface, AxInterface):
+            max_draws = arg_builder._tuner_interface.tuner_status['client'].generation_strategy.current_generator_run_limit()[0]
+        else:
+            max_draws = 25
+        for _ in range(max_draws):
             hp_attrs = arg_builder.sample()
             assert 10 <= hp_attrs.HPOne.hp_int <= 100
             assert isinstance(hp_attrs.HPOne.hp_int, int) is True

tests/tune/test_ax.py

@@ -14,6 +14,8 @@
 from tests.tune.attr_configs_test import *
 from tests.tune.base_asserts_test import *
 
+from ax.modelbridge.generation_strategy import GenerationStrategy, GenerationStep
+
 
 class TestAxBasic(AllTypes):
     @staticmethod