Add new setting for worker level configuration as well to enable better system utilization and tuning for shard-based ingestion

Author: BrianMichell

src/mdio/core/config.py

@@ -20,6 +20,11 @@ class MDIOSettings(BaseSettings):
         description="Number of CPUs to use for import operations",
         alias="MDIO__IMPORT__CPU_COUNT",
     )
+    import_zarr_threads: int = Field(
+        default=4,
+        description="Number of Zarr threads per worker during import operations",
+        alias="MDIO__IMPORT__ZARR_THREADS",
+    )
 
     # Grid sparsity configuration
     grid_sparsity_ratio_warn: float = Field(

src/mdio/segy/_workers.py

@@ -89,9 +89,11 @@ def trace_worker(  # noqa: PLR0913
     Returns:
         SummaryStatistics object containing statistics about the written traces.
     """
-    # Setting the zarr config to 1 thread to ensure we honor the `MDIO__IMPORT__CPU_COUNT` environment variable.
-    # The Zarr 3 engine utilizes multiple threads. This can lead to resource contention and unpredictable memory usage.
-    zarr_config.set({"threading.max_workers": 1})
+    # Configure zarr threading based on MDIO__IMPORT__ZARR_THREADS setting.
+    # With shard-aligned writes, fewer workers are used, so we allow more zarr threads per worker
+    # to better utilize available CPU resources.
+    settings = MDIOSettings()
+    zarr_config.set({"threading.max_workers": settings.import_zarr_threads})
 
     region_slices = tuple(region.values())
     local_grid_map = grid_map[region_slices[:-1]]  # minus last (vertical) axis

tests/unit/test_environment.py

@@ -16,6 +16,7 @@ class TestEnvironment:
         [
             ("MDIO__EXPORT__CPU_COUNT", "8", "export_cpus", 8),
             ("MDIO__IMPORT__CPU_COUNT", "4", "import_cpus", 4),
+            ("MDIO__IMPORT__ZARR_THREADS", "8", "import_zarr_threads", 8),
             ("MDIO__GRID__SPARSITY_RATIO_WARN", "3.5", "grid_sparsity_ratio_warn", 3.5),
             ("MDIO__GRID__SPARSITY_RATIO_LIMIT", "15.0", "grid_sparsity_ratio_limit", 15.0),
         ],