Improve diffexp for tiledb (#1388)

* Improve diffexp for tiledb

- The rows from the A and B sets are gathered and processed at the same time.  In this
  way the matrix is only accessed once instead of twice for each tile.
- There is now a single thread queue that gets shared between all callers of the diffexp.
  This will slow down work if diffexp gets too busy.
- There is a target_workunit amount of work given to each thread.  Previously the
  workunit was (rows selected * width of tile), which could be small.  Now multiple
  column tiles can be combined into one workunit.  If the target is too small then
  thread and other overheads may reduce performance.  If target_workunit is too large
  then the size of the gathered sub matrix may take up too much memory.
- add configuration parameters (max_workers, cpu_multiplier, and  target_workunit)

server/common/app_config.py

@@ -1,6 +1,6 @@
 from server import __version__ as cellxgene_version
 from flatten_dict import flatten
-from os import mkdir, environ
+import os
 from os.path import splitext, basename, isdir
 import sys
 from urllib.parse import urlparse
@@ -16,8 +16,9 @@ from server.common.utils import find_available_port, is_port_available
 import warnings
 from server.common.annotations import AnnotationsLocalFile
 from server.common.utils import custom_format_warning
+import server.compute.diffexp_cxg as diffexp_tiledb
 
-DEFAULT_SERVER_PORT = int(environ.get("CXG_SERVER_PORT", "5005"))
+DEFAULT_SERVER_PORT = int(os.environ.get("CXG_SERVER_PORT", "5005"))
 # anything bigger than this will generate a special message
 BIG_FILE_SIZE_THRESHOLD = 100 * 2 ** 20  # 100MB
 
@@ -85,6 +86,9 @@ class AppConfig(object):
             self.diffexp__enable = dc["diffexp"]["enable"]
             self.diffexp__lfc_cutoff = dc["diffexp"]["lfc_cutoff"]
             self.diffexp__top_n = dc["diffexp"]["top_n"]
+            self.diffexp__alg_cxg__max_workers = dc["diffexp"]["alg_cxg"]["max_workers"]
+            self.diffexp__alg_cxg__cpu_multiplier = dc["diffexp"]["alg_cxg"]["cpu_multiplier"]
+            self.diffexp__alg_cxg__target_workunit = dc["diffexp"]["alg_cxg"]["target_workunit"]
 
             self.data_locator__s3__region_name = dc["data_locator"]["s3"]["region_name"]
 
@@ -256,7 +260,7 @@ class AppConfig(object):
         # secret key:
         #   first, from CXG_SECRET_KEY environment variable
         #   second, from config file
-        self.server__flask_secret_key = environ.get("CXG_SECRET_KEY", self.server__flask_secret_key)
+        self.server__flask_secret_key = os.environ.get("CXG_SECRET_KEY", self.server__flask_secret_key)
 
     def handle_data_locator(self, context):
         self.__check_attr("data_locator__s3__region_name", (type(None), bool, str))
@@ -381,7 +385,7 @@ class AppConfig(object):
 
             if dirname is not None and not isdir(dirname):
                 try:
-                    mkdir(dirname)
+                    os.mkdir(dirname)
                 except OSError:
                     raise ConfigurationError("Unable to create directory specified by --annotations-dir")
 
@@ -433,6 +437,9 @@ class AppConfig(object):
         self.__check_attr("diffexp__enable", bool)
         self.__check_attr("diffexp__lfc_cutoff", float)
         self.__check_attr("diffexp__top_n", int)
+        self.__check_attr("diffexp__alg_cxg__max_workers", (str, int))
+        self.__check_attr("diffexp__alg_cxg__cpu_multiplier", int)
+        self.__check_attr("diffexp__alg_cxg__target_workunit", int)
 
         if self.single_dataset__datapath:
             with self.matrix_data_cache_manager.data_adaptor(self.single_dataset__datapath, self) as data_adaptor:
@@ -442,6 +449,14 @@ class AppConfig(object):
                         f"running differential expression may take longer or fail."
                     )
 
+        max_workers = self.diffexp__alg_cxg__max_workers
+        cpu_multiplier = self.diffexp__alg_cxg__cpu_multiplier
+        cpu_count = os.cpu_count()
+        max_workers = min(max_workers, cpu_multiplier * cpu_count)
+        diffexp_tiledb.set_config(
+            max_workers,
+            self.diffexp__alg_cxg__target_workunit)
+
     def handle_adaptor(self, context):
         # cxg
         self.__check_attr("adaptor__cxg_adaptor__tiledb_ctx", dict)

