refactor: unify env-var config in common/config.py

Before: SPARK_EXECUTOR_DATA_DIR, SPARK_EXECUTOR_JOBS_DIR, and
YARN_RESOURCE_MANAGER_URL were each read directly via os.environ.get()
inside the module that used them. Log level was hardcoded DEBUG in
common/logging.py. There was no single file showing what the full set
of env vars the app reads is.

After: common/config.py defines a single Settings dataclass that
reads all env vars at import time and exposes them as fields on a
module-level singleton. App code uses "from common.config import
settings; settings.data_dir" etc. New SPARK_EXECUTOR_LOG_LEVEL env
var controls stderr + info file verbosity (debug file always gets full
DEBUG).

Improvements:
  - One file lists every env var the app reads (was: grep the codebase)
  - Tests can monkeypatch fields on the settings singleton directly
    instead of monkeypatching the env + reloading
  - Adding a new env var means adding one field in config.py, not
    editing 3+ call sites
  - settings.reload() method for tests that prefer env-var style

Out of scope (kept where they are):
  - GUNICORN_* env vars live in gunicorn.conf.py (gunicorn concept)
  - PYTHONUNBUFFERED in Dockerfile (Python runtime flag)
  - SPARK_SUBMIT_OPTS not in config (JVM flag, not Python)

Test changes:
  - test_job_writer.py: settings.jobs_dir instead of monkeypatching
    SPARK_EXECUTOR_JOBS_DIR
  - test_yarn_client.py: settings.yarn_resource_manager_url instead of
    monkeypatching YARN_RESOURCE_MANAGER_URL
  - test_generate_tool.py: same as job_writer
  - Each test file gets an autouse fixture that snapshots+restores
    settings so one test mutation does not leak into the next

