From aa33c4cabc9038cb98d83f59ad9f8d7e981f4a4f Mon Sep 17 00:00:00 2001
From: PhilipDeegan <philip.deegan@gmail.com>
Date: Thu, 8 Feb 2024 16:04:50 +0000
Subject: [PATCH] detect mpirun from cmake tests for scheduling (#8)

v0.0.9
---
 phlop/testing/parallel_processor.py |  1 +
 phlop/testing/test_cases.py         | 57 +++++++++++++++++++++++++----
 pyproject.toml                      |  2 +-
 setup.py                            |  2 +-
 4 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/phlop/testing/parallel_processor.py b/phlop/testing/parallel_processor.py
index e350766..e2c7413 100644
--- a/phlop/testing/parallel_processor.py
+++ b/phlop/testing/parallel_processor.py
@@ -60,6 +60,7 @@ def process(batches, n_cores=None, print_only=False, fail_fast=False):
         batches = [batches]
     if sum([len(t.tests) for t in batches]) == 0:
         return  # nothing to do
+
     if print_only:
         print_tests(batches)
         return
diff --git a/phlop/testing/test_cases.py b/phlop/testing/test_cases.py
index 24ce3e0..45e7e9c 100644
--- a/phlop/testing/test_cases.py
+++ b/phlop/testing/test_cases.py
@@ -18,6 +18,9 @@
 
 _LOG_DIR = Path(os.environ.get("PHLOP_LOG_DIR", os.getcwd()))
 
+CMD_PREFIX = ""
+CMD_POSTFIX = ""
+
 
 @dataclass
 class TestCase:
@@ -25,6 +28,7 @@ class TestCase:
     env: dict = field(default_factory=lambda: {})  # dict[str, str] # eventually
     working_dir: str = field(default_factory=lambda: None)
     log_file_path: str = field(default_factory=lambda: None)
+    cores: int = field(default_factory=lambda: 1)
 
     def __post_init__(self):
         self.cmd = self.cmd.strip()
@@ -67,8 +71,8 @@ def __call__(self, ctest_test):
 
 class PythonUnitTestCaseExtractor:
     def __call__(self, ctest_test):
-        if "python3" in ctest_test.cmd:  # hacky
-            return load_test_cases_from_cmake(ctest_test)
+        if "python3 " in ctest_test.cmd:  # hacky
+            return load_py_test_cases_from_cmake(ctest_test)
         return None
 
 
@@ -80,7 +84,7 @@ def __call__(self, ctest_test):
 
 
 def python3_default_test_cmd(clazz, test_id):
-    return f"python3 -m {clazz.__module__} {clazz.__name__}.{test_id}"
+    return f"python3 -Oum {clazz.__module__} {clazz.__name__}.{test_id}"
 
 
 def load_test_cases_in(
@@ -104,10 +108,13 @@ def load_test_cases_in(
     return tests
 
 
-def load_test_cases_from_cmake(ctest_test):
+def load_py_test_cases_from_cmake(ctest_test):
     ppath = ctest_test.env.get("PYTHONPATH", "")
+    bits = ctest_test.cmd.split(" ")
+    idx = [i for i, x in enumerate(bits) if "python3" in x][0]
+    prefix = " ".join(bits[:idx])
     with extend_sys_path([ctest_test.working_dir] + ppath.split(env_sep())):
-        pyfile = ctest_test.cmd.split(" ")[-1]
+        pyfile = bits[-1]
         return load_test_cases_in(
             classes_in_file(pyfile, unittest.TestCase, fail_on_import_error=True),
             env=ctest_test.env,
@@ -115,10 +122,34 @@ def load_test_cases_from_cmake(ctest_test):
             log_file_path=_LOG_DIR
             / ".phlop"
             / f"{Path(ctest_test.working_dir).relative_to(_LOG_DIR)}",
+            test_cmd_pre=CMD_PREFIX + prefix + CMD_POSTFIX,
         )
 
 
-# probably return a list of TestBatch if we do some core count detection per test
+def determine_cores_for_test_case(test_case):
+    cores = 1
+
+    try:
+        if "mpirun -n" in test_case.cmd:
+            bits = test_case.cmd.split(" ")
+            # print(bits)
+            idx = [i for i, x in enumerate(bits) if "mpirun" in x][0]
+            test_case.cores = int(bits[idx + 2])
+    except Exception as e:
+        print("EXXXX", e)
+
+    return test_case
+
+
+def binless(test_case):
+    if test_case.cmd.startswith("/usr/bin/"):
+        test_case.cmd = test_case.cmd[9:]
+    return test_case
+
+
+MUTATORS = [determine_cores_for_test_case, binless]
+
+
 def load_cmake_tests(cmake_dir, cores=1, test_cmd_pre="", test_cmd_post=""):
     cmake_tests = get_cmake_tests(cmake_dir)
     tests = []
@@ -131,12 +162,22 @@ def load_cmake_tests(cmake_dir, cores=1, test_cmd_pre="", test_cmd_post=""):
             )
         ]
 
+    test_batches = {}
+
+    def _add(test_cases):
+        for test_case in test_cases:
+            for mutator in MUTATORS:
+                test_case = mutator(test_case)
+            if test_case.cores not in test_batches:
+                test_batches[test_case.cores] = []
+            test_batches[test_case.cores].append(test_case)
+
     test_cases = []
     for test in tests:
         for extractor in EXTRACTORS:
             res = extractor(test)
             if res:
-                test_cases += res
+                _add(res)
                 break
 
-    return TestBatch(test_cases, cores)
+    return [TestBatch(v, k) for k, v in test_batches.items()]
diff --git a/pyproject.toml b/pyproject.toml
index 374f9fa..e0877c8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "phlop"
-version = "0.0.8"
+version = "0.0.9"
 
 dependencies = [
 
diff --git a/setup.py b/setup.py
index f36fdd4..6684083 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="phlop",
-    version="0.0.8",
+    version="0.0.9",
     cmdclass={},
     classifiers=[],
     include_package_data=True,