Improve testing robustness on SLURM machines

smartsim/_core/config/config.py

@@ -24,6 +24,7 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import json
 import os
 import psutil
 import typing as t
@@ -179,6 +180,20 @@ def test_num_gpus(self) -> int:  # pragma: no cover
     def test_port(self) -> int:  # pragma: no cover
         return int(os.environ.get("SMARTSIM_TEST_PORT", 6780))
 
+    @property
+    def test_batch_resources(self) -> t.Dict[t.Any,t.Any]: # pragma: no cover
+        resource_str = os.environ.get("SMARTSIM_TEST_BATCH_RESOURCES", "{}")
+        resources = json.loads(resource_str)
+        if not isinstance(resources, dict):
+            raise TypeError(
+                (
+                    "SMARTSIM_TEST_BATCH_RESOURCES was not interpreted as a "
+                    "dictionary, check to make sure that it is a valid "
+                    f"JSON string: {resource_str}"
+                )
+            )
+        return resources
+
     @property
     def test_interface(self) -> t.List[str]:  # pragma: no cover
         if interfaces_cfg := os.environ.get("SMARTSIM_TEST_INTERFACE", None):

smartsim/_core/control/controller.py

@@ -26,6 +26,7 @@
 
 from __future__ import annotations
 
+import itertools
 import os.path as osp
 import pathlib
 import pickle
@@ -39,9 +40,16 @@
 
 from smartredis import Client, ConfigOptions
 
+from smartsim._core.utils.network import get_ip_from_host
+
 from ..._core.launcher.step import Step
 from ..._core.utils.helpers import unpack_colo_db_identifier, unpack_db_identifier
