Add memory & conn collector, collector manager, tests

.gitignore

@@ -36,3 +36,6 @@ smartsim/_core/lib
 **/*.err
 **/*.out
 **/.smartsim/*
+
+# optional dev tools
+.pre-commit-config.yaml

conftest.py

@@ -26,33 +26,38 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import os
-import pytest
-import psutil
+import pathlib
 import shutil
+import sys
+import tempfile
+import typing as t
+import uuid
+import warnings
+from subprocess import run
+
+import psutil
+import pytest
+
 import smartsim
 from smartsim import Experiment
-from smartsim.entity import Model
+from smartsim._core.config import CONFIG
+from smartsim._core.config.config import Config
+from smartsim._core.utils.telemetry.telemetry import JobEntity
 from smartsim.database import Orchestrator
+from smartsim.entity import Model
+from smartsim.error import SSConfigError
 from smartsim.settings import (
-    SrunSettings,
     AprunSettings,
     JsrunSettings,
-    MpirunSettings,
     MpiexecSettings,
+    MpirunSettings,
     PalsMpiexecSettings,
     RunSettings,
+    SrunSettings,
 )
-from smartsim._core.config import CONFIG
-from smartsim.error import SSConfigError
-from subprocess import run
-import sys
-import tempfile
-import typing as t
-import uuid
-import warnings
-
 
 # pylint: disable=redefined-outer-name,invalid-name,global-statement
 
@@ -66,12 +71,13 @@
 test_alloc_specs_path = os.getenv("SMARTSIM_TEST_ALLOC_SPEC_SHEET_PATH", None)
 test_port = CONFIG.test_port
 test_account = CONFIG.test_account or ""
-test_batch_resources: t.Dict[t.Any,t.Any] = CONFIG.test_batch_resources
+test_batch_resources: t.Dict[t.Any, t.Any] = CONFIG.test_batch_resources
 
 # Fill this at runtime if needed
 test_hostlist = None
 has_aprun = shutil.which("aprun") is not None
 
+
 def get_account() -> str:
     return test_account
 
@@ -368,6 +374,7 @@ def choose_host(rs: RunSettings) -> t.Optional[str]:
 
         return None
 
+
 @pytest.fixture
 def local_db(
     request: t.Any, wlmutils: t.Type[WLMUtils], test_dir: str
@@ -524,7 +531,7 @@ def _sanitize_caller_function(caller_function: str) -> str:
     # We split at the opening bracket, sanitize the string
     # to its right and then merge the function name and
     # the sanitized list with a dot.
-    caller_function = caller_function.replace("]","")
+    caller_function = caller_function.replace("]", "")
     caller_function_list = caller_function.split("[", maxsplit=1)
 
     def is_accepted_char(char: str) -> bool:
@@ -559,7 +566,8 @@ class FileUtils:
     @staticmethod
     def get_test_output_path(caller_function: str, caller_fspath: str) -> str:
         caller_file_to_dir = os.path.splitext(str(caller_fspath))[0]
-        rel_path = os.path.relpath(caller_file_to_dir, os.path.dirname(test_output_root))
+        dir_name = os.path.dirname(test_output_root)
+        rel_path = os.path.relpath(caller_file_to_dir, dir_name)
         dir_path = os.path.join(test_output_root, rel_path, caller_function)
         return dir_path
 
@@ -574,7 +582,9 @@ def get_test_dir_path(dirname: str) -> str:
         return dir_path
 
     @staticmethod
-    def make_test_file(file_name: str, file_dir: str, file_content: t.Optional[str] = None) -> str:
+    def make_test_file(
+        file_name: str, file_dir: str, file_content: t.Optional[str] = None
+    ) -> str:
         """Create a dummy file in the test output directory.
 
         :param file_name: name of file to create, e.g. "file.txt"
@@ -649,8 +659,8 @@ def setup_test_colo(
         if db_type == "uds" and colo_model_name is not None:
             tmp_dir = tempfile.gettempdir()
             socket_suffix = str(uuid.uuid4())[:7]
-            db_args["unix_socket"] = os.path.join(tmp_dir,
-                f"{colo_model_name}_{socket_suffix}.socket")
+            socket_name = f"{colo_model_name}_{socket_suffix}.socket"
+            db_args["unix_socket"] = os.path.join(tmp_dir, socket_name)
 
         colocate_fun: t.Dict[str, t.Callable[..., None]] = {
             "tcp": colo_model.colocate_db_tcp,
@@ -659,16 +669,172 @@ def setup_test_colo(
         }
         with warnings.catch_warnings():
             if db_type == "deprecated":
-                warnings.filterwarnings(
-                    "ignore",
-                    message="`colocate_db` has been deprecated"
-                )
+                message = "`colocate_db` has been deprecated"
+                warnings.filterwarnings("ignore", message=message)
             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
 
+
 @pytest.fixture
-def config() -> smartsim._core.config.Config:
+def config() -> Config:
     return CONFIG
+
+
+class MockSink:
+    """Telemetry sink that writes console output for testing purposes"""
+
+    def __init__(self, delay_ms: int = 0) -> None:
+        self._delay_ms = delay_ms
+        self.num_saves = 0
+        self.args: t.Any = None
+
+    async def save(self, *args: t.Any) -> None:
+        """Save all arguments as console logged messages"""
+        self.num_saves += 1
+        if self._delay_ms:
+            # mimic slow collection....
+            delay_s = self._delay_ms / 1000
+            await asyncio.sleep(delay_s)
+        self.args = args
+
+
+@pytest.fixture
+def mock_sink() -> t.Type[MockSink]:
+    return MockSink
+
+
+@pytest.fixture
+def mock_con() -> t.Callable[[int, int], t.Iterable[t.Any]]:
+    """Generates mock db connection telemetry"""
+
+    def _mock_con(min: int = 1, max: int = 254) -> t.Iterable[t.Any]:
+        for i in range(min, max):
+            yield [
+                {"addr": f"127.0.0.{i}:1234", "id": f"ABC{i}"},
+                {"addr": f"127.0.0.{i}:2345", "id": f"XYZ{i}"},
+            ]
+
+    return _mock_con
+
+
+@pytest.fixture
+def mock_mem() -> t.Callable[[int, int], t.Iterable[t.Any]]:
+    """Generates mock db memory usage telemetry"""
+
+    def _mock_mem(min: int = 1, max: int = 1000) -> t.Iterable[t.Any]:
+        for i in range(min, max):
+            yield {
+                "total_system_memory": 1000 * i,
+                "used_memory": 1111 * i,
+                "used_memory_peak": 1234 * i,
+            }
+
+    return _mock_mem
+
+
+@pytest.fixture
+def mock_redis() -> t.Callable[..., t.Any]:
+    def _mock_redis(
+        conn_side_effect=None,
+        mem_stats=None,
+        client_stats=None,
+        coll_side_effect=None,
+    ):
+        """Generate a mock object for the redis.Redis contract"""
+
+        class MockConn:
+            def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
+                if conn_side_effect is not None:
+                    conn_side_effect()
+
+            async def info(self, *args: t.Any, **kwargs: t.Any) -> t.Dict[str, t.Any]:
+                if coll_side_effect:
+                    await coll_side_effect()
+
+                if mem_stats:
+                    return next(mem_stats)
+                return {
+                    "total_system_memory": "111",
+                    "used_memory": "222",
+                    "used_memory_peak": "333",
+                }
+
+            async def client_list(
+                self, *args: t.Any, **kwargs: t.Any
+            ) -> t.Dict[str, t.Any]:
+                if coll_side_effect:
+                    await coll_side_effect()
+
+                if client_stats:
+                    return next(client_stats)
+                return {"addr": "127.0.0.1", "id": "111"}
+
+            async def ping(self):
+                return True
+
+        return MockConn
+
+    return _mock_redis
+
+
+class MockCollectorEntityFunc(t.Protocol):
+    @staticmethod
+    def __call__(
+        host: str = "127.0.0.1",
+        port: int = 6379,
+        name: str = "",
+        type: str = "",
+        telemetry_on: bool = False,
+    ) -> "JobEntity": ...
+
+
+@pytest.fixture
+def mock_entity(test_dir: str) -> MockCollectorEntityFunc:
+    def _mock_entity(
+        host: str = "127.0.0.1",
+        port: int = 6379,
+        name: str = "",
+        type: str = "",
+        telemetry_on: bool = False,
+    ) -> "JobEntity":
+        test_path = pathlib.Path(test_dir)
+
+        entity = JobEntity()
+        entity.name = name if name else str(uuid.uuid4())
+        entity.status_dir = str(test_path / entity.name)
+        entity.type = type
+        entity.telemetry_on = True
+        entity.collectors = {
+            "client": "",
+            "client_count": "",
+            "memory": "",
+        }
+        entity.config = {
+            "host": host,
+            "port": str(port),
+        }
+        entity.telemetry_on = telemetry_on
+        return entity
+
+    return _mock_entity
+
+
+class CountingCallable:
+    def __init__(self) -> None:
+        self._num: int = 0
+        self._details: t.List[t.Tuple[t.Tuple[t.Any, ...], t.Dict[str, t.Any]]] = []
+
+    def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
+        self._num += 1
+        self._details.append((args, kwargs))
+
+    @property
+    def num_calls(self) -> int:
+        return self._num
+
+    @property
+    def details(self) -> t.List[t.Tuple[t.Tuple[t.Any, ...], t.Dict[str, t.Any]]]:
+        return self._details

doc/changelog.rst

@@ -18,6 +18,7 @@ To be released at some future point in time
 
 Description
 
+- Update telemetry monitor, add telemetry collectors
 - Add method to specify node features for a Slurm job
 - Colo Orchestrator setup now blocks application start until setup finished
 - ExecArgs handling correction
@@ -32,6 +33,11 @@ Description
 
 Detailed Notes
 
+- Update the telemetry monitor to enable retrieval of metrics on a scheduled
+  interval. Switch basic experiment tracking telemetry to default to on. Add
+  database metric collectors. Improve telemetry monitor logging. Create
+  telemetry subpackage at `smartsim._core.utils.telemetry`. Refactor
+  telemetry monitor entrypoint. (SmartSim-PR460_)
 - Users can now specify node features for a Slurm job through
   ``SrunSettings.set_node_feature``. The method accepts a string
   or list of strings. (SmartSim-PR529_)
@@ -64,7 +70,8 @@ Detailed Notes
 - Remove previously deprecated behavior present in test suite on machines with
   Slurm and Open MPI. (SmartSim-PR520_)
 
-
+.. _SmartSim-PR460: https://github.com/CrayLabs/SmartSim/pull/460
+.. _SmartSim-PR512: https://github.com/CrayLabs/SmartSim/pull/512
 .. _SmartSim-PR529: https://github.com/CrayLabs/SmartSim/pull/529
 .. _SmartSim-PR522: https://github.com/CrayLabs/SmartSim/pull/522
 .. _SmartSim-PR524: https://github.com/CrayLabs/SmartSim/pull/524

setup.py

@@ -64,21 +64,22 @@
 #
 # This future is needed to print Python2 EOL message
 from __future__ import print_function
+
 import sys
 
 if sys.version_info < (3,):
     print("Python 2 has reached end-of-life and is not supported by SmartSim")
     sys.exit(-1)
 
 
-import os
 import importlib.util
+import os
 from pathlib import Path
 
 from setuptools import setup
-from setuptools.dist import Distribution
-from setuptools.command.install import install
 from setuptools.command.build_py import build_py
+from setuptools.command.install import install
+from setuptools.dist import Distribution
 
 # Some necessary evils we have to do to be able to use
 # the _install tools in smartsim/smartsim/_core/_install
@@ -181,6 +182,7 @@ def has_ext_modules(_placeholder):
         "pytest>=6.0.0",
         "pytest-cov>=2.10.1",
         "click==8.0.2",
+        "pytest-asyncio>=0.23.3",
     ],
     "mypy": [
         "mypy>=1.3.0",

smartsim/_core/_cli/validate.py

@@ -151,7 +151,8 @@ def test_install(
     with_onnx: bool,
 ) -> None:
     exp = Experiment("ValidationExperiment", exp_path=location, launcher="local")
-    exp.disable_telemetry()
+    exp.telemetry.disable()
+
     port = _find_free_port() if port is None else port
     with _make_managed_local_orc(exp, port) as client:
         logger.info("Verifying Tensor Transfer")