qsprpred.extra.gpu.utils package

Submodules

qsprpred.extra.gpu.utils.parallel module

class qsprpred.extra.gpu.utils.parallel.TorchJITGenerator(n_workers: int | None = None, worker_type: Literal['cpu', 'gpu'] = 'cpu', use_gpus: list[int] | None = None, jobs_per_gpu: int = 1)[source]

Bases: MultiprocessingJITGenerator

A variant of the MultiprocessingPoolGenerator that uses the torch.multiprocessing.Pool instead of the standard multiprocessing.Pool. This is needed when the parallel processing is done with PyTorch tensors or models, which require the torch.multiprocessing and using the spawn start method.

Configures the multiprocessing pool generator.

Parameters:

  • n_workers (int) – Number of workers to use.

  • worker_type (Literal["cpu", "gpu"]) – The type of worker to use.

  • use_gpus (list[int] | None) – A list of GPU indices to use. Only applicable if worker_type is ‘gpu’. If None, all available GPUs will be used.

  • jobs_per_gpu (int) – Number of jobs to run on each GPU.

checkProcess(process)

A simple check of a process or future before a result is attempted to be retrieved.

Parameters:

process (Any) – The process object or a future to check.

Returns:

None if the process is OK, otherwise raises an exception.

Raises:

Exception – If the process has a problem.

checkResultAvailable(process)

Check if the result of a process is available.

Parameters:

process (Any) – The process object or a future to check for a result.

Returns:

True if the result is available, otherwise False.

createJob(pool, process_func, *args, **kwargs)

Submit a job to the pool that applies the function to a generator item.

Parameters:

  • pool (Any) – The pool object to submit the job to.

  • process_func (Callable) – The function to apply to the input arguments.

  • args (tuple) – Positional arguments to pass to the function. The first argument should be the item from the generator.

  • kwargs (dict) – Additional keyword arguments to pass to the function.

Returns:

The process object or future that was submitted to the pool.

getPool()[source]

Create the pool of workers consuming the generator.

Returns:

A pool object that can be used to apply the function to the generator items in parallel.

getResult(process)

Get the result of a process.

Parameters:

process (Any) – The process object or a future to get the result from.

Returns:

The result of the process.

handleException(process, exception)

Handle an exception raised by a process. This is executed when the process raises an unexpected exception or the checkProcess method raises an exception.

Parameters:

  • process (Any) – The process object or a future that raised the exception.

  • exception (Exception) – The exception raised by the process.

Returns:

The result to yield instead of the result of the process.

make(generator: Generator, process_func: Callable, *args, **kwargs) → Generator

A parallel “JIT (Just In Time)” generator that yields the results of a function applied in parallel to each item of a supplied generator. The advantage of this JIT implementation is that it only evaluates the next item in the input generator as soon as it is needed for the calculation. This means that only one item from the iterable is loaded into memory at a time.

The generator also supports timeout for each job. The “pebble” pool_type can be used to support this feature. In all other cases, the “multiprocessing” pool_type is sufficient.

Parameters:

  • generator (SupportsNext) – An iterable object to apply the function to.

  • process_func (Callable) – The function to apply to each item in the iterable.

  • args (tuple | None) – Additional positional arguments to pass to the function.

  • kwargs (dict | None) – Additional keyword arguments to pass to the function.

Returns:

A generator that yields the results of the function applied to each item in the iterable. If a timeout is specified, a TimeoutError will be returned instead.

qsprpred.extra.gpu.utils.tests module

Module contents