diff --git a/conftest.py b/conftest.py index 123fafb4a..f2e73c028 100644 --- a/conftest.py +++ b/conftest.py @@ -113,8 +113,7 @@ def pytest_sessionfinish(session, exitstatus): returning the exit status to the system. """ if exitstatus == 0: - # shutil.rmtree(test_dir) - pass + shutil.rmtree(test_dir) else: # kill all spawned processes in case of error kill_all_test_spawned_processes() @@ -598,6 +597,7 @@ def coloutils(): return ColoUtils class ColoUtils: + @staticmethod def setup_test_colo(fileutils, db_type, exp, db_args): """Setup things needed for setting up the colo pinning tests""" # get test setup diff --git a/smartsim/_core/launcher/colocated.py b/smartsim/_core/launcher/colocated.py index ee9ec4863..de6946088 100644 --- a/smartsim/_core/launcher/colocated.py +++ b/smartsim/_core/launcher/colocated.py @@ -81,8 +81,7 @@ def _build_colocated_wrapper_cmd( extra_db_args: t.Optional[t.Dict[str, str]] = None, port: int = 6780, ifname: t.Optional[t.Union[str, t.List[str]]] = None, - limit_db_cpus: bool = False, - db_cpu_list: t.Optional[str] = None, + custom_pinning: t.Optional[str] = None, **kwargs: t.Any, ) -> str: """Build the command use to run a colocated DB application @@ -99,8 +98,6 @@ def _build_colocated_wrapper_cmd( :type port: int :param ifname: network interface(s) to bind DB to :type ifname: str | list[str], optional - :param limit_db_cpus: If True, limit the cpus that the database can run on - :type limit_db_cpus: bool, optional :param db_cpu_list: The list of CPUs that the database should be limited to :type db_cpu_list: str, optional :return: the command to run @@ -114,16 +111,11 @@ def _build_colocated_wrapper_cmd( # the lock on the file. lockfile = create_lockfile_name() - - - # create the command that will be used to launch the # database with the python entrypoint for starting # up the backgrounded db process - cmd = [] - cmd.extend( - [ + cmd = [ sys.executable, "-m", "smartsim._core.entrypoints.colocated", @@ -132,7 +124,6 @@ def _build_colocated_wrapper_cmd( "+db_cpus", str(cpus), ] - ) # Add in the interface if using TCP/IP if ifname: if isinstance(ifname, str): @@ -142,9 +133,9 @@ def _build_colocated_wrapper_cmd( # collect DB binaries and libraries from the config db_cmd = [] - if limit_db_cpus and db_cpu_list: + if custom_pinning: db_cmd.extend([ - 'taskset', '-c', db_cpu_list + 'taskset', '-c', custom_pinning ]) db_cmd.extend([ CONFIG.database_exe, diff --git a/smartsim/_core/launcher/step/step.py b/smartsim/_core/launcher/step/step.py index bba8d85ae..84d6411a2 100644 --- a/smartsim/_core/launcher/step/step.py +++ b/smartsim/_core/launcher/step/step.py @@ -83,8 +83,7 @@ def get_colocated_launch_script(self) -> str: else: db_log_file = "/dev/null" - if sys.platform == 'darwin': - if isinstance(self, LocalStep) and db_settings["limit_db_cpus"] is True: + if sys.platform == 'darwin' and db_settings["db_cpu_list"]: logger.warning( "DB pinning is not supported on MacOS, setting limit_db_cpus = False") db_settings["limit_db_cpus"] = False @@ -93,7 +92,7 @@ def get_colocated_launch_script(self) -> str: # entity currently being prepped to launch write_colocated_launch_script(script_path, db_log_file, db_settings) return script_path - + def add_to_batch(self, step: Step) -> None: """Add a job step to this batch diff --git a/smartsim/entity/model.py b/smartsim/entity/model.py index 078cc3de7..c7d4ff3a1 100644 --- a/smartsim/entity/model.py +++ b/smartsim/entity/model.py @@ -26,8 +26,10 @@ from __future__ import annotations +import itertools import typing as t import warnings +import collections.abc from .._core.utils.helpers import cat_arg_and_value, init_default from ..error import EntityExistsError, SSUnsupportedError @@ -154,10 +156,11 @@ def attach_generator_files( def colocate_db(self, *args: t.Any, **kwargs: t.Any) -> None: """An alias for ``Model.colocate_db_tcp``""" - logger.warning( + warnings.warn( ( "`colocate_db` has been deprecated and will be removed in a \n" - "future release. Please use `colocate_db_tcp` or `colocate_db_uds`." + "future release. Please use `colocate_db_tcp` or `colocate_db_uds`.", + FutureWarning ) ) self.colocate_db_tcp(*args, **kwargs) @@ -167,8 +170,7 @@ def colocate_db_uds( unix_socket: str = "/tmp/redis.socket", socket_permissions: int = 755, db_cpus: int = 1, - limit_db_cpus: bool = True, - db_cpu_list: t.Optional[str] = None, + custom_pinning: t.Optional[t.Iterable[t.Union(int, t.Iterable[int])]] = None, debug: bool = False, **kwargs: t.Any, ) -> None: @@ -200,11 +202,9 @@ def colocate_db_uds( :type socket_permissions: int, optional :param db_cpus: number of cpus to use for orchestrator, defaults to 1 :type db_cpus: int, optional - :param limit_db_cpus: whether to limit the number of cpus used by the database defaults to True - :type limit_db_cpus: bool, optional - :param db_cpu_list: lists the cpus which the database can be run on. Follows `taskset -c` syntax - e.g. '0-2,5' specifies processors 0, 1, 2, and 5 - :type db_cpu_list: str, optional + :param custom_pinning: CPUs to pin the orchestrator to. Passing an empty iterable + disables pinning + :type custom_pinning: iterable of ints or iterable of ints, optional :param debug: launch Model with extra debug information about the colocated db :type debug: bool, optional :param kwargs: additional keyword arguments to pass to the orchestrator database @@ -219,8 +219,7 @@ def colocate_db_uds( common_options = { "cpus": db_cpus, - "limit_db_cpus": limit_db_cpus, - "db_cpu_list": db_cpu_list, + "custom_pinning": custom_pinning, "debug": debug, } self._set_colocated_db_settings(uds_options, common_options, **kwargs) @@ -230,8 +229,7 @@ def colocate_db_tcp( port: int = 6379, ifname: str = "lo", db_cpus: int = 1, - limit_db_cpus: bool = True, - db_cpu_list: t.Optional[str] = None, + custom_pinning: t.Optional[t.Iterable[t.Union(int, t.Iterable[int])]] = None, debug: bool = False, **kwargs: t.Any, ) -> None: @@ -263,11 +261,9 @@ def colocate_db_tcp( :type ifname: str, optional :param db_cpus: number of cpus to use for orchestrator, defaults to 1 :type db_cpus: int, optional - :param limit_db_cpus: whether to limit the number of cpus used by the database, defaults to True - :type limit_db_cpus: bool, optional - :param db_cpu_list: lists the cpus which the database can be run on. Follows `taskset -c` syntax - e.g. '0-2,5' specifies processors 0, 1, 2, and 5 - :type db_cpu_list: str, optional + :param custom_pinning: CPUs to pin the orchestrator to. Passing an empty iterable + disables pinning + :type custom_pinning: iterable of ints or iterable of ints, optional :param debug: launch Model with extra debug information about the colocated db :type debug: bool, optional :param kwargs: additional keyword arguments to pass to the orchestrator database @@ -278,8 +274,7 @@ def colocate_db_tcp( tcp_options = {"port": port, "ifname": ifname} common_options = { "cpus": db_cpus, - "limit_db_cpus": limit_db_cpus, - "db_cpu_list": db_cpu_list, + "custom_pinning": custom_pinning, "debug": debug, } self._set_colocated_db_settings(tcp_options, common_options, **kwargs) @@ -305,36 +300,15 @@ def _set_colocated_db_settings( if "limit_app_cpus" in common_options: raise SSUnsupportedError( - "Pinning of app CPUs via limit_app_cpus is no supported. Modify RunSettings " + + "Pinning of app CPUs via limit_app_cpus is no supported. Modify RunSettings " "instead using the correct binding option for your launcher." ) # TODO list which db settings can be extras - cpus = common_options["cpus"] - # Deal with cases where the database should be pinned to cpus - # (1) if the user set a db_cpu_list, but not limit_db_cpus - if common_options["db_cpu_list"] and not common_options["limit_db_cpus"]: - logger.warning( - "limit_db_cpus is False, but db_cpu_list is not None. Setting limit_db_cpus=True" - ) - common_options["limit_db_cpus"] = True - # (2) limit_db_cpus, but not db_cpu_list is specified automatically set - # pin to the cpus 0:db_cpus-1 - if common_options["limit_db_cpus"] and not common_options["db_cpu_list"]: - if cpus > 1: - cpu_list = f"0-{cpus-1}" - elif cpus == 1: - cpu_list = "0" - else: - raise ValueError("db_cpus must be a positive number") - - logger.warning( - ( - "limit_db_cpus is True, but db_cpu_list was not specified. Automatically " - f"pinning to processors {cpu_list}" - ) - ) - common_options["db_cpu_list"] = cpu_list + common_options["custom_pinning"] = self._create_pinning_string( + common_options["custom_pinning"], + common_options["cpus"] + ) colo_db_config = {} colo_db_config.update(connection_options) @@ -359,6 +333,30 @@ def _set_colocated_db_settings( self.run_settings.colocated_db_settings = colo_db_config + @staticmethod + def _create_pinning_string( + pin_ids: t.Optional[t.Iterable[t.Union(int, t.Iterable[int])]], + cpus: int + ): + """Create a comma-separated string CPU ids. By default, None returns + 0,1,...,cpus-1; an empty iterable will disable pinning altogether, + and an iterable constructs a comma separate string (e.g. 0,2,5) + """ + if pin_ids is None: + return ','.join(str(i) for i in range(cpus)) + elif not pin_ids: + return None + elif isinstance(pin_ids, collections.abc.Iterable): + pin_list = [] + for i in pin_ids: + if isinstance(i, collections.abc.Iterable): + pin_list.extend([str(j) for j in i]) + else: + pin_list.append(str(i)) + return ','.join(pin_list) + else: + raise TypeError("pin_ids must be an iterable of ints") + def params_to_args(self) -> None: """Convert parameters to command line arguments and update run settings.""" if self.params_as_args is not None: diff --git a/tests/backends/test_dbmodel.py b/tests/backends/test_dbmodel.py index f1a2431ad..18b5213e9 100644 --- a/tests/backends/test_dbmodel.py +++ b/tests/backends/test_dbmodel.py @@ -369,7 +369,6 @@ def test_colocated_db_model_pytorch(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port(), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -412,7 +411,6 @@ def test_colocated_db_model_ensemble(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port(), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -424,7 +422,6 @@ def test_colocated_db_model_ensemble(fileutils, wlmutils): entity.colocate_db( port=wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -454,7 +451,6 @@ def test_colocated_db_model_ensemble(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port() + len(colo_ensemble), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -508,7 +504,6 @@ def test_colocated_db_model_ensemble_reordered(fileutils, wlmutils): entity.colocate_db( wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -527,7 +522,6 @@ def test_colocated_db_model_ensemble_reordered(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port() + len(colo_ensemble), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -564,7 +558,6 @@ def test_colocated_db_model_errors(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port(), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -584,7 +577,6 @@ def test_colocated_db_model_errors(fileutils, wlmutils): entity.colocate_db( port=wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -613,7 +605,6 @@ def test_colocated_db_model_errors(fileutils, wlmutils): entity.colocate_db( port=wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) diff --git a/tests/backends/test_dbscript.py b/tests/backends/test_dbscript.py index d9e6c4b4f..46a4c2aa4 100644 --- a/tests/backends/test_dbscript.py +++ b/tests/backends/test_dbscript.py @@ -153,7 +153,6 @@ def test_colocated_db_script(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port(), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -204,7 +203,6 @@ def test_colocated_db_script_ensemble(fileutils, wlmutils): entity.colocate_db( port=wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -214,7 +212,6 @@ def test_colocated_db_script_ensemble(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port() + len(colo_ensemble), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -268,7 +265,6 @@ def test_colocated_db_script_ensemble_reordered(fileutils, wlmutils): entity.colocate_db( port=wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -278,7 +274,6 @@ def test_colocated_db_script_ensemble_reordered(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port() + len(colo_ensemble), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -315,7 +310,6 @@ def test_db_script_errors(fileutils, wlmutils): colo_model.colocate_db( port=wlmutils.get_test_port(), db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -335,7 +329,6 @@ def test_db_script_errors(fileutils, wlmutils): entity.colocate_db( port=wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) @@ -358,7 +351,6 @@ def test_db_script_errors(fileutils, wlmutils): entity.colocate_db( port=wlmutils.get_test_port() + i, db_cpus=1, - limit_db_cpus=False, debug=True, ifname="lo", ) diff --git a/tests/conftest/setup_test_colo/colocated_model.out b/tests/conftest/setup_test_colo/colocated_model.out deleted file mode 100644 index b8829b436..000000000 --- a/tests/conftest/setup_test_colo/colocated_model.out +++ /dev/null @@ -1,4 +0,0 @@ -SmartRedis Library@19-19-36:WARNING: Environment variable SR_LOG_FILE is not set. Defaulting to stdout -SmartRedis Library@19-19-36:WARNING: Environment variable SR_LOG_LEVEL is not set. Defaulting to INFO -default@19-19-36:ERROR: Unable to connect to backend database: Failed to connect to Redis: Connection refused -Test worked! Sent and received array: [1 2 3 4] diff --git a/tests/on_wlm/test_colocated_model.py b/tests/on_wlm/test_colocated_model.py index 2db956a1d..5336e0cca 100644 --- a/tests/on_wlm/test_colocated_model.py +++ b/tests/on_wlm/test_colocated_model.py @@ -25,25 +25,30 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import sys -import warnings import pytest from smartsim import Experiment, status +from smartsim.entity import Model -launcher = pytest.test_launcher -supported_dbs = ["uds", "tcp", "deprecated"] +if sys.platform == "darwin": + supported_dbs = ["tcp", "deprecated"] +else: + supported_dbs = ["uds", "tcp", "deprecated"] + supported_dbs = ["uds"] # retrieved from pytest fixtures +launcher = pytest.test_launcher if launcher not in pytest.wlm_options: pytestmark = pytest.mark.skip(reason="Not testing WLM integrations") -@pytest.mark.parametrize("db_type", ["uds", "tcp", "deprecated"]) -def test_launch_colocated_model(fileutils, coloutils, db_type): - """Test the launch of a model with a colocated database""" +@pytest.mark.parametrize("db_type", supported_dbs) +def test_launch_colocated_model_defaults(fileutils, coloutils, db_type): + """Test the launch of a model with a colocated database and local launcher""" - exp = Experiment("colocated_model_wlm", launcher=launcher) db_args = { } + + exp = Experiment("colocated_model_defaults", launcher=launcher) colo_model = coloutils.setup_test_colo( fileutils, db_type, @@ -51,6 +56,7 @@ def test_launch_colocated_model(fileutils, coloutils, db_type): db_args, ) + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0" exp.start(colo_model, block=True) statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) @@ -60,15 +66,36 @@ def test_launch_colocated_model(fileutils, coloutils, db_type): statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) -@pytest.mark.parametrize("db_type", ["uds", "tcp", "deprecated"]) -def test_launch_colocated_pinned_model(fileutils, coloutils, db_type): +@pytest.mark.parametrize("db_type", supported_dbs) +def test_colocated_model_disable_pinning(fileutils, coloutils, db_type): + + exp = Experiment("colocated_model_pinning_auto_1cpu", launcher=launcher) db_args = { - "limit_db_cpus": True, - "db_cpus": 2, + "db_cpus": 1, + "custom_pinning": [], } + # Check to make sure that the CPU mask was correctly generated + colo_model = coloutils.setup_test_colo( + fileutils, + db_type, + exp, + db_args, + ) + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] is None + exp.start(colo_model, block=True) + statuses = exp.get_status(colo_model) + assert all([stat == status.STATUS_COMPLETED for stat in statuses]) + +@pytest.mark.parametrize("db_type", supported_dbs) +def test_colocated_model_pinning_auto_2cpu(fileutils, coloutils, db_type): + exp = Experiment("colocated_model_pinning_auto_2cpu", launcher=launcher) + db_args = { + "db_cpus": 2, + } + # Check to make sure that the CPU mask was correctly generated colo_model = coloutils.setup_test_colo( fileutils, @@ -76,31 +103,76 @@ def test_launch_colocated_pinned_model(fileutils, coloutils, db_type): exp, db_args, ) - assert colo_model.run_settings.colocated_db_settings["db_cpu_list"] == "0-1" - assert colo_model.run_settings.colocated_db_settings["limit_db_cpus"] + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0,1" exp.start(colo_model, block=True) statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) @pytest.mark.parametrize("db_type", supported_dbs) -def test_colocated_model_pinning_manual(fileutils, coloutils, db_type): +def test_colocated_model_pinning_range(fileutils, coloutils, db_type): # Check to make sure that the CPU mask was correctly generated + # Assume that there are at least 4 cpus on the node exp = Experiment("colocated_model_pinning_manual", launcher=launcher) + + db_args = { + "db_cpus": 4, + "custom_pinning": range(5) + } + + colo_model = coloutils.setup_test_colo( + fileutils, + db_type, + exp, + db_args, + ) + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0,1,2,3" + exp.start(colo_model, block=True) + statuses = exp.get_status(colo_model) + assert all([stat == status.STATUS_COMPLETED for stat in statuses]) + +@pytest.mark.parametrize("db_type", supported_dbs) +def test_colocated_model_pinning_list(fileutils, coloutils, db_type): + # Check to make sure that the CPU mask was correctly generated + # note we presume that this has more than 2 CPUs on the supercomputer node + + exp = Experiment("colocated_model_pinning_manual", launcher=launcher) + db_args = { - "limit_db_cpus": True, "db_cpus": 2, - "db_cpu_list": "0,2" + "custom_pinning": [0,2] } - colo_model = coloutils._setup_test_colo( + colo_model = coloutils.setup_test_colo( fileutils, db_type, - "colocated_model_pinning_manual", + exp, db_args, ) - assert colo_model.run_settings.colocated_db_settings["db_cpu_list"] == "0,2" - assert colo_model.run_settings.colocated_db_settings["limit_db_cpus"] + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0,2" exp.start(colo_model, block=True) statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) + +@pytest.mark.parametrize("db_type", supported_dbs) +def test_colocated_model_pinning_mixed(fileutils, coloutils, db_type): + # Check to make sure that the CPU mask was correctly generated + # note we presume that this at least 4 CPUs on the supercomputer node + + exp = Experiment("colocated_model_pinning_manual", launcher=launcher) + + db_args = { + "db_cpus": 2, + "custom_pinning": [range(2), 3] + } + + colo_model = coloutils.setup_test_colo( + fileutils, + db_type, + exp, + db_args, + ) + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0,1,3" + exp.start(colo_model, block=True) + statuses = exp.get_status(colo_model) + assert all([stat == status.STATUS_COMPLETED for stat in statuses]) \ No newline at end of file diff --git a/tests/test_colo_model_local.py b/tests/test_colo_model_local.py index 01263aef1..6024fcd6f 100644 --- a/tests/test_colo_model_local.py +++ b/tests/test_colo_model_local.py @@ -29,11 +29,24 @@ import pytest from smartsim import Experiment, status +from smartsim.entity import Model if sys.platform == "darwin": supported_dbs = ["tcp", "deprecated"] else: supported_dbs = ["uds", "tcp", "deprecated"] + +def test_create_pinning_string(): + # Automatic creation of pinned cpu list + assert Model._create_pinning_string(None,2) == '0,1' + # Individual ids only + assert Model._create_pinning_string([1,2],2) == '1,2' + # Mixed ranges and individual ids + assert Model._create_pinning_string([range(2),3],3) == '0,1,3' + # Range only + assert Model._create_pinning_string(range(3),3) == '0,1,2' + + @pytest.mark.parametrize("db_type", supported_dbs) def test_launch_colocated_model_defaults(fileutils, coloutils, db_type, launcher="local"): """Test the launch of a model with a colocated database and local launcher""" @@ -48,6 +61,7 @@ def test_launch_colocated_model_defaults(fileutils, coloutils, db_type, launcher db_args, ) + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0" exp.start(colo_model, block=True) statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) @@ -58,12 +72,12 @@ def test_launch_colocated_model_defaults(fileutils, coloutils, db_type, launcher assert all([stat == status.STATUS_COMPLETED for stat in statuses]) @pytest.mark.parametrize("db_type", supported_dbs) -def test_colocated_model_pinning_auto_1cpu(fileutils, coloutils, db_type, launcher="local"): +def test_colocated_model_disable_pinning(fileutils, coloutils, db_type, launcher="local"): exp = Experiment("colocated_model_pinning_auto_1cpu", launcher=launcher) db_args = { - "limit_db_cpus": True, "db_cpus": 1, + "custom_pinning": [], } # Check to make sure that the CPU mask was correctly generated @@ -73,8 +87,7 @@ def test_colocated_model_pinning_auto_1cpu(fileutils, coloutils, db_type, launch exp, db_args, ) - assert colo_model.run_settings.colocated_db_settings["db_cpu_list"] == "0" - assert colo_model.run_settings.colocated_db_settings["limit_db_cpus"] + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] is None exp.start(colo_model, block=True) statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) @@ -85,7 +98,6 @@ def test_colocated_model_pinning_auto_2cpu(fileutils, coloutils, db_type, launch exp = Experiment("colocated_model_pinning_auto_2cpu", launcher=launcher) db_args = { - "limit_db_cpus": True, "db_cpus": 2, } @@ -96,22 +108,42 @@ def test_colocated_model_pinning_auto_2cpu(fileutils, coloutils, db_type, launch exp, db_args, ) - assert colo_model.run_settings.colocated_db_settings["db_cpu_list"] == "0-1" - assert colo_model.run_settings.colocated_db_settings["limit_db_cpus"] + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0,1" exp.start(colo_model, block=True) statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) @pytest.mark.parametrize("db_type", supported_dbs) -def test_colocated_model_pinning_manual(fileutils, coloutils, db_type, launcher="local"): +def test_colocated_model_pinning_range(fileutils, coloutils, db_type, launcher="local"): # Check to make sure that the CPU mask was correctly generated exp = Experiment("colocated_model_pinning_manual", launcher=launcher) db_args = { - "limit_db_cpus": True, "db_cpus": 2, - "db_cpu_list": "0,2" + "custom_pinning": range(2) + } + + colo_model = coloutils.setup_test_colo( + fileutils, + db_type, + exp, + db_args, + ) + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "0,1" + exp.start(colo_model, block=True) + statuses = exp.get_status(colo_model) + assert all([stat == status.STATUS_COMPLETED for stat in statuses]) + +@pytest.mark.parametrize("db_type", supported_dbs) +def test_colocated_model_pinning_list(fileutils, coloutils, db_type, launcher="local"): + # Check to make sure that the CPU mask was correctly generated + + exp = Experiment("colocated_model_pinning_manual", launcher=launcher) + + db_args = { + "db_cpus": 1, + "custom_pinning": [1] } colo_model = coloutils.setup_test_colo( @@ -120,8 +152,7 @@ def test_colocated_model_pinning_manual(fileutils, coloutils, db_type, launcher= exp, db_args, ) - assert colo_model.run_settings.colocated_db_settings["db_cpu_list"] == "0,2" - assert colo_model.run_settings.colocated_db_settings["limit_db_cpus"] + assert colo_model.run_settings.colocated_db_settings["custom_pinning"] == "1" exp.start(colo_model, block=True) statuses = exp.get_status(colo_model) assert all([stat == status.STATUS_COMPLETED for stat in statuses]) \ No newline at end of file