Allow pinning of the co-located database (#306)

This PR removes the `limit_app_cpus` functionality from the colocated
database. On deeper inspection, this was not functioning as expected
on multiple workfoad manager. The database was being pinned to the
first N logical CPUs on the machine (without the user being able to
change the pinning) and when `limit_app_cpus` was True, only the
launch command and not the launched applications was being pinned.
Additionally, any pinning that the user was trying to do for their Slurm-
launched application, was being silently removed. Both these factors
led to significant performance degradation especially on Cray EX-like
platforms.

To solve this problem, users can now set a `custom_pinning` argument
for the co-located database. By default on Linux, the database will still
be assigned to the first N CPUs, however for machines like Frontier
which reserve some CPUs for system processes, or for advanced users
who might want to pin the DB CPUs to a NUMA node and/or by proximity
to GPUs can now pass in a list of CPU ids for the database to be
assigned.

[ committed by @ashao ]
[ reviewed by @MattToast and @mellis13 ]

conftest.py

@@ -27,7 +27,6 @@
 import json
 import os
 import inspect
-import shutil
 import pytest
 import psutil
 import shutil
@@ -44,6 +43,7 @@
 from smartsim._core.config import CONFIG
 from smartsim.error import SSConfigError
 from subprocess import run
+import sys
 
 
 # Globals, yes, but its a testing file
@@ -513,7 +513,7 @@ def _test_dir_path(caller_function, caller_fspath):
         return dir_path
 
     @staticmethod
-    def get_test_dir(caller_function=None, caller_fspath=None):
+    def get_test_dir(caller_function=None, caller_fspath=None, level=1):
         """Get path to test output.
 
         This function should be called without arguments from within
@@ -532,7 +532,7 @@ def get_test_dir(caller_function=None, caller_fspath=None):
         :rtype: str
         """
         if not caller_function or not caller_fspath:
-            caller_frame = inspect.stack()[1]
+            caller_frame = inspect.stack()[level]
             caller_fspath = caller_frame.filename
             caller_function = caller_frame.function
 
@@ -543,7 +543,7 @@ def get_test_dir(caller_function=None, caller_fspath=None):
         return dir_path
 
     @staticmethod
-    def make_test_dir(caller_function=None, caller_fspath=None):
+    def make_test_dir(caller_function=None, caller_fspath=None, level=1):
         """Create test output directory and return path to it.
 
         This function should be called without arguments from within
@@ -560,7 +560,7 @@ def make_test_dir(caller_function=None, caller_fspath=None):
         :rtype: str
         """
         if not caller_function or not caller_fspath:
-            caller_frame = inspect.stack()[1]
+            caller_frame = inspect.stack()[level]
             caller_fspath = caller_frame.filename
             caller_function = caller_frame.function
 
@@ -597,3 +597,34 @@ def get_test_device():
     @staticmethod
     def get_test_num_gpus():
         return test_num_gpus
+
+@pytest.fixture
+def coloutils():
+    return ColoUtils
+
+class ColoUtils:
+    def setup_test_colo(fileutils, db_type, exp, db_args):
+        """Setup things needed for setting up the colo pinning tests"""
+        # get test setup
+        test_dir = fileutils.make_test_dir(level=2)
+        sr_test_script = fileutils.get_test_conf_path("send_data_local_smartredis.py")
+
+        # Create an app with a colo_db which uses 1 db_cpu
+        colo_settings = exp.create_run_settings(exe=sys.executable, exe_args=sr_test_script)
+        colo_model = exp.create_model(f"colocated_model", colo_settings)
+        colo_model.set_path(test_dir)
+
+        if db_type in ['tcp', "deprecated"]:
+            db_args["port"] = 6780
+            db_args["ifname"] = "lo"
+
+        colocate_fun = {
+            "tcp": colo_model.colocate_db_tcp,
+            "deprecated": colo_model.colocate_db,
+            "uds":colo_model.colocate_db_uds
+        }
+        colocate_fun[db_type](**db_args)
+        # assert model will launch with colocated db
+        assert colo_model.colocated
+        # Check to make sure that limit_db_cpus made it into the colo settings
+        return colo_model

doc/changelog.rst

@@ -26,6 +26,7 @@ Description
 A full list of changes and detailed notes can be found below:
 
 - Fix add_ml_model() and add_script() documentation, tests, and code
+- Replace `limit_app_cpus` with `limit_db_cpus` for co-located orchestrators
 - Remove wait time associated with Experiment launch summary
 - Update and rename Redis conf file
 - Migrate from redis-py-cluster to redis-py
@@ -47,6 +48,10 @@ Detailed notes
   testing interface has been changed to lo instead of ipogif. (PR304_)
 - Typehints have been added. A makefile target `make check-mypy` executes static
   analysis with mypy. (PR295_, PR303_)
+- Replace `limit_app_cpus` with `limit_db_cpus` for co-located orchestrators.
+  This resolves some incorrect behavior/assumptions about how the application
+  would be pinned.  Instead, users should directly specify the binding options in
+  their application using the options appropriate for their launcher
 - Simplify code in `random_permutations` parameter generation strategy (PR300_)
 - Remove wait time associated with Experiment launch summary (PR298_)
 - Update Redis conf file to conform with Redis v7.0.5 conf file (PR293_)
@@ -62,6 +67,8 @@ Detailed notes
   return/error codes. These have now all been updated. (PR284_)
 - Orchestrator and Colocated DB now accept a list of interfaces to bind to. The
   argument name is still `interface` for backward compatibility reasons. (PR281_)
+- Typehints have been added to public APIs. A makefile target to execute static
+  analysis with mypy is available `make check-mypy`. (PR295_)
 
 .. _PR305: https://github.com/CrayLabs/SmartSim/pull/305
 .. _PR304: https://github.com/CrayLabs/SmartSim/pull/304

doc/orchestrator.rst

@@ -84,7 +84,7 @@ the ``Model.colocate_db_tcp`` or ``Model.colocate_db_uds`` function.
     :noindex:
 
 Here is an example of creating a simple model that is colocated with an
-``Orchestrator`` deployment
+``Orchestrator`` deployment using Unix Domain Sockets
 
 .. code-block:: python
 
@@ -94,20 +94,30 @@ Here is an example of creating a simple model that is colocated with an
   colo_settings = exp.create_run_settings(exe="./some_mpi_app")
 
   colo_model = exp.create_model("colocated_model", colo_settings)
-  colo_model.colocate_db_tcp(
-          port=6780,              # database port
+  colo_model.colocate_db_uds(
           db_cpus=1,              # cpus given to the database on each node
           debug=False             # include debug information (will be slower)
-          limit_app_cpus=False,   # don't oversubscribe app with database cpus
           ifname=network_interface # specify network interface(s) to use (i.e. "ib0" or ["ib0", "lo"])
   )
   exp.start(colo_model)
 
 
-By default, SmartSim will attempt to make sure that the database and the application
-do not fight over resources by taking over the affinity mapping process locally on
-each node. This can be disabled by setting ``limit_app_cpus`` to ``False``.
+By default, SmartSim will pin the database to the first _N_ CPUs according to ``db_cpus``. By
+specifying the optional argument ``custom_pinning``, an alternative pinning can be specified
+by sending in a list of CPU ids (e.g [0,2,range(5,8)]). For optimal performance, most users
+will want to also modify the RunSettings for the model to pin their application to cores not
+occupied by the database.
 
+.. warning::
+
+  Pinning is not supported on MacOS X. Setting ``custom_pinning`` to anything
+  other than ``None`` will raise a warning and the input will be ignored.
+
+.. note::
+
+  Pinning _only_ affects the co-located deployment because both the application and the database
+  are sharing the same compute node. For the clustered deployment, a shard occupies the entirerty
+  of the node.
 
 Redis
 =====

smartsim/_core/entrypoints/colocated.py

@@ -202,21 +202,11 @@ def main(
         raise SSInternalError("Colocated process failed to start") from e
 
     try:
-        if sys.platform != "darwin":
-            # Set CPU affinity to the last $db_cpus CPUs
-            affinity = p.cpu_affinity()
-            cpus_to_use = affinity[-db_cpus:] if affinity else None
-            p.cpu_affinity(cpus_to_use)
-        else:
-            # psutil doesn't support pinning on MacOS
-            cpus_to_use = "CPU pinning disabled on MacOS"
-
         logger.debug(
             "\n\nColocated database information\n"
             + "\n".join(
                 (
                     f"\tIP Address(es): {' '.join(ip_addresses + [lo_address])}",
-                    f"\tAffinity: {cpus_to_use}",
                     f"\tCommand: {' '.join(cmd)}\n\n",
                     f"\t# of Database CPUs: {db_cpus}",
                 )

smartsim/_core/launcher/colocated.py

@@ -70,11 +70,8 @@ def write_colocated_launch_script(file_name: str, db_log: str, colocated_setting
         f.write(f"{colocated_cmd}\n")
         f.write(f"DBPID=$!\n\n")
 
-        if colocated_settings["limit_app_cpus"]:
-            cpus = colocated_settings["cpus"]
-            f.write(f"taskset -c 0-$(nproc --ignore={str(cpus+1)}) $@\n\n")
-        else:
-            f.write(f"$@\n\n")
+        # Write the actual launch command for the app
+        f.write(f"$@\n\n")
 
 
 def _build_colocated_wrapper_cmd(
@@ -84,12 +81,13 @@ 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,
+    custom_pinning: t.Optional[str] = None,
     **kwargs: t.Any,
 ) -> str:
     """Build the command use to run a colocated DB application
 
     :param db_log: log file for the db
-    :type db_log: str    
+    :type db_log: str
     :param cpus: db cpus, defaults to 1
     :type cpus: int, optional
     :param rai_args: redisai args, defaults to None
@@ -100,6 +98,8 @@ def _build_colocated_wrapper_cmd(
     :type port: int
     :param ifname: network interface(s) to bind DB to
     :type ifname: str | list[str], 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
     :rtype: str
     """
@@ -114,23 +114,37 @@ def _build_colocated_wrapper_cmd(
     # create the command that will be used to launch the
     # database with the python entrypoint for starting
     # up the backgrounded db process
+
     cmd = [
-        sys.executable,
-        "-m",
-        "smartsim._core.entrypoints.colocated",
-        "+lockfile",
-        lockfile,
-        "+db_cpus",
-        str(cpus),
-    ]
+            sys.executable,
+            "-m",
+            "smartsim._core.entrypoints.colocated",
+            "+lockfile",
+            lockfile,
+            "+db_cpus",
+            str(cpus),
+        ]
     # Add in the interface if using TCP/IP
     if ifname:
         if isinstance(ifname, str):
             ifname = [ifname]
         cmd.extend(["+ifname", ",".join(ifname)])
     cmd.append("+command")
     # collect DB binaries and libraries from the config
-    db_cmd = [CONFIG.database_exe, CONFIG.database_conf, "--loadmodule", CONFIG.redisai]
+
+    db_cmd = []
+    if custom_pinning:
+        db_cmd.extend([
+            'taskset', '-c', custom_pinning
+        ])
+    db_cmd.extend(
+        [
+            CONFIG.database_exe,
+            CONFIG.database_conf,
+            "--loadmodule",
+            CONFIG.redisai
+        ]
+    )
 
     # add extra redisAI configurations
     for arg, value in (rai_args or {}).items():

smartsim/_core/launcher/step/slurmStep.py

@@ -157,10 +157,7 @@ def get_launch_cmd(self) -> t.List[str]:
         srun_cmd += self.run_settings.format_run_args()
 
         if self.run_settings.colocated_db_settings:
-            # disable cpu binding as the entrypoint will set that
-            # for the application and database process now
-            srun_cmd.append("--cpu-bind=none")
-
+
             # Replace the command with the entrypoint wrapper script
             bash = shutil.which("bash")
             launch_script_path = self.get_colocated_launch_script()

smartsim/_core/launcher/step/step.py

@@ -27,6 +27,7 @@
 from __future__ import annotations
 
 import os.path as osp
+import sys
 import time
 import typing as t
 
@@ -81,21 +82,11 @@ def get_colocated_launch_script(self) -> str:
         else:
             db_log_file = "/dev/null"
 
-        # if user specified to use taskset with local launcher
-        # (not allowed b/c MacOS doesn't support it)
-        # TODO: support this only on linux
-        if (
-            self.__class__.__name__ == "LocalStep"
-            and db_settings["limit_app_cpus"] is True
-        ):  # pragma: no cover
-            logger.warning("Setting limit_app_cpus=False for local launcher")
-            db_settings["limit_app_cpus"] = False
-
         # write the colocated wrapper shell script to the directory for this
         # 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