Remove JIT dep in memory leak check and reorganize tests (#830)

  *  Runs DataFrames tests that do not have jit dependencies in separate sessions on CI to avoid importing JIT before a non-jit test
  *  Removes dependency on numba/compiler in memory_leak_check
  *  Once scatter/gather don't use JIT, we can remove a majority of the tests from the "JIT dependency"

Author: scott-routledge2

.github/workflows/_test_python_source.yml

@@ -23,7 +23,7 @@ on:
         required: false
         default: 'ubuntu-latest'
         type: string
-      # Options: DF_LIB, SPAWN, NORMAL
+      # Options: DF_LIB, DF_LIB_NO_JIT, SPAWN, NORMAL
       test-type:
         type: string
         description: The kind of tests to run e.g. spawn tests
@@ -130,7 +130,7 @@ jobs:
             # Disabling the DataFrame library for spawn tests since some of the tests
             # create Pandas manager states for testing that are not fully functional.
             BODO_ENABLE_DATAFRAME_LIBRARY: ${{ inputs.test-type != 'SPAWN' && '1' || '0' }}
-            BODO_ENABLE_TEST_DATAFRAME_LIBRARY: ${{ inputs.test-type == 'DF_LIB' && '1' || '0' }}
+            BODO_ENABLE_TEST_DATAFRAME_LIBRARY: ${{ startsWith(inputs.test-type, 'DF_LIB') && '1' || '0' }}
             BODOSQL_PY4J_GATEWAY_PORT: "auto"
             BODO_SPAWN_MODE: "0"
             BODO_BUFFER_POOL_REMOTE_MODE: "1"

.github/workflows/pr_ci.yml

@@ -137,6 +137,19 @@ jobs:
       test-type: SPAWN
     secrets: inherit
 
+  df-lib-non-jit-ci:
+    needs: [compile-bodo]
+    name: Test DF Library Non-JIT
+    uses: ./.github/workflows/_test_python_source.yml
+    with:
+      batch: 1
+      total-batches: 1
+      pytest-marker: "df_lib and not jit_dependency and not weekly and not iceberg"
+      collect-coverage: true
+      os: ${{ inputs.runner_os || 'ubuntu-latest' }}
+      test-type: "DF_LIB_NO_JIT"
+    secrets: inherit
+
   df-lib-ci:
     needs: [compile-bodo]
     name: Test DF Library
@@ -149,7 +162,7 @@ jobs:
     with:
       batch: ${{ matrix.batch }}
       total-batches: 2
-      pytest-marker: "df_lib and not weekly and not iceberg"
+      pytest-marker: "df_lib and jit_dependency and not weekly and not iceberg"
       collect-coverage: true
       os: ${{ inputs.runner_os || 'ubuntu-latest' }}
       test-type: "DF_LIB"
@@ -241,6 +254,7 @@ jobs:
     needs:
       - collect-results
       - df-lib-ci
+      - df-lib-non-jit-ci
       - validate
     runs-on: ubuntu-latest
     steps:

bodo/libs/_array.cpp

@@ -955,9 +955,38 @@ table_info* retrieve_table_py_entry(table_info* in_table,
     }
 }
 
