mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-10-04 15:08:11 +08:00
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)
This commit is contained in:
bmccandless
@@ -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)
|
||||
@@ -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))
|
||||
Reference in New Issue
Block a user