-from ..._core.utils.redis import db_is_active, set_ml_model, set_script, shutdown_db
+from ..._core.utils.redis import (
+    db_is_active,
+    set_ml_model,
+    set_script,
+    shutdown_db_node,
+)
 from ...database import Orchestrator
 from ...entity import (
     Ensemble,
@@ -235,12 +243,22 @@ def stop_db(self, db: Orchestrator) -> None:
         if db.batch:
             self.stop_entity(db)
         else:
-            shutdown_db(db.hosts, db.ports)
             with JM_LOCK:
-                for entity in db:
-                    job = self._jobs[entity.name]
-                    job.set_status(STATUS_CANCELLED, "", 0, output=None, error=None)
-                    self._jobs.move_to_completed(job)
+                for node in db.entities:
+                    for host_ip, port in itertools.product(
+                        (get_ip_from_host(host) for host in node.hosts), db.ports
+                    ):
+                        retcode, _, _ = shutdown_db_node(host_ip, port)
+                        # Sometimes the DB will not shutdown (unless we force NOSAVE)
+                        if retcode != 0:
+                            self.stop_entity(node)
+                            continue
+
+                        job = self._jobs[node.name]
+                        job.set_status(STATUS_CANCELLED, "", 0, output=None, error=None)
+                        self._jobs.move_to_completed(job)
+
+        db.reset_hosts()
 
     def stop_entity_list(self, entity_list: EntitySequence[SmartSimEntity]) -> None:
         """Stop an instance of an entity list
@@ -358,9 +376,9 @@ def _launch(
         for orchestrator in manifest.dbs:
             for key in self._jobs.get_db_host_addresses():
                 _, db_id = unpack_db_identifier(key, "_")
-                if orchestrator.name == db_id:
+                if orchestrator.db_identifier == db_id:
                     raise SSDBIDConflictError(
-                        f"Database identifier {orchestrator.name}"
+                        f"Database identifier {orchestrator.db_identifier}"
                         " has already been used. Pass in a unique"
                         " name for db_identifier"
                     )
@@ -600,30 +618,27 @@ def _prep_entity_client_env(self, entity: Model) -> None:
 
         for db_id, addresses in address_dict.items():
             db_name, _ = unpack_db_identifier(db_id, "_")
-
             if addresses:
-                if len(addresses) <= 128:
-                    client_env[f"SSDB{db_name}"] = ",".join(addresses)
-                else:
-                    # Cap max length of SSDB
-                    client_env[f"SSDB{db_name}"] = ",".join(addresses[:128])
-                if entity.incoming_entities:
-                    client_env[f"SSKEYIN{db_name}"] = ",".join(
-                        [in_entity.name for in_entity in entity.incoming_entities]
-                    )
-                if entity.query_key_prefixing():
-                    client_env[f"SSKEYOUT{db_name}"] = entity.name
+                # Cap max length of SSDB
+                client_env[f"SSDB{db_name}"] = ",".join(addresses[:128])
+
+                # Retrieve num_shards to append to client env
+                client_env[f"SR_DB_TYPE{db_name}"] = (
+                    CLUSTERED if len(addresses) > 1 else STANDALONE
+                )
 
-            # Retrieve num_shards to append to client env
-            client_env[f"SR_DB_TYPE{db_name}"] = (
-                CLUSTERED if len(addresses) > 1 else STANDALONE
+        if entity.incoming_entities:
+            client_env["SSKEYIN"] = ",".join(
+                [in_entity.name for in_entity in entity.incoming_entities]
             )
+        if entity.query_key_prefixing():
+            client_env["SSKEYOUT"] = entity.name
 
         # Set address to local if it's a colocated model
         if entity.colocated and entity.run_settings.colocated_db_settings is not None:
             db_name_colo = entity.run_settings.colocated_db_settings["db_identifier"]
 
-            for key in self._jobs.get_db_host_addresses():
+            for key in address_dict:
                 _, db_id = unpack_db_identifier(key, "_")
                 if db_name_colo == db_id:
                     raise SSDBIDConflictError(

smartsim/_core/control/jobmanager.py

@@ -34,9 +34,8 @@
 
 from ...database import Orchestrator
 from ...entity import DBNode, SmartSimEntity, EntitySequence
-from ...error import SmartSimError
 from ...log import get_logger
-from ...status import TERMINAL_STATUSES
+from ...status import TERMINAL_STATUSES, STATUS_NEVER_STARTED
 from ..config import CONFIG
 from ..launcher import LocalLauncher, Launcher
 from ..utils.network import get_ip_from_host
@@ -160,6 +159,13 @@ def __call__(self) -> t.Dict[str, Job]:
         all_jobs = {**self.jobs, **self.db_jobs}
         return all_jobs
 
+    def __contains__(self, key: str) -> bool:
+        try:
+            self[key]  # pylint: disable=pointless-statement
+            return True
+        except KeyError:
+            return False
+
     def add_job(
         self,
         job_name: str,
@@ -242,17 +248,14 @@ def get_status(
         :returns: tuple of status
         """
         with self._lock:
-            try:
-                if entity.name in self.completed:
-                    return self.completed[entity.name].status
+            if entity.name in self.completed:
+                return self.completed[entity.name].status
 
+            if entity.name in self:
                 job: Job = self[entity.name]  # locked
-            except KeyError:
-                raise SmartSimError(
-                    f"Entity {entity.name} has not been launched in this Experiment"
-                ) from None
+                return job.status
 
-            return job.status
+            return STATUS_NEVER_STARTED
 
     def set_launcher(self, launcher: Launcher) -> None:
         """Set the launcher of the job manager to a specific launcher instance
@@ -312,7 +315,7 @@ def get_db_host_addresses(self) -> t.Dict[str, t.List[str]]:
         :rtype: Dict[str, list]
         """
 
-        address_dict = {}
+        address_dict: t.Dict[str, t.List[str]] = {}
         for db_job in self.db_jobs.values():
             addresses = []
             if isinstance(db_job.entity, (DBNode, Orchestrator)):
@@ -321,7 +324,9 @@ def get_db_host_addresses(self) -> t.Dict[str, t.List[str]]:
                     ip_addr = get_ip_from_host(combine[0])
                     addresses.append(":".join((ip_addr, str(combine[1]))))
 
-            address_dict.update({db_entity.name: addresses})
+                dict_entry: t.List[str] = address_dict.get(db_entity.db_identifier, [])
+                dict_entry.extend(addresses)
+                address_dict[db_entity.db_identifier] = dict_entry
 
         return address_dict
 

smartsim/_core/utils/helpers.py

@@ -40,28 +40,28 @@
 
 
 def unpack_db_identifier(db_id: str, token: str) -> t.Tuple[str, str]:
-    """Unpack the unformatted database identifier using the token,
-    and format for env variable suffix
-    :db_id: the unformatted database identifier eg. identifier_1_0
-    :token: character '_' or '-' to use to unpack the database identifier
-    :return: db suffix, and formatted db_id eg. _identifier_1, identifier_1
+    """Unpack the unformatted database identifier
+    and format for env variable suffix using the token
+    :param db_id: the unformatted database identifier eg. identifier_1
+    :type db_id: str
+    :param token: character to use to construct the db suffix
+    :type token: str
+    :return: db id suffix and formatted db_id e.g. ("_identifier_1", "identifier_1")
+    :rtype: (str, str)
     """
 
     if db_id == "orchestrator":
         return "", ""
-    db_id = "_".join(db_id.split(token)[:-1])
-    # if unpacked db_id is default, return empty
-    if db_id == "orchestrator":
-        # if db_id is default after unpack, return empty
-        return "", ""
-    db_name_suffix = "_" + db_id
+    db_name_suffix = token + db_id
     return db_name_suffix, db_id
 
 
 def unpack_colo_db_identifier(db_id: str) -> str:
     """Create database identifier suffix for colocated database
-    :db_id: the unformatted database identifier
+    :param db_id: the unformatted database identifier
+    :type db_id: str
     :return: db suffix
+    :rtype: str
     """
     return "_" + db_id if db_id else ""
 

smartsim/_core/utils/redis.py

@@ -24,7 +24,6 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-import itertools
 import logging
 import redis
 import time
@@ -219,31 +218,30 @@ def set_script(db_script: DBScript, client: Client) -> None:
             raise error
 
 
-def shutdown_db(hosts: t.List[str], ports: t.List[int]) -> None:  # cov-wlm
-    """Send shutdown signal to cluster instances.
+def shutdown_db_node(host_ip: str, port: int) -> t.Tuple[int, str, str]:  # cov-wlm
+    """Send shutdown signal to DB node.
 
     Should only be used in the case where cluster deallocation
-    needs to occur manually. Usually, the SmartSim task manager
+    needs to occur manually. Usually, the SmartSim job manager
     will take care of this automatically.
 
-    :param hosts: List of hostnames to connect to
-    :type hosts: List[str]
-    :param ports: List of ports for each hostname
-    :type ports: List[int]
-    :raises SmartSimError: if cluster creation fails
+    :param host_ip: IP of host to connect to
+    :type hosts: str
+    :param ports: Port to which node is listening
+    :type ports: int
+    :return: returncode, output, and error of the process
+    :rtype: tuple of (int, str, str)
     """
-    for host_ip, port in itertools.product(
-        (get_ip_from_host(host) for host in hosts), ports
-    ):
-        # call cluster command
-        redis_cli = CONFIG.database_cli
-        cmd = [redis_cli, "-h", host_ip, "-p", str(port), "shutdown"]
-        returncode, out, err = execute_cmd(
-            cmd, proc_input="yes", shell=False, timeout=10
-        )
+    redis_cli = CONFIG.database_cli
+    cmd = [redis_cli, "-h", host_ip, "-p", str(port), "shutdown"]
+    returncode, out, err = execute_cmd(
+        cmd, proc_input="yes", shell=False, timeout=10
+    )
+
+    if returncode != 0:
+        logger.error(out)
+        logger.error(err)
+    elif out:
+        logger.debug(out)
 
-        if returncode != 0:
-            logger.error(out)
-            logger.error(err)
-        else:
-            logger.debug(out)
+    return returncode, out, err