decentriq_util 0.1.0.dev8__tar.gz → 0.1.0.dev10__tar.gz

{decentriq_util-0.1.0.dev8 → decentriq_util-0.1.0.dev10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: decentriq_util
-Version: 0.1.0.dev8
+Version: 0.1.0.dev10
 Summary: Utility functions for Python code run in Decentriq workers
 Author: hafenr
 Author-email: robin.hafen@gmail.com

{decentriq_util-0.1.0.dev8 → decentriq_util-0.1.0.dev10}/decentriq_util/spark/__init__.py

@@ -11,6 +11,9 @@ import pyarrow.parquet as pq
 from pyspark.sql import SparkSession
 
 
+_DEFAULT_HEAP_SIZE_EXTRA_ROOM = 2 * 512 * 1024 * 1024
+
+
 def _get_cgroup_memory() -> int:
     max_memory_str = None
     if os.path.isfile("/sys/fs/cgroup/memory.max"):
@@ -126,12 +129,12 @@ def _get_num_partitions(
 
 def _determine_optimal_spark_settings(
         files: list[str],
-        heap_size_extra_room: Optional[int] = 2 * 512 * 1024 * 1024
+        heap_size_extra_room = _DEFAULT_HEAP_SIZE_EXTRA_ROOM
 ) -> list[Tuple[str, str]]:
     def get_spark_memory() -> int:
         """Determine the amount of memory available to spark"""
         cgroup_memory = _get_cgroup_memory()
-        spark_memory = cgroup_memory - (heap_size_extra_room or 512 * 1024 * 1024)
+        spark_memory = cgroup_memory - heap_size_extra_room
         return spark_memory
 
     cpu_count = multiprocessing.cpu_count()
@@ -145,7 +148,6 @@ def _determine_optimal_spark_settings(
         ("spark.driver.cores", str(cpu_count)),
         ("spark.driver.memory", str(spark_memory_4096)),
     ]
-    print("Spark settings:\n" + f"\n".join([str(x) for x in settings]))
 
     return settings
 
@@ -154,7 +156,7 @@ def _create_spark_session(
         name: str = "local_spark_session",
         parallelism: int = 8,
         heap_size: Optional[int] = None,
-        heap_size_extra_room: int = 512 * 1024 * 1024,
+        heap_size_extra_room: int = _DEFAULT_HEAP_SIZE_EXTRA_ROOM,
         config: list[Tuple[str, Any]] = [],
         java_tmp_dir = "/output",
 ) -> SparkSession:
@@ -189,6 +191,9 @@ def _create_spark_session(
     )
     for (key, value) in config:
         ss = ss.config(key, value)
+
+    print("Spark settings:\n" + f"\n".join([str(x) for x in config]))
+
     return ss.getOrCreate()
 
 
@@ -231,7 +236,10 @@ def spark_session(temp_dir: str = "/scratch", input_files: list[str] = None, **k
             config_dict = dict(config)
             if input_files is not None and input_files:
                 optimal_settings = _determine_optimal_spark_settings(
-                    input_files, heap_size_extra_room=kwargs.get("heap_size_extra_room")
+                    input_files,
+                    heap_size_extra_room=kwargs.get(
+                        "heap_size_extra_room", _DEFAULT_HEAP_SIZE_EXTRA_ROOM
+                    )
                 )
                 for key, value in optimal_settings:
                     if key not in config_dict:

{decentriq_util-0.1.0.dev8 → decentriq_util-0.1.0.dev10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "decentriq_util"
-version = "0.1.0-dev8"
+version = "0.1.0-dev10"
 description = "Utility functions for Python code run in Decentriq workers"
 authors = ["hafenr <robin.hafen@gmail.com>"]