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)