Refine Presto worker config validation

tools/deployment/presto-clp/scripts/init.py

@@ -1,9 +1,12 @@
 import argparse
 import logging
+import math
+import os
 import sys
 from pathlib import Path
 from typing import Any, Dict, Optional
 
+import psutil
 import yaml
 from dotenv import dotenv_values
 
@@ -226,21 +229,86 @@ def _add_worker_env_vars(coordinator_common_env_file_path: Path, env_vars: Dict[
     config = dotenv_values(coordinator_common_env_file_path)
 
     try:
-        env_vars["PRESTO_COORDINATOR_CONFIGPROPERTIES_DISCOVERY_URI"] = (
-            f'http://{config["PRESTO_COORDINATOR_SERVICENAME"]}'
-            f':{config["PRESTO_COORDINATOR_HTTPPORT"]}'
+        coordinator_service_name = _get_required_env_var(
+            config, "PRESTO_COORDINATOR_SERVICENAME", coordinator_common_env_file_path
         )
-    except KeyError as e:
-        logger.error(
-            "Missing required key %s in '%s'",
-            e,
-            coordinator_common_env_file_path,
+        coordinator_http_port = _get_required_env_var(
+            config, "PRESTO_COORDINATOR_HTTPPORT", coordinator_common_env_file_path
         )
+    except KeyError:
+        return False
+
+    env_vars["PRESTO_COORDINATOR_CONFIGPROPERTIES_DISCOVERY_URI"] = (
+        f"http://{coordinator_service_name}:{coordinator_http_port}"
+    )
+
+    if not _configure_worker_memory_env_vars(env_vars):
         return False
 
     return True
 
 
+def _configure_worker_memory_env_vars(env_vars: Dict[str, str]) -> bool:
+    """Adds memory-related environment variables for the worker."""
+
+    total_memory_bytes = psutil.virtual_memory().total
+    total_memory_gb = total_memory_bytes / (1024**3)
+    total_memory_gb_floor = max(1, math.floor(total_memory_gb))
+
+    system_memory_gb = min(
+        total_memory_gb_floor,
+        max(1, math.floor(total_memory_gb * 0.90)),
+    )
+    query_memory_gb = min(
+        system_memory_gb,
+        max(1, math.floor(system_memory_gb * (2 / 3))),
+    )
+    system_mem_limit_gb = min(
+        total_memory_gb_floor,
+        max(system_memory_gb, max(1, math.floor(total_memory_gb * 0.94))),
+    )
+
+    env_vars["PRESTO_WORKER_CONFIGPROPERTIES_QUERY_MEMORY_GB"] = str(query_memory_gb)
+    env_vars["PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEMORY_GB"] = str(system_memory_gb)
+    env_vars["PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEM_LIMIT_GB"] = str(system_mem_limit_gb)
+
+    logger.info(
+        (
+            "Detected %.2f GB of system memory; configured Presto worker memory limits "
+            "(query-memory-gb=%s, system-memory-gb=%s, system-mem-limit-gb=%s)."
+        ),
+        total_memory_gb,
+        env_vars["PRESTO_WORKER_CONFIGPROPERTIES_QUERY_MEMORY_GB"],
+        env_vars["PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEMORY_GB"],
+        env_vars["PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEM_LIMIT_GB"],
+    )
+
+    return True
+
+
+def _get_required_env_var(config: Dict[str, Optional[str]], key: str, env_file_path: Path) -> str:
+    """
+    Fetches a required entry from a dotenv-style config.
+
+    :param config: Mapping of config keys to their string values.
+    :param key: The key that must exist in `config` with a non-empty value.
+    :param env_file_path: Path to the file that produced `config`, for logging context.
+    :return: The string value associated with `key`.
+    :raises KeyError: If `key` is missing or empty.
+    """
+
+    value = config.get(key)
+    if value is None:
+        logger.error("Missing required key %s in '%s'", key, env_file_path)
+        raise KeyError(key)
+
+    if value == "":
+        logger.error("Empty value for required key %s in '%s'", key, env_file_path)
+        raise KeyError(key)
+
+    return value
+
+
 def _generate_worker_clp_properties(
     worker_config_template_path: Path, env_vars: Dict[str, str]
 ) -> bool:

tools/deployment/presto-clp/scripts/requirements.txt

@@ -1,2 +1,3 @@
+psutil
 python-dotenv
 PyYAML

tools/deployment/presto-clp/worker.env

@@ -1,4 +1,7 @@
 PRESTO_WORKER_HTTP_PORT=8080
 
-# node.properties
+PRESTO_WORKER_CONFIGPROPERTIES_QUERY_MEMORY_GB=38
+PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEMORY_GB=57
+PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEM_LIMIT_GB=60
+
 PRESTO_WORKER_NODEPROPERTIES_LOCATION=worker-location

tools/deployment/presto-clp/worker/config-template/config.properties

@@ -3,3 +3,6 @@ http-server.http.port=${PRESTO_WORKER_HTTP_PORT}
 shutdown-onset-sec=1
 register-test-functions=false
 runtime-metrics-collection-enabled=false
+system-memory-gb=${PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEMORY_GB}
+query-memory-gb=${PRESTO_WORKER_CONFIGPROPERTIES_QUERY_MEMORY_GB}
+system-mem-limit-gb=${PRESTO_WORKER_CONFIGPROPERTIES_SYSTEM_MEM_LIMIT_GB}