From 5c0b8c629661803f85238288f87f538acfa1873a Mon Sep 17 00:00:00 2001 From: bmccandless Date: Mon, 13 Apr 2020 18:53:13 -0700 Subject: [PATCH] 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 | 23 ++++-- server/common/default_config.py | 9 +++ server/common/errors.py | 8 +++ server/compute/diffexp_cxg.py | 110 +++++++++++++++++++++++++++++ server/compute/diffexp_tiledb.py | 117 ------------------------------- server/data_cxg/cxg_adaptor.py | 4 +- server/test/run_diffexp.py | 12 ++-- server/test/test_diffexp.py | 80 +++++++++++++++++++++ 8 files changed, 233 insertions(+), 130 deletions(-) create mode 100644 server/compute/diffexp_cxg.py delete mode 100644 server/compute/diffexp_tiledb.py create mode 100644 server/test/test_diffexp.py diff --git a/server/common/app_config.py b/server/common/app_config.py index 10862e44..2090b889 100644 --- a/server/common/app_config.py +++ b/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) diff --git a/server/common/default_config.py b/server/common/default_config.py index 2a5f1c35..f378b7d5 100644 --- a/server/common/default_config.py +++ b/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: diff --git a/server/common/errors.py b/server/common/errors.py index cbc9e08c..53b2bab0 100644 --- a/server/common/errors.py +++ b/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 diff --git a/server/compute/diffexp_cxg.py b/server/compute/diffexp_cxg.py new file mode 100644 index 00000000..833946dd --- /dev/null +++ b/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) diff --git a/server/compute/diffexp_tiledb.py b/server/compute/diffexp_tiledb.py deleted file mode 100644 index 11517ee1..00000000 --- a/server/compute/diffexp_tiledb.py +++ /dev/null @@ -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)) diff --git a/server/data_cxg/cxg_adaptor.py b/server/data_cxg/cxg_adaptor.py index cd7f5bfc..857ace28 100644 --- a/server/data_cxg/cxg_adaptor.py +++ b/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) diff --git a/server/test/run_diffexp.py b/server/test/run_diffexp.py index 62200379..9faba5ac 100644 --- a/server/test/run_diffexp.py +++ b/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) diff --git a/server/test/test_diffexp.py b/server/test/test_diffexp.py new file mode 100644 index 00000000..11aa1d35 --- /dev/null +++ b/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)