server/common/default_config.py

@@ -66,6 +66,15 @@ diffexp:
   enable: true
   lfc_cutoff: 0.01
   top_n: 10
+  alg_cxg:
+    # The number of threads to use is computed from: min(max_workers, cpu_multipler * cpu_count).
+    # Where cpu_count is determined at runtime.
+    max_workers: 64
+    cpu_multiplier: 4
+
+    # The target number of matrix elements that are evaluated
+    # together in one thread.
+    target_workunit: 16_000_000
 
 data_locator:
   s3:

server/common/errors.py

@@ -76,3 +76,11 @@ class ExceedsLimitError(Exception):
     """
 
     pass
+
+
+class ComputeError(Exception):
+    """
+    Raised when an error occurs during a compute algorithm (such as diffexp)
+    """
+
+    pass

server/compute/diffexp_cxg.py

@@ -0,0 +1,110 @@
+import concurrent.futures
+import numpy as np
+from server.compute.diffexp_generic import diffexp_ttest_from_mean_var, mean_var_n
+from server.data_cxg.cxg_util import pack_selector_from_indices
+from server.common.errors import ComputeError
+
+"""
+See the comments in diffexp_generic for a description of this algorithm
+
+This implementation runs directly in-process.  It is multi- threaded, but not particularly scalable.
+Longer term, will likely move to a distributed framework for this.
+
+There are currently no global throttles on simultaneous workers.
+"""
+
+diffexp_thread_executor = None
+max_workers = None
+target_workunit = None
+
+
+def set_config(config_max_workers, config_target_workunit):
+    global max_workers
+    global target_workunit
+    max_workers = config_max_workers
+    target_workunit = config_target_workunit
+
+
+def get_thread_executor():
+    global diffexp_thread_executor
+    if diffexp_thread_executor is None:
+        diffexp_thread_executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
+    return diffexp_thread_executor
+
+
+def diffexp_ttest(adaptor, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01):
+    row_selector_A = np.where(maskA)[0]
+    row_selector_B = np.where(maskB)[0]
+    nA = len(row_selector_A)
+    nB = len(row_selector_B)
+    matrix = adaptor.open_array("X")
+
+    dtype = matrix.dtype
+    cols = matrix.shape[1]
+    tile_extent = [dim.tile for dim in matrix.schema.domain]
+
+    # The rows from both row_selector_A and row_selector_B are gathered at the
+    # same time, then the mean and variance are computed by subsetting on that
+    # combined submatrix.  Combining the gather reduces number of requests/bandwidth
+    # to the data source.
+    row_selector_AB = np.union1d(row_selector_A, row_selector_B)
+    row_selector_A_in_AB = np.in1d(row_selector_AB, row_selector_A, assume_unique=True)
+    row_selector_B_in_AB = np.in1d(row_selector_AB, row_selector_B, assume_unique=True)
+    row_selector_AB = pack_selector_from_indices(row_selector_AB)
+
+    # because all IO is done per-tile, and we are always dense and col-major,
+    # use the tile column size as the unit of partition.  Possibly access
+    # more than one column tile at a time based on the target_workunit.
+    # Revisit partitioning if we change the X layout, or start using a non-local execution environment
+    # which may have other constraints.
+
+    # TODO: If the number of row selections is large enough, then the cells_per_coltile will exceed
+    # the target_workunit.  A potential improvement would be to partition by both columns and rows.
+    # However partitioning the rows is slightly more complex due to the arbitrary distribution
+    # of row selections that are passed into this algorithm.
+    cells_per_coltile = (nA + nB) * tile_extent[1]
+    cols_per_partition = max(1, int(target_workunit / cells_per_coltile)) * tile_extent[1]
+    col_partitions = [(c, min(c + cols_per_partition, cols)) for c in range(0, cols, cols_per_partition)]
+
+    meanA = np.zeros((cols,), dtype=np.float64)
+    varA = np.zeros((cols,), dtype=np.float64)
+    meanB = np.zeros((cols,), dtype=np.float64)
+    varB = np.zeros((cols,), dtype=np.float64)
+
+    executor = get_thread_executor()
+    futures = []
+    for cols in col_partitions:
+        futures.append(
+            executor.submit(_mean_var_ab, matrix, row_selector_AB, row_selector_A_in_AB, row_selector_B_in_AB, cols)
+        )
+
+    for future in futures:
+        # returns tuple: (meanA, varA, meanB, varB, cols)
+        try:
+            result = future.result()
+            part_meanA, part_varA, part_meanB, part_varB, cols = result
+            meanA[cols[0]: cols[1]] += part_meanA
+            varA[cols[0]: cols[1]] += part_varA
+            meanB[cols[0]: cols[1]] += part_meanB
+            varB[cols[0]: cols[1]] += part_varB
+        except Exception as e:
+            for future in futures:
+                future.cancel()
+            raise ComputeError(str(e))
+
+    r = diffexp_ttest_from_mean_var(
+        meanA.astype(dtype),
+        varA.astype(dtype),
+        nA,
+        meanB.astype(dtype),
+        varB.astype(dtype),
+        nB, top_n, diffexp_lfc_cutoff)
+
+    return r
+
+
+def _mean_var_ab(matrix, row_selector_AB, row_selector_A_in_AB, row_selector_B_in_AB, col_range):
+    X = matrix.multi_index[row_selector_AB, col_range[0] : col_range[1] - 1][""]
+    meanA, varA, n = mean_var_n(X[row_selector_A_in_AB])
+    meanB, varB, n = mean_var_n(X[row_selector_B_in_AB])
+    return (meanA, varA, meanB, varB, col_range)

server/compute/diffexp_tiledb.py

@@ -1,117 +0,0 @@
-import os
-import concurrent.futures
-from itertools import repeat
-import numpy as np
-from server.compute.diffexp_generic import diffexp_ttest_from_mean_var, mean_var_n
-from server.data_cxg.cxg_util import pack_selector_from_indices
-
-"""
-See the comments in diffexp_generic for a description of this algorithm
-
-This implementation runs directly in-process.  It is multi- threaded, but not particularly scalable.
-Longer term, will likely move to a distributed framework for this.
-
-There are currently no global throttles on simultaneous workers.
-"""
-
-# number of simultaneous workers, per HTTP request
-MAX_WORKERS = 16
-
-
-def diffexp_ttest(adaptor, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01):
-    row_selector_A = np.where(maskA)[0]
-    row_selector_B = np.where(maskB)[0]
-    nA = len(row_selector_A)
-    nB = len(row_selector_B)
-    matrix = adaptor.open_array("X")
-    # mean, variance, N - calculate for both selections
-    with MyThreadPoolExecutor(max_workers=2) as executor:
-        A = executor.submit(mean_var, matrix, row_selector_A)
-        B = executor.submit(mean_var, matrix, row_selector_B)
-        meanA, varA = A.result()
-        meanB, varB = B.result()
-
-    return diffexp_ttest_from_mean_var(meanA, varA, nA, meanB, varB, nB, top_n, diffexp_lfc_cutoff)
-
-
-class DispatchMixins:
-    def parallel_dispatch(self, fn, *iterables):
-        """
-        Dispatch jobs via concurrent.futures.Executor.submit() and wait for their
-        completion, return result via future.result().  Primary purpose is to throttle
-        dispatch rate so that only 'MAX_JOBS_QUEUE_LENGTH' jobs are running at any
-        given time, reducing overall memory footprint.
-        """
-        MAX_JOBS_QUEUE_LENGTH = self._max_workers + 4
-
-        def submit_more_jobs(jobs, active_jobs):
-            for job in jobs:
-                future = self.submit(fn, *job)
-                active_jobs[future] = job
-                if len(active_jobs) >= MAX_JOBS_QUEUE_LENGTH:
-                    return True
-            return False
-
-        def result_iterator(jobs):
-            active_jobs = {}  # map of future -> args
-            try:
-                while submit_more_jobs(jobs, active_jobs) or len(active_jobs) > 0:
-                    for future in concurrent.futures.as_completed(active_jobs.keys()):
-                        job = active_jobs[future]
-                        result = future.result()
-                        # be careful to not retain dangling references
-                        del active_jobs[future], future
-                        yield (result, job)
-            except Exception as e:
-                print(str(e))
-                raise
-            finally:
-                for future in active_jobs.keys():
-                    future.cancel()
-
-        return result_iterator(zip(*iterables))
-
-
-class MyThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor, DispatchMixins):
-    pass
-
-
-def _mean_var(matrix, row_selector, col_range):
-    X = matrix.multi_index[row_selector, col_range[0] : col_range[1] - 1][""]
-    mean, var, n = mean_var_n(X)
-    return (mean, var)
-
-
-def mean_var(matrix, row_selector):
-    """
-    row_selector: list of row indices
-    """
-    dtype = matrix.dtype
-    rows, cols = matrix.shape
-    tile_extent = [dim.tile for dim in matrix.schema.domain]
-
-    dispatch_func = _mean_var
-    row_selector = pack_selector_from_indices(row_selector)
-
-    # because all IO is done per-tile, and we are always dense and col-major,
-    # use the tile column size as the partition.  Revisit partitioning if we
-    # change the X layout, or start using a non-local execution environment
-    # which may have other constraints.
-    cols_per_partition = tile_extent[1]
-    col_partitions = [(c, min(c + cols_per_partition, cols)) for c in range(0, cols, cols_per_partition)]
-
-    max_workers = max(1, min(MAX_WORKERS, os.cpu_count()))  # throttle max_workers
-
-    mean = np.zeros((cols,), dtype=np.float64)
-    var = np.zeros((cols,), dtype=np.float64)
-    dispatch_args = [repeat(matrix), repeat(row_selector), col_partitions]
-    with MyThreadPoolExecutor(max_workers=max_workers) as exec:
-        for result in exec.parallel_dispatch(dispatch_func, *dispatch_args):
-            # returns tuple: (return_val, dispatch_args)
-            m, v = result[0]
-            cols = result[1][2]
-            mean[cols[0] : cols[1]] += m
-            var[cols[0] : cols[1]] += v
-            del result, m, v
-
-    return (mean.astype(dtype), var.astype(dtype))

server/data_cxg/cxg_adaptor.py

@@ -8,7 +8,7 @@ from server.common.constants import Axis
 from server.data_common.data_adaptor import DataAdaptor
 from server.data_common.fbs.matrix import encode_matrix_fbs
 from server.data_cxg.cxg_util import pack_selector_from_mask
-import server.compute.diffexp_tiledb as diffexp_tiledb
+import server.compute.diffexp_cxg as diffexp_cxg
 from server.common.immutable_kvcache import ImmutableKVCache
 import tiledb
 import numpy as np
@@ -192,7 +192,7 @@ class CxgAdaptor(DataAdaptor):
             top_n = self.config.diffexp__top_n
         if lfc_cutoff is None:
             lfc_cutoff = self.config.diffexp__lfc_cutoff
-        return diffexp_tiledb.diffexp_ttest(self, maskA, maskB, top_n, lfc_cutoff)
+        return diffexp_cxg.diffexp_ttest(self, maskA, maskB, top_n, lfc_cutoff)
 
     def get_X_array(self, obs_mask=None, var_mask=None):
         obs_items = pack_selector_from_mask(obs_mask)

server/test/run_diffexp.py

@@ -4,7 +4,7 @@ import random
 import time
 import numpy as np
 
-import server.compute.diffexp_tiledb as diffexp_tiledb
+import server.compute.diffexp_cxg as diffexp_cxg
 import server.compute.diffexp_generic as diffexp_generic
 
 from server.common.app_config import AppConfig
@@ -19,7 +19,7 @@ def main():
     parser.add_argument("-nb", "--numB", type=int, required=True, help="number of rows in group B")
     parser.add_argument("-t", "--trials", default=1, type=int, help="number of trials")
     parser.add_argument(
-        "-a", "--alg", choices=("default", "generic", "tiledb"), default="default", help="algorithm to use"
+        "-a", "--alg", choices=("default", "generic", "cxg"), default="default", help="algorithm to use"
     )
     parser.add_argument("-s", "--show", default=False, action="store_true", help="show the results")
     parser.add_argument(
@@ -30,8 +30,6 @@ def main():
     args = parser.parse_args()
 
     app_config = AppConfig()
-    app_config.data_locator__s3__region_name = "us-west-2"
-    app_config.adaptor__cxg_adaptor__tiledb_ctx["vfs.s3.region"] = "us-west-2"
     app_config.single_dataset__datapath = args.dataset
     app_config.server__verbose = True
     app_config.complete_config()
@@ -70,11 +68,11 @@ def main():
             results = adaptor.compute_diffexp_ttest(maskA, maskB)
         elif args.alg == "generic":
             results = diffexp_generic.diffexp_ttest(adaptor, maskA, maskB)
-        elif args.alg == "tiledb":
+        elif args.alg == "cxg":
             if not isinstance(adaptor, CxgAdaptor):
-                print("tiledb only works with CxgAdaptor")
+                print("cxg only works with CxgAdaptor")
                 sys.exit(1)
-            results = diffexp_tiledb.diffexp_ttest(adaptor, maskA, maskB)
+            results = diffexp_cxg.diffexp_ttest(adaptor, maskA, maskB)
 
         t2 = time.time()
         print("TIME=", t2 - t1)

server/test/test_diffexp.py

@@ -0,0 +1,80 @@
+import unittest
+from server.data_common.matrix_loader import MatrixDataLoader
+from server.common.app_config import AppConfig
+import server.compute.diffexp_cxg as diffexp_cxg
+import server.compute.diffexp_generic as diffexp_generic
+import numpy as np
+
+
+class DiffExpTest(unittest.TestCase):
+    """Tests the diffexp returns the expected results for one test case, using different
+    adaptor types and different algorithms."""
+
+    def load_dataset(self, path):
+        app_config = AppConfig()
+        app_config.single_dataset__datapath = path
+        app_config.server__verbose = True
+        app_config.complete_config()
+        loader = MatrixDataLoader(path)
+        adaptor = loader.open(app_config)
+        return adaptor
+
+    def get_mask(self, adaptor, start, stride):
+        """Simple function to return a mask or rows"""
+        rows = adaptor.get_shape()[0]
+        sel = list(range(start, rows, stride))
+        mask = np.zeros(rows, dtype=bool)
+        mask[sel] = True
+        return mask
+
+    def check_1_10_2_10(self, results):
+        """Checks the results for a specific set of rows selections"""
+        expects = [
+            [956, 0.016060986, 0.0008649321884808977, 1.0],
+            [1124, 0.96602094, 0.0011717216548271284, 1.0],
+            [1809, 1.1110606, 0.0019304405196777848, 1.0],
+            [1712, -0.5525154, 0.0051788902660723345, 1.0],
+            [1754, 0.5201581, 0.005691734062127954, 1.0],
+            [948, 1.6390722, 0.006622111055981219, 1.0],
+            [1810, 0.78618884, 0.007055917428377063, 1.0],
+            [779, 1.5241305, 0.007202934422407284, 1.0],
+            [1575, 1.0317602, 0.007830310753043345, 1.0],
+            [576, 0.97873515, 0.008272092578813124, 1.0],
+        ]
+        self.assertEqual(len(results), len(expects))
+        for result, expect in zip(results, expects):
+            self.assertEqual(result[0], expect[0])
+            self.assertAlmostEqual(result[1], expect[1])
+            self.assertAlmostEqual(result[2], expect[2])
+            self.assertAlmostEqual(result[3], expect[3])
+
+    def test_anndata_default(self):
+        """Test an anndata adaptor with its default diffexp algorithm (diffexp_generic)"""
+        adaptor = self.load_dataset("../example-dataset/pbmc3k.h5ad")
+        maskA = self.get_mask(adaptor, 1, 10)
+        maskB = self.get_mask(adaptor, 2, 10)
+        results = adaptor.compute_diffexp_ttest(maskA, maskB, 10)
+        self.check_1_10_2_10(results)
+
+    def test_cxg_default(self):
+        """Test a cxg adaptor with its default diffexp algorithm (diffexp_cxg)"""
+        adaptor = self.load_dataset("test/test_datasets/pbmc3k.cxg")
+        maskA = self.get_mask(adaptor, 1, 10)
+        maskB = self.get_mask(adaptor, 2, 10)
+
+        # run it through the adaptor
+        results = adaptor.compute_diffexp_ttest(maskA, maskB, 10)
+        self.check_1_10_2_10(results)
+
+        # run it directly
+        results = diffexp_cxg.diffexp_ttest(adaptor, maskA, maskB, 10)
+        self.check_1_10_2_10(results)
+
+    def test_cxg_generic(self):
+        """Test a cxg adaptor with the generic adaptor"""
+        adaptor = self.load_dataset("test/test_datasets/pbmc3k.cxg")
+        maskA = self.get_mask(adaptor, 1, 10)
+        maskB = self.get_mask(adaptor, 2, 10)
+        # run it directly
+        results = diffexp_generic.diffexp_ttest(adaptor, maskA, maskB, 10)
+        self.check_1_10_2_10(results)