mdales.tngl.sh
An implementation of the STAR biodiversity metric
Merge pull request #26 from quantifyearth/mwd-limit-db-workers
Limit how many CPU cores we use in DB extraction phase
author
Michael Dales
committer
GitHub
date
(Feb 3, 2026, 1:11 PM UTC)
commit
8d57df3f
8d57df3f7c5692e93e03852a66299db13c6ce46b
parent
c3f18c4e
c3f18c4e71ef7f71f4eca2dbe650465faef3b910
+20
-6
+18
-5
prepare_species/extract_species_data_psql.py
+18
-5
prepare_species/extract_species_data_psql.py
···
5
5
import os
6
6
import sys
7
7
from functools import partial
8
-
from multiprocessing import Pool
8
+
from multiprocessing import Pool, cpu_count
9
9
from pathlib import Path
10
10
from typing import Optional
11
11
···
257
257
excludes_path: Optional[Path],
258
258
output_directory_path: Path,
259
259
target_projection: Optional[str],
260
+
processes_count: int,
260
261
) -> None:
262
+
# The limiting amount here is how many concurrent connections the database can take
263
+
worker_count=min(processes_count, 20)
264
+
logger.info("Using %d workers", worker_count)
261
265
262
266
connection = psycopg2.connect(DB_CONFIG)
263
267
cursor = connection.cursor()
···
293
297
if overrides_path:
294
298
results = apply_overrides(overrides_path, results)
295
299
296
-
# The limiting amount here is how many concurrent connections the database can take
297
300
try:
298
-
with Pool(processes=20) as pool:
301
+
with Pool(processes=worker_count) as pool:
299
302
reports = pool.map(
300
303
partial(process_row, class_name, era_output_directory_path, target_projection, presence),
301
-
results
304
+
results,
302
305
)
303
306
except psycopg2.OperationalError:
304
307
sys.exit("Database connection failed for some rows, aborting")
···
316
319
"excludes": "input.excludes",
317
320
"output_directory_path": "params.output_dir",
318
321
"target_projection": "params.projection",
322
+
"processes_count": "threads",
319
323
})
320
324
def main() -> None:
321
325
parser = argparse.ArgumentParser(description="Process agregate species data to per-species-file.")
···
357
361
dest="target_projection",
358
362
default="ESRI:54017"
359
363
)
364
+
parser.add_argument(
365
+
"-j",
366
+
type=int,
367
+
required=False,
368
+
default=cpu_count() // 2,
369
+
dest="processes_count",
370
+
help="Number of concurrent threads to use."
371
+
)
360
372
args = parser.parse_args()
361
373
362
374
extract_data_per_species(
···
364
376
args.overrides,
365
377
args.excludes,
366
378
args.output_directory_path,
367
-
args.target_projection
379
+
args.target_projection,
380
+
args.processes_count,
368
381
)
369
382
370
383
if __name__ == "__main__":
+1
-1
threats/threat_summation.py
+1
-1
threats/threat_summation.py
+1
workflow/rules/species.smk
+1
workflow/rules/species.smk