116/116 still pass. Live verified: SPARK_EXECUTOR_LOG_LEVEL=INFO
suppresses DEBUG loguru output as expected.
This commit is contained in:
Claude
2026-06-25 10:52:53 +08:00
parent 3ec9ea37fd
commit f8e367d43c
9 changed files with 204 additions and 48 deletions
+90
View File
@@ -0,0 +1,90 @@
# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
Single source of truth for environment-driven application configuration.
Why this module exists: every prior `os.environ.get("SPARK_EXECUTOR_*")`
was scattered across `connection_store.py`, `pending_store.py`,
`yarn_client.py`, `job_writer.py`. The same env var name appeared in
multiple files, and there was no single place to see "what's
configurable?". This file is that place.
All application code should:
from common.config import settings
... settings.data_dir / settings.yarn_resource_manager_url / ...
...rather than `os.environ.get(...)` directly. This way:
- The env var name appears in exactly one place.
- Ops can audit the full config surface from one file.
- Tests can monkeypatch fields on the `settings` singleton directly
(more reliable than monkeypatching env + reloading).
NOT in this module:
- Gunicorn-specific env vars (workers, threads, bind, timeout) live in
gunicorn.conf.py because they configure gunicorn, not the app.
- `PYTHONUNBUFFERED` — Python runtime flag, set in Dockerfile.
- `SPARK_SUBMIT_OPTS` — JVM flags, not Python config; users pass them
to spark-submit directly.
"""
import os
from dataclasses import dataclass
@dataclass
class Settings:
"""Application settings loaded from environment variables at import time.
Mutable so tests can reassign fields directly. For tests that use
`monkeypatch.setenv(...)`, call `settings.reload()` to re-read from
the (now-patched) environment.
"""
# --- Data persistence (./data/ in dev, /var/lib/... in prod) ---
data_dir: str = "./data"
# --- Job files (LLM-generated PySpark) ---
# Defaults to <data_dir>/jobs; can be pointed at a larger disk via
# SPARK_EXECUTOR_JOBS_DIR (e.g. /var/spark-jobs on a big-disk host).
jobs_dir: str = ""
# --- YARN REST client (fallback for Job.yarn_rm_url snapshot) ---
# Set in the env OR per-Connection via save_connection.
yarn_resource_manager_url: str | None = None
# --- Loguru ---
# Controls stderr verbosity and the info-level file sink. The debug
# file sink always captures DEBUG (full audit trail).
log_level: str = "DEBUG"
@classmethod
def from_env(cls) -> "Settings":
data_dir = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data")
jobs_dir = os.environ.get(
"SPARK_EXECUTOR_JOBS_DIR", os.path.join(data_dir, "jobs")
)
return cls(
data_dir=data_dir,
jobs_dir=jobs_dir,
# `or None` collapses empty string to None for the URL fallback
yarn_resource_manager_url=os.environ.get("YARN_RESOURCE_MANAGER_URL") or None,
log_level=os.environ.get("SPARK_EXECUTOR_LOG_LEVEL", "DEBUG"),
)
def reload(self) -> "Settings":
"""Re-read from environment, mutate in place, return self (chainable).
Useful in tests that `monkeypatch.setenv(...)` and want the
singleton to pick up the new value without a re-import.
"""
fresh = self.from_env()
self.data_dir = fresh.data_dir
self.jobs_dir = fresh.jobs_dir
self.yarn_resource_manager_url = fresh.yarn_resource_manager_url
self.log_level = fresh.log_level
return self
# Module-level singleton. Loaded once at import time.
settings = Settings.from_env()
+10 -3
View File
@@ -6,6 +6,10 @@
Process-wide loguru configuration. Import `logger` from here in every Process-wide loguru configuration. Import `logger` from here in every
module instead of instantiating new loggers. module instead of instantiating new loggers.
Log level is controlled via `SPARK_EXECUTOR_LOG_LEVEL` (see common/config.py):
DEBUG - default; full verbosity
INFO - quieter; recommended for production
Levels used in this project: Levels used in this project:
DEBUG - entry/exit of public tools, subprocess commands, file I/O paths DEBUG - entry/exit of public tools, subprocess commands, file I/O paths
INFO - business events (job submitted, status changed, connection saved) INFO - business events (job submitted, status changed, connection saved)
@@ -15,13 +19,16 @@ Levels used in this project:
import sys import sys
from pathlib import Path from pathlib import Path
from loguru import logger from loguru import logger
from common.config import settings
Path("data/logs/debug").mkdir(parents=True, exist_ok=True) Path("data/logs/debug").mkdir(parents=True, exist_ok=True)
Path("data/logs/info").mkdir(parents=True, exist_ok=True) Path("data/logs/info").mkdir(parents=True, exist_ok=True)
logger.remove() logger.remove()
logger.add( logger.add(
sys.stderr, sys.stderr,
level="DEBUG", level=settings.log_level,
format=( format=(
"<green>{time:HH:mm:ss.SSS}</green> | " "<green>{time:HH:mm:ss.SSS}</green> | "
"<level>{level: <7}</level> | " "<level>{level: <7}</level> | "
@@ -34,7 +41,7 @@ logger.add(
logger.add( logger.add(
"data/logs/debug/{time:YYYY-MM-DD}.log", "data/logs/debug/{time:YYYY-MM-DD}.log",
level="DEBUG", level="DEBUG", # always full at the file level for audit
enqueue=True, enqueue=True,
retention="30 days", retention="30 days",
compression="gz", compression="gz",
@@ -43,7 +50,7 @@ logger.add(
logger.add( logger.add(
"data/logs/info/{time:YYYY-MM-DD}.log", "data/logs/info/{time:YYYY-MM-DD}.log",
level="INFO", level=settings.log_level,
enqueue=True, enqueue=True,
retention="30 days", retention="30 days",
compression="gz", compression="gz",
+2 -1
View File
@@ -9,10 +9,11 @@ import tempfile
from pathlib import Path from pathlib import Path
from threading import Lock from threading import Lock
from common.config import settings
from common.logging import logger from common.logging import logger
from spark_executor.models import Connection from spark_executor.models import Connection
DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") DEFAULT_DATA_DIR = settings.data_dir
DEFAULT_FILE_NAME = "connections.json" DEFAULT_FILE_NAME = "connections.json"
+7 -6
View File
@@ -16,20 +16,21 @@ import os
import secrets import secrets
from datetime import datetime from datetime import datetime
from common.config import settings
from common.logging import logger from common.logging import logger
DEFAULT_JOBS_DIR = "./data/jobs"
ENV_JOBS_DIR = "SPARK_EXECUTOR_JOBS_DIR" ENV_JOBS_DIR = "SPARK_EXECUTOR_JOBS_DIR"
def resolve_jobs_dir(jobs_dir: str | None = None) -> str: def resolve_jobs_dir(jobs_dir: str | None = None) -> str:
"""Pick the effective jobs directory in priority order: arg > env > default.""" """Pick the effective jobs directory in priority order: arg > settings.jobs_dir.
`settings.jobs_dir` itself is computed in `common/config.py` as:
SPARK_EXECUTOR_JOBS_DIR env var (if set) > <data_dir>/jobs
"""
if jobs_dir is not None: if jobs_dir is not None:
return jobs_dir return jobs_dir
env = os.environ.get(ENV_JOBS_DIR) return settings.jobs_dir
if env:
return env
return DEFAULT_JOBS_DIR
def write_job_file(code: str, jobs_dir: str | None = None) -> str: def write_job_file(code: str, jobs_dir: str | None = None) -> str:
+2 -1
View File
@@ -9,10 +9,11 @@ import tempfile
from pathlib import Path from pathlib import Path
from threading import Lock from threading import Lock
from common.config import settings
from common.logging import logger from common.logging import logger
from spark_executor.models import PendingSubmission from spark_executor.models import PendingSubmission
DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") DEFAULT_DATA_DIR = settings.data_dir
DEFAULT_FILE_NAME = "pending_jobs.json" DEFAULT_FILE_NAME = "pending_jobs.json"
+2 -2
View File
@@ -18,10 +18,10 @@ var if unset. This matches the pattern the original Connection.yarn_rm_url
field was designed for, but no longer requires the `yarn` CLI to interpret it. field was designed for, but no longer requires the `yarn` CLI to interpret it.
""" """
import json import json
import os
import httpx import httpx
from common.config import settings
from common.logging import logger from common.logging import logger
@@ -35,7 +35,7 @@ class YarnConfigError(YarnError):
def _base_url(yarn_rm_url: str | None) -> str: def _base_url(yarn_rm_url: str | None) -> str:
"""Resolve and validate the RM URL. Raises YarnConfigError if unusable.""" """Resolve and validate the RM URL. Raises YarnConfigError if unusable."""
url = yarn_rm_url or os.environ.get("YARN_RESOURCE_MANAGER_URL") url = yarn_rm_url or settings.yarn_resource_manager_url
if not url: if not url:
raise YarnConfigError( raise YarnConfigError(
"YARN ResourceManager URL is not configured. " "YARN ResourceManager URL is not configured. "
+21 -3
View File
@@ -2,11 +2,29 @@
import os import os
from pathlib import Path from pathlib import Path
import pytest
from common import config
from spark_executor.tools import generate from spark_executor.tools import generate
@pytest.fixture(autouse=True)
def _restore_settings():
snapshot = config.Settings(
data_dir=config.settings.data_dir,
jobs_dir=config.settings.jobs_dir,
yarn_resource_manager_url=config.settings.yarn_resource_manager_url,
log_level=config.settings.log_level,
)
yield
config.settings.data_dir = snapshot.data_dir
config.settings.jobs_dir = snapshot.jobs_dir
config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
config.settings.log_level = snapshot.log_level
def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path): def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
monkeypatch.chdir(tmp_path) config.settings.jobs_dir = str(tmp_path / "data" / "jobs")
out = generate.generate_job_file( out = generate.generate_job_file(
"from pyspark.sql import SparkSession\n" "from pyspark.sql import SparkSession\n"
"spark = SparkSession.builder.getOrCreate()\n" "spark = SparkSession.builder.getOrCreate()\n"
@@ -20,8 +38,8 @@ def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
assert "SparkSession.builder.getOrCreate()" in f.read() assert "SparkSession.builder.getOrCreate()" in f.read()
def test_generate_uses_env_var_when_set(monkeypatch, tmp_path: Path): def test_generate_uses_settings_jobs_dir(tmp_path: Path):
monkeypatch.setenv("SPARK_EXECUTOR_JOBS_DIR", str(tmp_path / "custom")) config.settings.jobs_dir = str(tmp_path / "custom")
out = generate.generate_job_file("x = 1\n") out = generate.generate_job_file("x = 1\n")
assert out["script_path"].startswith(str(tmp_path / "custom")) assert out["script_path"].startswith(str(tmp_path / "custom"))
assert os.path.isfile(out["script_path"]) assert os.path.isfile(out["script_path"])
+45 -25
View File
@@ -2,29 +2,48 @@
import os import os
from pathlib import Path from pathlib import Path
import pytest
from common import config
from spark_executor.core.job_writer import ( from spark_executor.core.job_writer import (
DEFAULT_JOBS_DIR,
ENV_JOBS_DIR, ENV_JOBS_DIR,
resolve_jobs_dir, resolve_jobs_dir,
write_job_file, write_job_file,
) )
# --- resolve_jobs_dir (priority: arg > env > default) --- @pytest.fixture(autouse=True)
def _restore_settings():
"""job_writer uses settings.jobs_dir which is captured at module import.
Tests that set env vars then call settings.reload() — this fixture saves
and restores the original settings so one test doesn't leak into another."""
snapshot = config.Settings(
data_dir=config.settings.data_dir,
jobs_dir=config.settings.jobs_dir,
yarn_resource_manager_url=config.settings.yarn_resource_manager_url,
log_level=config.settings.log_level,
)
yield
config.settings.data_dir = snapshot.data_dir
config.settings.jobs_dir = snapshot.jobs_dir
config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
config.settings.log_level = snapshot.log_level
# --- resolve_jobs_dir (priority: arg > settings.jobs_dir) ---
def test_resolve_explicit_arg_wins(tmp_path: Path, monkeypatch): def test_resolve_explicit_arg_wins(tmp_path: Path, monkeypatch):
monkeypatch.setenv(ENV_JOBS_DIR, "/from/env") config.settings.jobs_dir = "/from/settings"
assert resolve_jobs_dir(str(tmp_path)) == str(tmp_path) assert resolve_jobs_dir(str(tmp_path)) == str(tmp_path)
def test_resolve_env_var_used_when_no_arg(monkeypatch, tmp_path: Path): def test_resolve_uses_settings_jobs_dir_when_no_arg(monkeypatch, tmp_path: Path):
monkeypatch.setenv(ENV_JOBS_DIR, str(tmp_path)) config.settings.jobs_dir = str(tmp_path)
assert resolve_jobs_dir() == str(tmp_path) assert resolve_jobs_dir() == str(tmp_path)
def test_resolve_default_when_neither_set(monkeypatch): def test_resolve_default_when_settings_unset(monkeypatch):
monkeypatch.delenv(ENV_JOBS_DIR, raising=False) config.settings.jobs_dir = "./data/jobs"
assert resolve_jobs_dir() == DEFAULT_JOBS_DIR
assert resolve_jobs_dir() == "./data/jobs" assert resolve_jobs_dir() == "./data/jobs"
@@ -48,27 +67,28 @@ def test_write_creates_jobs_dir_if_missing(tmp_path: Path):
assert os.path.isfile(out) assert os.path.isfile(out)
def test_write_uses_env_var_when_no_arg(monkeypatch, tmp_path: Path): def test_write_uses_settings_jobs_dir_when_no_arg(tmp_path: Path):
monkeypatch.setenv(ENV_JOBS_DIR, str(tmp_path)) config.settings.jobs_dir = str(tmp_path)
out = write_job_file("env-driven\n") out = write_job_file("settings-driven\n")
assert out.startswith(str(tmp_path)) assert out.startswith(str(tmp_path))
with open(out) as f: with open(out) as f:
assert f.read() == "env-driven\n" assert f.read() == "settings-driven\n"
def test_write_uses_default_when_neither_set(monkeypatch, tmp_path: Path): def test_write_uses_default_jobs_dir(tmp_path: Path):
"""Default is ./data/jobs relative to cwd. Run in a tmp dir so the test """When settings.jobs_dir is the default './data/jobs', writes resolve
doesn't pollute the real project and stays self-contained.""" relative to cwd. chdir to tmp so the test doesn't pollute the project."""
monkeypatch.delenv(ENV_JOBS_DIR, raising=False) config.settings.jobs_dir = "./data/jobs"
monkeypatch.chdir(tmp_path) # pytest built-in: chdir for this test only old_cwd = os.getcwd()
out = write_job_file("default\n") try:
assert os.path.isabs(out) os.chdir(tmp_path)
# Should resolve to <tmp_path>/data/jobs/job_*.py out = write_job_file("default\n")
expected_dir = tmp_path / "data" / "jobs" assert os.path.isabs(out)
assert expected_dir.is_dir() expected_dir = tmp_path / "data" / "jobs"
assert str(out).startswith(str(expected_dir)) assert expected_dir.is_dir()
with open(out) as f: assert str(out).startswith(str(expected_dir))
assert f.read() == "default\n" finally:
os.chdir(old_cwd)
def test_write_returns_unique_paths_for_concurrent_calls(tmp_path: Path): def test_write_returns_unique_paths_for_concurrent_calls(tmp_path: Path):
+25 -7
View File
@@ -4,6 +4,7 @@ from unittest.mock import patch
import httpx import httpx
import pytest import pytest
from common import config
from spark_executor.core.yarn_client import ( from spark_executor.core.yarn_client import (
YarnConfigError, YarnConfigError,
YarnError, YarnError,
@@ -16,6 +17,23 @@ from spark_executor.core.yarn_client import (
RM = "http://rm:8088" RM = "http://rm:8088"
@pytest.fixture(autouse=True)
def _restore_settings():
"""Each test gets a clean copy of `settings` so env-var-style overrides
in one test don't leak into the next."""
snapshot = config.Settings(
data_dir=config.settings.data_dir,
jobs_dir=config.settings.jobs_dir,
yarn_resource_manager_url=config.settings.yarn_resource_manager_url,
log_level=config.settings.log_level,
)
yield
config.settings.data_dir = snapshot.data_dir
config.settings.jobs_dir = snapshot.jobs_dir
config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
config.settings.log_level = snapshot.log_level
def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response: def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response:
if json_data is not None: if json_data is not None:
return httpx.Response(status, json=json_data) return httpx.Response(status, json=json_data)
@@ -57,16 +75,16 @@ def test_status_raises_when_state_field_missing():
def test_status_requires_rm_url(): def test_status_requires_rm_url():
with patch.dict("os.environ", {}, clear=True): config.settings.yarn_resource_manager_url = None
with pytest.raises(YarnConfigError): with pytest.raises(YarnConfigError):
get_application_status("application_1", None) get_application_status("application_1", None)
def test_status_falls_back_to_env_var(): def test_status_falls_back_to_settings_yarn_rm_url():
config.settings.yarn_resource_manager_url = "http://env-rm:8088"
fake = _resp(200, json_data={"app": {"state": "FINISHED"}}) fake = _resp(200, json_data={"app": {"state": "FINISHED"}})
with patch.dict("os.environ", {"YARN_RESOURCE_MANAGER_URL": "http://env-rm:8088"}, clear=False): with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: state, _ = get_application_status("application_1", None)
state, _ = get_application_status("application_1", None)
assert state == "FINISHED" assert state == "FINISHED"
assert "env-rm:8088" in m.call_args.args[1] assert "env-rm:8088" in m.call_args.args[1]