update max_tasks

Author: marshallweekley-eng

apps/bfd-pipeline-idr/src/idr_pipeline/__init__.py

@@ -27,7 +27,7 @@
 from .pipeline_stages import StagedIdrPipeline
 from .settings import (
     INCREMENTAL_IDR_JOB_GRACE_PERIOD,
-    MAX_TASKS,
+    max_tasks,
     TABLES_TO_LOAD,
     bfd_test_date,
 )
@@ -124,7 +124,7 @@ def run(source: Source, load_mode: LoadMode, load_type: LoadType, job_id: int =
     atexit.register(worker_manager.cleanup)
 
     staged_pipeline = StagedIdrPipeline(
-        max_workers=MAX_TASKS,
+        max_workers=max_tasks(),
         load_mode=load_mode,
         start_time=start_time,
         load_type=load_type,

apps/bfd-pipeline-idr/src/idr_pipeline/settings.py

@@ -72,7 +72,9 @@ def enable_prior_auth_ingestion() -> bool:
 This is useful if you've already loaded some data and you do not want to reprocess
 any batches that have already completed before this date."""
 
-MAX_TASKS = int(getenv("IDR_MAX_TASKS", "32"))
+def max_tasks() -> int:
+    return int(getenv("IDR_MAX_TASKS", "32"))
+
 """Maximum concurrent tasks to run.
 Changing this has a drastic effect on the runtime.
 In prod, we want to run as many tasks as possible without running out of memory."""

apps/bfd-pipeline-idr/test/test_pipeline.py

@@ -786,6 +786,7 @@ def _setup_pipeline_environment(info: psycopg.ConnectionInfo) -> None:
     os.environ["IDR_PER_BATCH_MIN_CONNECTIONS"] = "1"
     os.environ["IDR_PER_BATCH_MAX_CONNECTIONS"] = "1"
     os.environ["IDR_ENABLE_PRIOR_AUTH"] = "1"
+    os.environ["IDR_MAX_TASKS"] = "4"
 
 
 @pytest.fixture(scope="module")