decentriq_util 0.1.0.dev11__tar.gz → 0.1.0.dev13__tar.gz

{decentriq_util-0.1.0.dev11 → decentriq_util-0.1.0.dev13}/PKG-INFO

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

{decentriq_util-0.1.0.dev11 → decentriq_util-0.1.0.dev13}/decentriq_util/spark/__init__.py

@@ -11,7 +11,7 @@ import pyarrow.parquet as pq
 from pyspark.sql import SparkSession
 
 
-_DEFAULT_HEAP_SIZE_EXTRA_ROOM = 2 * 512 * 1024 * 1024
+_DEFAULT_HEAP_SIZE_EXTRA_ROOM_PER_CORE = 268435456
 
 
 def _get_cgroup_memory() -> int:
@@ -101,7 +101,7 @@ def _get_num_partitions(
 
     # Try to make partitions as big as possible s.t. each CPU is saturated,
     # while accounting for 50% overhead in spark.
-    task_per_cpu = 3  # try to get at least this many tasks/partitions per CPU
+    task_per_cpu = 2  # try to get at least this many tasks/partitions per CPU
     max_target_partition_size = max(
         int(0.5 * available_memory / (task_per_cpu * cpu_count)) // megabyte_in_bytes * megabyte_in_bytes,
         min_partition_size
@@ -129,16 +129,25 @@ def _get_num_partitions(
 
 def _determine_optimal_spark_settings(
         files: list[str],
-        heap_size_extra_room = _DEFAULT_HEAP_SIZE_EXTRA_ROOM
+        heap_size: Optional[int] = None,
+        heap_size_extra_room: Optional[int] = None,
+        heap_size_extra_room_per_core: Optional[int] = None
 ) -> list[Tuple[str, str]]:
-    def get_spark_memory() -> int:
-        """Determine the amount of memory available to spark"""
+    cpu_count = multiprocessing.cpu_count()
+
+    if heap_size is not None:
+        spark_memory = heap_size
+    else:
         cgroup_memory = _get_cgroup_memory()
-        spark_memory = cgroup_memory - heap_size_extra_room
-        return spark_memory
+        two_gb_in_bytes = 2147483648
+        if heap_size_extra_room is not None:
+            pass
+        elif heap_size_extra_room_per_core is not None:
+            heap_size_extra_room = heap_size_extra_room_per_core * cpu_count
+        else:
+            heap_size_extra_room = _DEFAULT_HEAP_SIZE_EXTRA_ROOM_PER_CORE * cpu_count
+        spark_memory = max(cgroup_memory - heap_size_extra_room, two_gb_in_bytes)
 
-    cpu_count = multiprocessing.cpu_count()
-    spark_memory = get_spark_memory()
     spark_memory_4096 = (spark_memory // 4096) * 4096
     num_partitions = _get_num_partitions(files, spark_memory, cpu_count)
 
@@ -155,34 +164,20 @@ def _determine_optimal_spark_settings(
 def _create_spark_session(
         name: str = "local_spark_session",
         parallelism: int = 8,
-        heap_size: Optional[int] = None,
-        heap_size_extra_room: int = _DEFAULT_HEAP_SIZE_EXTRA_ROOM,
         config: list[Tuple[str, Any]] = [],
-        java_tmp_dir = "/output",
+        java_temp_dir: str = "/scratch",
+        java_user_home_dir: str = "/scratch",
 ) -> SparkSession:
     """
     :param name: The name of the spark session.
     :param parallelism: The size of the executor pool computing internal spark tasks.
-    :param heap_size: If set then this determines the JVM heap size. Note that not all the heap size is used by spark.
-    :param heap_size_extra_room: If heap_size is unset then it is determined by the available memory in the current
-        cgroup. Then heap_size_extra_room is subtracted to arrive at the final heap size. This number should account for
-        the python runtime itself as well as other structures that may be needed during spark compute.
     :param config: Additional config settings to set on the spark config.
-    :param java_tmp_dir: Location for the JVM to store temp files.
+    :param java_temp_dir: Location for the JVM to store temp files.
+    :param java_user_home_dir: Location for the user home directory as seen by the JVM.
     :return: The spark session.
     """
     os.environ["SPARK_EXECUTOR_POOL_SIZE"] = str(parallelism)
-    os.environ["JDK_JAVA_OPTIONS"] = f'-Duser.home=/output -Djava.io.tmpdir="{java_tmp_dir}"'
-    spark_memory = None
-    if heap_size is None:
-        cgroup_memory = _get_cgroup_memory()
-        spark_memory = cgroup_memory - heap_size_extra_room
-    else:
-        spark_memory = heap_size
-
-    spark_memory_4096 = (spark_memory // 4096) * 4096
-    if "spark.driver.memory" not in dict(config):
-        config.append(("spark.driver.memory", str(spark_memory_4096)))
+    os.environ["JDK_JAVA_OPTIONS"] = f'-Duser.home="{java_user_home_dir}" -Djava.io.tmpdir="{java_temp_dir}"'
 
     ss = (
         SparkSession.builder
@@ -198,15 +193,28 @@ def _create_spark_session(
 
 
 @contextmanager
-def spark_session(temp_dir: str = "/scratch", input_files: list[str] = None, **kwargs):
+def spark_session(
+        temp_dir: str = "/scratch",
+        name: str = "Spark",
+        input_files: list[str] = None,
+        heap_size: Optional[int] = None,
+        heap_size_extra_room: Optional[int] = None,
+        heap_size_extra_room_per_core: Optional[int] = None,
+        config: list[Tuple[str, str]] = [],
+):
     """
     Create a spark session and configure it according to the enclave environment.
 
     **Parameters**:
     - `temp_dir`: Where to store temporary data such as persisted data frames
       or shuffle data.
+    - `name`: An optional name for this spark session.
     - `input_files`: A list of input files on the basis of which the partition
       size is determined.
+    - `heap_size`: The heap size to use for Spark. If not set, sensible defaults are chosen.
+    - `heap_size_extra_room`: How much extra room to leave for non-JVM tasks. If not set, sensible defaults are chosen.
+    - `heap_size_extra_room_per_core`: Same as `heap_size_extra_room` but scaled with the number of cores. This is the default behavior.
+    - `config`: Extra settings to pass to the Spark session builder.
 
     **Example**:
 
@@ -232,14 +240,13 @@ def spark_session(temp_dir: str = "/scratch", input_files: list[str] = None, **k
     """
     with tempfile.TemporaryDirectory(dir=temp_dir, prefix="java-") as java_tmp:
         with tempfile.TemporaryDirectory(dir=temp_dir, prefix="spark-") as spark_tmp:
-            config = kwargs.get("config", [])
             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", _DEFAULT_HEAP_SIZE_EXTRA_ROOM
-                    )
+                    heap_size=heap_size,
+                    heap_size_extra_room=heap_size_extra_room,
+                    heap_size_extra_room_per_core=heap_size_extra_room_per_core,
                 )
                 for key, value in optimal_settings:
                     if key not in config_dict:
@@ -248,10 +255,7 @@ def spark_session(temp_dir: str = "/scratch", input_files: list[str] = None, **k
                 config.append(
                     ("spark.local.dir", spark_tmp)
                 )
-            kwargs["config"] = config
-            if "java_tmp_dir" not in kwargs:
-                kwargs["java_tmp_dir"] = java_tmp
-            ss = _create_spark_session(**kwargs)
+            ss = _create_spark_session(name=name, java_temp_dir=java_tmp, java_user_home_dir=java_tmp, config=config)
             try:
                 yield ss
             finally:

{decentriq_util-0.1.0.dev11 → decentriq_util-0.1.0.dev13}/pyproject.toml

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