Add SPARGLIM_SERVER_CUSTOM_CONFIG for custom config spark (#15)

* Add SPARGLIM_SERVER_CUSTOM_CONFIG for custom config spark * Fix custom config
Wh1isper · Dec 22, 2023 · 99a1a42 · 99a1a42
1 parent 0358fa8
commit 99a1a42
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git a/dev/sparglim-server/k8s/deployment.yaml b/dev/sparglim-server/k8s/deployment.yaml
@@ -33,3 +33,6 @@ spec:
               valueFrom:
                 fieldRef:
                   fieldPath: metadata.name
+            - name: SPARGLIM_SERVER_CUSTOM_CONFIG
+              # json string for custom config, see https://spark.apache.org/docs/latest/configuration.html
+              value: "{}"
diff --git a/examples/sparglim-server/k8s/deployment.yaml b/examples/sparglim-server/k8s/deployment.yaml
@@ -33,3 +33,6 @@ spec:
               valueFrom:
                 fieldRef:
                   fieldPath: metadata.name
+            - name: SPARGLIM_SERVER_CUSTOM_CONFIG
+              # json string for custom config, see https://spark.apache.org/docs/latest/configuration.html
+              value: "{}"
diff --git a/sparglim/server/cli.py b/sparglim/server/cli.py
@@ -1,5 +1,6 @@
 #  Copyright (c) 2023 Wh1isper
 #  Licensed under the BSD 3-Clause License
+import json
 import os
 import signal
 
@@ -12,9 +13,12 @@
 @click.command()
 @click.option("--mode", default=None)
 @click.option("--root_dir", default="./")
-def start(mode, root_dir):
+@click.option("--custom_config", default=None)
+def start(mode, root_dir, custom_config):
+    if custom_config is not None:
+        custom_config = json.loads(custom_config)
     if mode:
-        Daemon(mode=mode, root_dir=root_dir).start_and_daemon()
+        Daemon(mode=mode, root_dir=root_dir, custom_config=custom_config).start_and_daemon()
     else:
         Daemon(root_dir=root_dir).start_and_daemon()
 

diff --git a/sparglim/server/daemon.py b/sparglim/server/daemon.py
@@ -1,5 +1,6 @@
 #  Copyright (c) 2023 Wh1isper
 #  Licensed under the BSD 3-Clause License
+import json
 import os
 import signal
 import subprocess
@@ -9,7 +10,7 @@
 from datetime import datetime
 from functools import wraps
 from pathlib import Path
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 import psutil
 
@@ -31,6 +32,7 @@ class Daemon:
     # Stop the daemon when the KeyboardInterrupt is raised
     # This will not start duplicate connect server, unless using different root_dir and port
     ENV_MASTER_MODE = "SPARGLIM_SERVER_MODE"
+    ENV_CUSTOM_CONFIG = "SPARGLIM_SERVER_CUSTOM_CONFIG"
     DEFAULT_MODE = "local"
     SPARK_IDENT_STRING = "sparglim-connect-server"  # SPARK_IDENT_STRING
 
@@ -40,8 +42,10 @@ def __init__(
         root_dir: str = "./",
         *,
         k8s_config_path: Optional[str] = None,
+        custom_config: Optional[Dict[str, Any]] = None,
     ):
         self.mode = mode or os.getenv(self.ENV_MASTER_MODE, self.DEFAULT_MODE)
+        self.custom_config = custom_config or json.loads(os.getenv(self.ENV_CUSTOM_CONFIG, "{}"))
 
         self.root_dir = Path(root_dir)
 
@@ -68,7 +72,7 @@ def __init__(
                 f"org.apache.spark:spark-connect_{SCALA_VERSION}:{SPARK_VERSION}"
             )
         self.builder = SparkEnvConfiger().config_connect_server(
-            self.mode, k8s_config_path=k8s_config_path
+            self.mode, custom_config=self.custom_config, k8s_config_path=k8s_config_path
         )
 
         self._tailer: Optional[Tailer] = None