Prefer len(os.sched_getaffinity(0)) over os.cpu_count()

docs/source/developer_guide/guides/2_real_world_phishing.md

@@ -388,7 +388,7 @@ To start, we will need to instantiate and set a few attributes of the `Config` c
 config = Config()
 config.mode = PipelineModes.NLP
 
-config.num_threads = os.cpu_count()
+config.num_threads = len(os.sched_getaffinity(0))
 config.feature_length = model_fea_length
 
 with open(labels_file, encoding='UTF-8') as fh:
@@ -563,7 +563,7 @@ def run_pipeline(use_stage_function: bool,
     config.mode = PipelineModes.NLP
 
     # Set the thread count to match our cpu count
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
     config.feature_length = model_fea_length
 
     with open(labels_file, encoding='UTF-8') as fh:

docs/source/developer_guide/guides/6_digital_fingerprinting_reference.md

@@ -41,7 +41,7 @@ from morpheus.cli.utils import load_labels_file
 CppConfig.set_should_use_cpp(False)
 
 config = Config()
-config.num_threads = os.cpu_count()
+config.num_threads = len(os.sched_getaffinity(0))
 config.ae = ConfigAutoEncoder()
 config.ae.feature_columns = load_labels_file(get_package_relative_file("data/columns_ae_azure.txt"))
 ```

examples/abp_pcap_detection/run.py

@@ -39,7 +39,7 @@
 @click.command()
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use.",
 )

examples/developer_guide/2_1_real_world_phishing/run.py

@@ -98,7 +98,7 @@ def run_pipeline(use_stage_function: bool,
     config.mode = PipelineModes.NLP
 
     # Set the thread count to match our cpu count
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
     config.feature_length = model_fea_length
 
     with open(labels_file, encoding='UTF-8') as fh:

examples/developer_guide/2_2_rabbitmq/read_simple.py

@@ -38,7 +38,7 @@ def run_pipeline(use_source_function: bool):
     configure_logging(log_level=logging.DEBUG)
 
     config = Config()
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     # Create a linear pipeline object
     pipeline = LinearPipeline(config)

examples/developer_guide/2_2_rabbitmq/write_simple.py

@@ -32,7 +32,7 @@ def run_pipeline():
     input_file = os.path.join(root_dir, 'examples/data/email.jsonlines')
 
     config = Config()
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     # Create a linear pipeline object
     pipeline = LinearPipeline(config)

examples/developer_guide/4_rabbitmq_cpp_stage/README.md

@@ -18,7 +18,7 @@ limitations under the License.
 # Example RabbitMQ stages
 This example builds upon the `examples/developer_guide/2_2_rabbitmq` example adding a C++ implementation for the `RabbitMQSourceStage` along with adding package install scripts.
 
-This example adds two flags to the `read_simple.py` script. A `--use_cpp` flag which defaults to `True` and a `--num_threads` flag which defaults to the number of cores on the system as returned by `os.cpu_count()`.
+This example adds two flags to the `read_simple.py` script. A `--use_cpp` flag which defaults to `True` and a `--num_threads` flag which defaults to the number of cores on the system as returned by `len(os.sched_getaffinity(0))`.
 
 ## Supported Environments
 | Environment | Supported | Notes |

examples/developer_guide/4_rabbitmq_cpp_stage/src/read_simple.py

@@ -32,7 +32,7 @@
 @click.option('--use_cpp', default=True)
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use",
 )

examples/developer_guide/4_rabbitmq_cpp_stage/src/write_simple.py

@@ -32,7 +32,7 @@ def run_pipeline():
     input_file = os.path.join(root_dir, 'examples/data/email.jsonlines')
 
     config = Config()
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     # Create a linear pipeline object
     pipeline = LinearPipeline(config)

examples/digital_fingerprinting/production/grafana/run.py

@@ -245,7 +245,7 @@ def run_pipeline(train_users,
 
     CppConfig.set_should_use_cpp(False)
 
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     config.ae = ConfigAutoEncoder()
 

examples/digital_fingerprinting/production/morpheus/dfp/utils/config_generator.py

@@ -174,7 +174,7 @@ def generate_ae_config(source: str,
                        pipeline_batch_size: int = 0,
                        edge_buffer_size: int = 0,
                        use_cpp: bool = False,
-                       num_threads: int = os.cpu_count()):
+                       num_threads: int = len(os.sched_getaffinity(0))):
     config = Config()
 
     CppConfig.set_should_use_cpp(use_cpp)

examples/digital_fingerprinting/production/morpheus/dfp_azure_pipeline.py

@@ -233,7 +233,7 @@ def run_pipeline(train_users,
 
     CppConfig.set_should_use_cpp(False)
 
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     config.ae = ConfigAutoEncoder()
 

examples/digital_fingerprinting/production/morpheus/dfp_duo_pipeline.py

@@ -230,7 +230,7 @@ def run_pipeline(train_users,
 
     CppConfig.set_should_use_cpp(False)
 
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     config.ae = ConfigAutoEncoder()
 

examples/digital_fingerprinting/production/morpheus/notebooks/dfp_azure_inference.ipynb

@@ -219,7 +219,7 @@
     "\n",
     "CppConfig.set_should_use_cpp(False)\n",
     "\n",
-    "config.num_threads = os.cpu_count()\n",
+    "config.num_threads = len(os.sched_getaffinity(0))\n",
     "\n",
     "config.ae = ConfigAutoEncoder()\n",
     "\n",

examples/digital_fingerprinting/production/morpheus/notebooks/dfp_azure_training.ipynb

@@ -216,7 +216,7 @@
     "\n",
     "CppConfig.set_should_use_cpp(False)\n",
     "\n",
-    "config.num_threads = os.cpu_count()\n",
+    "config.num_threads = len(os.sched_getaffinity(0))\n",
     "\n",
     "config.ae = ConfigAutoEncoder()\n",
     "\n",

examples/digital_fingerprinting/production/morpheus/notebooks/dfp_duo_inference.ipynb

@@ -218,7 +218,7 @@
     "\n",
     "CppConfig.set_should_use_cpp(False)\n",
     "\n",
-    "config.num_threads = os.cpu_count()\n",
+    "config.num_threads = len(os.sched_getaffinity(0))\n",
     "\n",
     "config.ae = ConfigAutoEncoder()\n",
     "\n",

examples/digital_fingerprinting/production/morpheus/notebooks/dfp_duo_training.ipynb

@@ -217,7 +217,7 @@
     "\n",
     "CppConfig.set_should_use_cpp(False)\n",
     "\n",
-    "config.num_threads = os.cpu_count()\n",
+    "config.num_threads = len(os.sched_getaffinity(0))\n",
     "\n",
     "config.ae = ConfigAutoEncoder()\n",
     "\n",

examples/digital_fingerprinting/starter/run_cloudtrail_dfp.py

@@ -39,7 +39,7 @@
 @click.command()
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use",
 )

examples/digital_fingerprinting/visualization/dfp_viz_azure_pipeline.py

@@ -183,7 +183,7 @@ def run_pipeline(train_users,
 
     CppConfig.set_should_use_cpp(False)
 
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     config.ae = ConfigAutoEncoder()
 

examples/digital_fingerprinting/visualization/dfp_viz_duo_pipeline.py

@@ -186,7 +186,7 @@ def run_pipeline(train_users,
 
     CppConfig.set_should_use_cpp(False)
 
-    config.num_threads = os.cpu_count()
+    config.num_threads = len(os.sched_getaffinity(0))
 
     config.ae = ConfigAutoEncoder()
 

examples/doca/run_udp_convert.py

@@ -43,7 +43,7 @@
 )
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     show_default=True,
     help="Number of internal pipeline threads to use.",

examples/gnn_fraud_detection_pipeline/run.py

@@ -38,7 +38,7 @@
 @click.command()
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use.",
 )

examples/llm/agents/run.py

@@ -27,7 +27,7 @@ def run():
 @run.command(help="Runs a simple finite pipeline with a single execution of a LangChain agent from a fixed input")
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use",
 )
@@ -67,7 +67,7 @@ def simple(**kwargs):
 @run.command(help="Runs a pipeline LangChain agents which pulls inputs from a Kafka message bus")
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use",
 )

examples/llm/completion/run.py

@@ -28,7 +28,7 @@ def run():
 @run.command()
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use",
 )

examples/llm/rag/run.py

@@ -27,7 +27,7 @@ def run():
 @run.command()
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use",
 )

examples/llm/vdb_upload/run.py

@@ -78,7 +78,7 @@ def run():
 )
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use",
 )

examples/log_parsing/run.py

@@ -35,7 +35,7 @@
 @click.command()
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use.",
 )

examples/ransomware_detection/run.py

@@ -41,7 +41,7 @@
 @click.option('--use_cpp', default=False, help="Enable C++ execution for this pipeline, currently this is unsupported.")
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use.",
 )

examples/sid_visualization/run.py

@@ -123,7 +123,7 @@ def _generate_frames(self):
 @click.option('--use_cpp', default=True)
 @click.option(
     "--num_threads",
-    default=os.cpu_count(),
+    default=len(os.sched_getaffinity(0)),
     type=click.IntRange(min=1),
     help="Number of internal pipeline threads to use.",
 )

python/morpheus/morpheus/cli/commands.py

@@ -265,7 +265,7 @@ def install(**kwargs):
 
 @cli.group(short_help="Run one of the available pipelines", no_args_is_help=True, cls=AliasedGroup)
 @click.option('--num_threads',
-              default=os.cpu_count(),
+              default=len(os.sched_getaffinity(0)),
               type=click.IntRange(min=1),
               help="Number of internal pipeline threads to use")
 @click.option('--pipeline_batch_size',

python/morpheus/morpheus/stages/input/http_server_source_stage.py

@@ -68,7 +68,7 @@ class HttpServerSourceStage(PreallocatorMixin, SingleOutputSource):
         Maximum number of requests to queue before rejecting requests. If `None` then `config.edge_buffer_size` will be
         used.
     num_server_threads : int, default None
-        Number of threads to use for the HTTP server. If `None` then `os.cpu_count()` will be used.
+        Number of threads to use for the HTTP server. If `None` then `len(os.sched_getaffinity(0))` will be used.
     max_payload_size : int, default 10
         The maximum size in megabytes of the payload that the server will accept in a single request.
     request_timeout_secs : int, default 30
@@ -117,7 +117,7 @@ def __init__(self,
         self._sleep_time = sleep_time
         self._queue_timeout = queue_timeout
         self._max_queue_size = max_queue_size or config.edge_buffer_size
-        self._num_server_threads = num_server_threads or os.cpu_count()
+        self._num_server_threads = num_server_threads or len(os.sched_getaffinity(0))
         self._max_payload_size_bytes = max_payload_size * 1024 * 1024
         self._request_timeout_secs = request_timeout_secs
         self._lines = lines

python/morpheus/morpheus/stages/output/http_server_sink_stage.py

@@ -65,7 +65,7 @@ class HttpServerSinkStage(PassThruTypeMixin, SinglePortStage):
         Maximum number of requests to queue before rejecting requests. If `None` then `config.edge_buffer_size` will be
         used. Once the queue is full, the incoming edge buffer will begin to fill up.
     num_server_threads : int, default None
-        Number of threads to use for the HTTP server. If `None` then `os.cpu_count()` will be used.
+        Number of threads to use for the HTTP server. If `None` then `len(os.sched_getaffinity(0))` will be used.
     max_rows_per_response : int, optional
         Maximum number of rows to include in a single response, by default 10000.
     overflow_pct: float, optional
@@ -103,7 +103,7 @@ def __init__(self,
         self._port = port
         self._endpoint = endpoint
         self._method = method
-        self._num_server_threads = num_server_threads or os.cpu_count()
+        self._num_server_threads = num_server_threads or len(os.sched_getaffinity(0))
         self._max_rows_per_response = max_rows_per_response
         self._overflow_pct = overflow_pct
         self._request_timeout_secs = request_timeout_secs

tests/common/test_http_server.py

@@ -42,7 +42,7 @@ def make_parse_fn(status: HTTPStatus = HTTPStatus.OK,
 @pytest.mark.parametrize("method", ["GET", "POST", "PUT"])
 @pytest.mark.parametrize("use_callback", [True, False])
 @pytest.mark.parametrize("use_context_mgr", [True, False])
-@pytest.mark.parametrize("num_threads", [1, 2, min(8, os.cpu_count())])
+@pytest.mark.parametrize("num_threads", [1, 2, min(8, len(os.sched_getaffinity(0)))])
 @pytest.mark.parametrize("status,content_type,content",
                          [(HTTPStatus.OK, MimeTypes.TEXT.value, "OK"),
                           (HTTPStatus.OK, MimeTypes.JSON.value, '{"test": "OK"}'),