+// Macro to define a Python wrapper for a memory stats functions
+// (Avoids numba JIT dependency in testing)
+#define DEFINE_NOARG_STATS_WRAPPER(py_name, func_name)                \
+    static PyObject* py_name(PyObject* self, PyObject* args) {        \
+        if (PyTuple_Size(args) != 0) {                                \
+            PyErr_SetString(PyExc_TypeError,                          \
+                            #func_name "() does not take arguments"); \
+            return nullptr;                                           \
+        }                                                             \
+        return PyLong_FromSize_t(func_name());                        \
+    }
+
+DEFINE_NOARG_STATS_WRAPPER(get_stats_alloc_py_wrapper, get_stats_alloc)
+DEFINE_NOARG_STATS_WRAPPER(get_stats_free_py_wrapper, get_stats_free)
+DEFINE_NOARG_STATS_WRAPPER(get_stats_mi_alloc_py_wrapper, get_stats_mi_alloc)
+DEFINE_NOARG_STATS_WRAPPER(get_stats_mi_free_py_wrapper, get_stats_mi_free)
+
+#undef DEFINE_NOARG_STATS_WRAPPER
+
+static PyMethodDef array_ext_methods[] = {
+#define declmethod(func) {#func, (PyCFunction)func, METH_VARARGS, NULL}
+    declmethod(get_stats_alloc_py_wrapper),
+    declmethod(get_stats_free_py_wrapper),
+    declmethod(get_stats_mi_alloc_py_wrapper),
+    declmethod(get_stats_mi_free_py_wrapper),
+    {nullptr},
+#undef declmethod
+};
+
 PyMODINIT_FUNC PyInit_array_ext(void) {
     PyObject* m;
-    MOD_DEF(m, "array_ext", "No docs", nullptr);
+    MOD_DEF(m, "array_ext", "No docs", array_ext_methods);
     if (m == nullptr) {
         return nullptr;
     }
@@ -1040,10 +1069,6 @@ PyMODINIT_FUNC PyInit_array_ext(void) {
     SetAttrStringFromVoidPtr(m, pd_pyarrow_array_from_bodo_array_py_entry);
     SetAttrStringFromVoidPtr(m, string_array_from_sequence);
     SetAttrStringFromVoidPtr(m, is_na_value);
-    SetAttrStringFromVoidPtr(m, get_stats_alloc);
-    SetAttrStringFromVoidPtr(m, get_stats_free);
-    SetAttrStringFromVoidPtr(m, get_stats_mi_alloc);
-    SetAttrStringFromVoidPtr(m, get_stats_mi_free);
     SetAttrStringFromVoidPtr(m, array_info_getitem);
     SetAttrStringFromVoidPtr(m, array_info_getdata1);
     SetAttrStringFromVoidPtr(m, array_info_unpin);

bodo/pytest.ini

@@ -27,6 +27,8 @@ markers =
     tz_aware: timezone-aware datetime tests
     flaky: flaky tests
     df_lib: DataFrame library tests
+    jit_dependency: marks Bodo DataFrames tests that require jit dependencies
+
 
     # Used for Azure NP=3
     bodosql_1of7: part 1/7 of the test suite

bodo/tests/conftest.py

@@ -8,6 +8,7 @@
 import os
 import shutil
 import subprocess
+import sys
 import time
 import traceback
 from collections.abc import Callable, Generator
@@ -94,7 +95,6 @@ def memory_leak_check():
     Equivalent to Numba's MemoryLeakMixin:
     https://github.com/numba/numba/blob/13ece9b97e6f01f750e870347f231282325f60c3/numba/tests/support.py#L688
     """
-    import bodo.decorators  # isort:skip # noqa
     import bodo.tests.utils
     import bodo.utils.allocation_tracking
 
@@ -120,6 +120,16 @@ def memory_leak_check():
         assert total_mi_alloc == total_mi_free
 
 
+@pytest.fixture(scope="function")
+def jit_import_check():
+    """Fixture used to assert that a test should not import JIT."""
+    jit_already_imported = "bodo.decorators" in sys.modules
+    yield
+    assert jit_already_imported or "bodo.decorators" not in sys.modules, (
+        "Test was not explicitly marked as a JIT dependency, but imported JIT."
+    )
+
+
 def item_file_name(item):
     """Get the name of a pytest item. Uses the default pytest implementation, except for C++ tests, where we return the cached name"""
     if isinstance(item, (CppTestFile, CppTestItem)):
@@ -259,6 +269,24 @@ def pytest_collection_modifyitems(items):
         for marker in markers:
             item.add_marker(marker)
 
+    # If running tests in DataFrames mode, add a check that JIT is not imported to all tests
+    # that don't explicitly depend on JIT via the jit_dependency marker.
+    if bodo.test_dataframe_library_enabled:
+        assert "bodo.decorators" not in sys.modules, (
+            "JIT was imported before pytest_collection_modifyitems finished."
+        )
+
+        no_jit_dep = [
+            item for item in items if not item.get_closest_marker("jit_dependency")
+        ]
+
+        for item in no_jit_dep:
+            if (
+                hasattr(item, "fixturenames")
+                and "jit_import_check" not in item.fixturenames
+            ):
+                item.fixturenames = ["jit_import_check"] + item.fixturenames
+
 
 def group_from_hash(testname, num_groups):
     """

bodo/tests/test_dataframe_part2.py

@@ -25,7 +25,7 @@
     pytest_pandas,
 )
 
-pytestmark = pytest_pandas
+pytestmark = pytest_pandas + [pytest.mark.jit_dependency]
 
 
 def test_pd_isna_getitem(memory_leak_check):

bodo/tests/test_dataframe_part3.py

@@ -31,7 +31,7 @@
     pytest_pandas,
 )
 
-pytestmark = pytest_pandas
+pytestmark = pytest_pandas + [pytest.mark.jit_dependency]
 
 
 def test_dataframe_apply_method_str(memory_leak_check):

bodo/tests/test_df_lib/test_ai_funcs.py

@@ -17,6 +17,8 @@
 from bodo.spawn.spawner import spawn_process_on_nodes
 from bodo.tests.utils import _test_equal
 
+pytestmark = pytest.mark.jit_dependency
+
 
 def test_write_s3_vectors(datapath):
     """Test writing to S3 Vectors using Bodo DataFrame API."""

bodo/tests/test_df_lib/test_direct_series_methods.py

@@ -7,6 +7,8 @@
 from bodo.pandas.plan import assert_executed_plan_count
 from bodo.tests.utils import _test_equal
 
+pytestmark = pytest.mark.jit_dependency
+
 
 @pytest.mark.parametrize("use_index1", [True, False])
 @pytest.mark.parametrize("use_index2", [True, False])

bodo/tests/test_df_lib/test_end_to_end.py

@@ -24,6 +24,8 @@
 )
 from bodo.tests.utils import _test_equal, pytest_mark_spawn_mode, temp_config_override
 
+pytestmark = pytest.mark.jit_dependency
+
 # Various Index kinds to use in test data (assuming maximum size of 100 in input)
 MAX_DATA_SIZE = 100