From 5911121eee64d7d91f4e5b5c3efedb8fe557c59f Mon Sep 17 00:00:00 2001
From: Jordan Dialpuri <44945647+Dialpuri@users.noreply.github.com>
Date: Sun, 15 Mar 2026 15:19:14 +0000
Subject: [PATCH] Add memory-based thread calculation and `--max-memory`
argument
- Introduced `calculate_n_from_max_memory` for determining thread count based on maximum memory.
- Added `--max-memory` argument for thread calculation when `nthreads` is unspecified.
- Updated documentation to include the new `--max-memory` option.
- Default `nthreads` updated to 0, enabling smart thread computation logic.
---
.../topics/Command-Line-Usage.topic | 5 ++
.../src/nucleofind/prediction/arguments.py | 8 +++-
package/src/nucleofind/prediction/load.py | 4 +-
package/src/nucleofind/prediction/memory.py | 47 +++++++++++++++++++
package/src/nucleofind/prediction/predict.py | 12 +++++
5 files changed, 74 insertions(+), 2 deletions(-)
create mode 100644 package/src/nucleofind/prediction/memory.py
diff --git a/docs/Writerside/topics/Command-Line-Usage.topic b/docs/Writerside/topics/Command-Line-Usage.topic
index 6b4d2ce..65b4436 100644
--- a/docs/Writerside/topics/Command-Line-Usage.topic
+++ b/docs/Writerside/topics/Command-Line-Usage.topic
@@ -91,6 +91,11 @@
Use GPU acceleration. (Linux only)
Advanced
+
+ Set the number of threads based on the maximum amount of memory you would like NucleoFind to use. This option is
+ only active if `nthreads` is not specified.
+ This feature does not explicitly limit memory usage. Use with caution.
+
diff --git a/package/src/nucleofind/prediction/arguments.py b/package/src/nucleofind/prediction/arguments.py
index 9333e89..bcdf5b4 100644
--- a/package/src/nucleofind/prediction/arguments.py
+++ b/package/src/nucleofind/prediction/arguments.py
@@ -28,7 +28,7 @@ def parse_arguments() -> SimpleNamespace:
"-n",
"--nthreads",
nargs="?",
- default=1,
+ default=0,
type=int,
help="Number of threads to use",
)
@@ -62,6 +62,12 @@ def parse_arguments() -> SimpleNamespace:
parser.add_argument(
"--gpu", action=argparse.BooleanOptionalAction, help="Use GPU (experimental)"
)
+ parser.add_argument(
+ "--max-memory",
+ nargs="?",
+ help="Maximum memory to use in Gigabytes (experimental). This option is only active when nthreads is not specified.",
+ type=float,
+ )
parser.add_argument(
"--debug", action=argparse.BooleanOptionalAction, help="Turn on debug logging"
)
diff --git a/package/src/nucleofind/prediction/load.py b/package/src/nucleofind/prediction/load.py
index 75de1d4..ebb83dd 100644
--- a/package/src/nucleofind/prediction/load.py
+++ b/package/src/nucleofind/prediction/load.py
@@ -63,7 +63,9 @@ def load_onnx_model(
if use_gpu:
providers.insert(0, "CUDAExecutionProvider")
sess_options = rt.SessionOptions()
- sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+ sess_options.graph_optimization_level = (
+ rt.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+ )
try:
return rt.InferenceSession(
diff --git a/package/src/nucleofind/prediction/memory.py b/package/src/nucleofind/prediction/memory.py
new file mode 100644
index 0000000..2b374e9
--- /dev/null
+++ b/package/src/nucleofind/prediction/memory.py
@@ -0,0 +1,47 @@
+import logging
+
+
+def calculate_n_from_max_memory(max_memory: float):
+ """
+ Calculates the number of processing cores (n) based on the provided maximum memory
+ constraint and a heuristic mapping of memory per core.
+
+ This function uses a predefined heuristic mapping of memory per core to determine
+ the number of cores that can fit within the specified maximum memory. The input
+ memory value is expected to be a string representing the maximum available memory.
+
+ Parameters:
+ max_memory (str): The maximum allowable memory in string format.
+
+ Returns:
+ int: The calculated number of cores based on the memory constraint.
+ """
+ if not max_memory:
+ return 1
+
+ try:
+ max_memory = float(max_memory)
+ except ValueError:
+ logging.error(
+ f"Invalid max_memory value: {max_memory}. Please provide a valid number in GB, ignoring for now."
+ )
+ return 1
+
+ # Check to ensure the max_memory isn't too big (i.e. passed in MB rather than GB)
+ if max_memory > 1024:
+ max_memory = max_memory // 1024
+ logging.warning(
+ f"NucleoFind has converted the max_memory argument to GB. Max Memory = {max_memory} GB"
+ )
+
+ if max_memory < 1:
+ logging.error(
+ f"Invalid max_memory value: {max_memory}. Please provide a valid number in GB, ignoring for now."
+ )
+ return 1
+
+ if max_memory < 4 :
+ return 1
+
+ max_cpus = (max_memory - 1.47) // 2.37
+ return max_cpus
diff --git a/package/src/nucleofind/prediction/predict.py b/package/src/nucleofind/prediction/predict.py
index 3c0ce09..51759ea 100644
--- a/package/src/nucleofind/prediction/predict.py
+++ b/package/src/nucleofind/prediction/predict.py
@@ -13,6 +13,7 @@
from .save import save_grid
from ..logs import setup_logging
from .config import Configuration, MapType
+from .memory import calculate_n_from_max_memory
class NucleoFind:
@@ -188,6 +189,17 @@ def run():
n_threads=args.nthreads,
**vars(model_configuration),
)
+
+ # Set number of threads
+ # Default to 1 thread if no nthreads specified and no max_memory specified
+ # If max_memory specified, calculate nthreads based on max_memory
+ if args.nthreads == 0 and args.max_memory:
+ configuration.n_threads = calculate_n_from_max_memory(args.max_memory)
+ elif args.nthreads == 0 and not args.max_memory:
+ configuration.n_threads = 1
+ else:
+ configuration.n_threads = args.nthreads
+
nucleofind = NucleoFind(model_path, configuration)
nucleofind.predict(
args.input,