diff --git a/.gitignore b/.gitignore index b5279afd86..8ce5bc503d 100644 --- a/.gitignore +++ b/.gitignore @@ -22,7 +22,7 @@ venv/ .venv/ env/ .env/ -smartsim/_core/config/.env +**/.env # written upon install smartsim/version.py diff --git a/smartsim/_core/_cli/scripts/dragon_install.py b/smartsim/_core/_cli/scripts/dragon_install.py index 532939dd68..676ad71412 100644 --- a/smartsim/_core/_cli/scripts/dragon_install.py +++ b/smartsim/_core/_cli/scripts/dragon_install.py @@ -9,6 +9,7 @@ from smartsim._core._cli.utils import pip from smartsim._core._install.builder import WebTGZ +from smartsim._core.config import CONFIG from smartsim._core.utils.helpers import expand_exe_path from smartsim.error.errors import SmartSimCLIActionCancelled from smartsim.log import get_logger @@ -16,16 +17,8 @@ logger = get_logger(__name__) -def config_dir() -> pathlib.Path: - """Return a directory path where SmartSim configuration files can be written - - :returns: path to the configuration directory""" - return pathlib.Path(__file__).parents[2] / "config" - - def create_dotenv(dragon_root_dir: pathlib.Path) -> None: """Create a .env file with required environment variables for the Dragon runtime""" - dotenv_path = config_dir() / ".env" dragon_vars = { "DRAGON_BASE_DIR": str(dragon_root_dir), "DRAGON_ROOT_DIR": str(dragon_root_dir), # note: same as base_dir @@ -36,7 +29,10 @@ def create_dotenv(dragon_root_dir: pathlib.Path) -> None: lines = [f"{k}={v}\n" for k, v in dragon_vars.items()] - with dotenv_path.open("w", encoding="utf-8") as dotenv: + if not CONFIG.dragon_dotenv.parent.exists(): + CONFIG.dragon_dotenv.parent.mkdir(parents=True) + + with CONFIG.dragon_dotenv.open("w", encoding="utf-8") as dotenv: dotenv.writelines(lines) diff --git a/smartsim/_core/config/config.py b/smartsim/_core/config/config.py index fd4821d503..d4543c3462 100644 --- a/smartsim/_core/config/config.py +++ b/smartsim/_core/config/config.py @@ -100,6 +100,7 @@ def __init__(self) -> None: self.lib_path = Path(dependency_path, "lib").resolve() self.bin_path = Path(dependency_path, "bin").resolve() self.conf_path = Path(dependency_path, "config", "redis.conf") + self.conf_dir = Path(self.core_path, "config") @property def redisai(self) -> str: @@ -153,6 +154,11 @@ def database_file_parse_trials(self) -> int: def database_file_parse_interval(self) -> int: return int(os.getenv("SMARTSIM_DB_FILE_PARSE_INTERVAL", "2")) + @property + def dragon_dotenv(self) -> Path: + """Returns the path to a .env file containing dragon environment variables""" + return self.conf_dir / "dragon" / ".env" + @property def dragon_server_path(self) -> t.Optional[str]: return os.getenv("SMARTSIM_DRAGON_SERVER_PATH") diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index ad46996208..dd5f777031 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -145,14 +145,13 @@ def _load_persisted_env(self) -> t.Dict[str, str]: # use previously loaded env vars. return self._env_vars - dotenv_path = Path(__file__).parents[2] / ".env" - if not dotenv_path.exists(): + if not CONFIG.dragon_dotenv.exists(): self._env_vars = {} return self._env_vars - with open(dotenv_path, encoding="utf-8") as dot_env: + with open(CONFIG.dragon_dotenv, encoding="utf-8") as dot_env: for kvp in dot_env.readlines(): - split = kvp.split("=", maxsplit=1) + split = kvp.strip().split("=", maxsplit=1) key, value = split[0], split[-1] self._env_vars[key] = value diff --git a/tests/test_dragon_installer.py b/tests/test_dragon_installer.py index 37bc703eee..e67690f3f9 100644 --- a/tests/test_dragon_installer.py +++ b/tests/test_dragon_installer.py @@ -38,6 +38,7 @@ from smartsim._core._cli.scripts.dragon_install import ( check_for_utility, cleanup, + create_dotenv, install_dragon, install_package, is_crayex_platform, @@ -372,3 +373,100 @@ def test_install_macos(monkeypatch: pytest.MonkeyPatch, extraction_dir: pathlib. result = install_dragon(extraction_dir) assert result == 1 + + +def test_create_dotenv(monkeypatch: pytest.MonkeyPatch, test_dir: str): + """Verify that attempting to create a .env file without any existing + file or container directory works""" + test_path = pathlib.Path(test_dir) + mock_dragon_root = pathlib.Path(test_dir) / "dragon" + exp_env_path = pathlib.Path(test_dir) / "dragon" / ".env" + + with monkeypatch.context() as ctx: + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + + # ensure no .env exists before trying to create it. + assert not exp_env_path.exists() + + create_dotenv(mock_dragon_root) + + # ensure the .env is created as side-effect of create_dotenv + assert exp_env_path.exists() + + +def test_create_dotenv_existing_dir(monkeypatch: pytest.MonkeyPatch, test_dir: str): + """Verify that attempting to create a .env file in an existing + target dir works""" + test_path = pathlib.Path(test_dir) + mock_dragon_root = pathlib.Path(test_dir) / "dragon" + exp_env_path = pathlib.Path(test_dir) / "dragon" / ".env" + + # set up the parent directory that will contain the .env + exp_env_path.parent.mkdir(parents=True) + + with monkeypatch.context() as ctx: + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + + # ensure no .env exists before trying to create it. + assert not exp_env_path.exists() + + create_dotenv(mock_dragon_root) + + # ensure the .env is created as side-effect of create_dotenv + assert exp_env_path.exists() + + +def test_create_dotenv_existing_dotenv(monkeypatch: pytest.MonkeyPatch, test_dir: str): + """Verify that attempting to create a .env file when one exists works as expected""" + test_path = pathlib.Path(test_dir) + mock_dragon_root = pathlib.Path(test_dir) / "dragon" + exp_env_path = pathlib.Path(test_dir) / "dragon" / ".env" + + # set up the parent directory that will contain the .env + exp_env_path.parent.mkdir(parents=True) + + # write something into file to verify it is overwritten + var_name = "DRAGON_BASE_DIR" + exp_env_path.write_text(f"{var_name}=/foo/bar") + + with monkeypatch.context() as ctx: + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + + # ensure .env exists so we can update it + assert exp_env_path.exists() + + create_dotenv(mock_dragon_root) + + # ensure the .env is created as side-effect of create_dotenv + assert exp_env_path.exists() + + # ensure file was overwritten and env vars are not duplicated + dotenv_content = exp_env_path.read_text(encoding="utf-8") + split_content = dotenv_content.split(var_name) + + # split to confirm env var only appars once + assert len(split_content) == 2 + + +def test_create_dotenv_format(monkeypatch: pytest.MonkeyPatch, test_dir: str): + """Verify that created .env files are correctly formatted""" + test_path = pathlib.Path(test_dir) + mock_dragon_root = pathlib.Path(test_dir) / "dragon" + exp_env_path = pathlib.Path(test_dir) / "dragon" / ".env" + + with monkeypatch.context() as ctx: + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + + create_dotenv(mock_dragon_root) + + # ensure the .env is created as side-effect of create_dotenv + content = exp_env_path.read_text(encoding="utf-8") + + # ensure we have values written, but ignore empty lines + lines = [line for line in content.split("\n") if line] + assert lines + + # ensure each line is formatted as key=value + for line in lines: + line_split = line.split("=") + assert len(line_split) == 2 diff --git a/tests/test_dragon_launcher.py b/tests/test_dragon_launcher.py index 22191351ae..5fca317752 100644 --- a/tests/test_dragon_launcher.py +++ b/tests/test_dragon_launcher.py @@ -27,12 +27,15 @@ import logging import multiprocessing as mp import os +import pathlib import sys import typing as t import pytest import zmq +import smartsim._core.config +from smartsim._core._cli.scripts.dragon_install import create_dotenv from smartsim._core.config.config import get_config from smartsim._core.launcher.dragon.dragonLauncher import DragonLauncher from smartsim._core.launcher.dragon.dragonSockets import ( @@ -52,7 +55,13 @@ class MockPopen: - def __init__(self, *args: t.Any, **kwargs: t.Any) -> None: ... + calls = [] + + def __init__(self, *args: t.Any, **kwargs: t.Any) -> None: + self.args = args + self.kwargs = kwargs + + MockPopen.calls.append((args, kwargs)) @property def pid(self) -> int: @@ -177,9 +186,10 @@ def mock_dragon_env(test_dir, *args, **kwargs): logger.info(f"exception occurred while configuring mock handshaker: {ex}") -def test_dragon_connect_bind_address(monkeypatch: pytest.MonkeyPatch, test_dir: str): +def test_dragon_connect_attributes(monkeypatch: pytest.MonkeyPatch, test_dir: str): """Test the connection to a dragon environment dynamically selects an open port - in the range supplied""" + in the range supplied and passes the correct environment""" + test_path = pathlib.Path(test_dir) with monkeypatch.context() as ctx: # make sure we don't touch "real keys" during a test @@ -204,7 +214,15 @@ def test_dragon_connect_bind_address(monkeypatch: pytest.MonkeyPatch, test_dir: # avoid starting a real zmq socket ctx.setattr("zmq.Context.socket", mock_socket) # avoid starting a real process for dragon entrypoint - ctx.setattr("subprocess.Popen", lambda *args, **kwargs: MockPopen()) + ctx.setattr( + "subprocess.Popen", lambda *args, **kwargs: MockPopen(*args, **kwargs) + ) + + # avoid reading "real" config in test... + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + dotenv_path = smartsim._core.config.CONFIG.dragon_dotenv + dotenv_path.parent.mkdir(parents=True) + dotenv_path.write_text("FOO=BAR\nBAZ=BOO") dragon_launcher = DragonLauncher() dragon_launcher.connect_to_dragon(test_dir) @@ -212,6 +230,14 @@ def test_dragon_connect_bind_address(monkeypatch: pytest.MonkeyPatch, test_dir: chosen_port = int(mock_socket.bind_address.split(":")[-1]) assert chosen_port >= 5995 + # grab the kwargs env=xxx from the mocked popen to check what was passed + env = MockPopen.calls[0][1].get("env", None) + + # confirm the environment values were passed from .env file to dragon process + assert "PYTHONUNBUFFERED" in env + assert "FOO" in env + assert "BAZ" in env + @pytest.mark.parametrize( "socket_type, is_server", @@ -355,4 +381,66 @@ def fn(*args, **kwargs): launcher.connect_to_dragon(test_dir) finally: launcher.cleanup() - ... + + +def test_load_env_no_file(monkeypatch: pytest.MonkeyPatch, test_dir: str): + """Ensure an empty dragon .env file doesn't break the launcher""" + test_path = pathlib.Path(test_dir) + # mock_dragon_root = pathlib.Path(test_dir) / "dragon" + # exp_env_path = pathlib.Path(test_dir) / "dragon" / ".env" + + with monkeypatch.context() as ctx: + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + + dragon_conf = smartsim._core.config.CONFIG.dragon_dotenv + # verify config doesn't exist + assert not dragon_conf.exists() + + launcher = DragonLauncher() + + loaded_env = launcher._load_persisted_env() + assert not loaded_env + + +def test_load_env_env_file_created(monkeypatch: pytest.MonkeyPatch, test_dir: str): + """Ensure a populated dragon .env file is loaded correctly by the launcher""" + test_path = pathlib.Path(test_dir) + mock_dragon_root = pathlib.Path(test_dir) / "dragon" + + with monkeypatch.context() as ctx: + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + create_dotenv(mock_dragon_root) + dragon_conf = smartsim._core.config.CONFIG.dragon_dotenv + + # verify config does exist + assert dragon_conf.exists() + + # load config w/launcher + launcher = DragonLauncher() + loaded_env = launcher._load_persisted_env() + assert loaded_env + + # confirm .env was parsed as expected by inspecting a key + assert "DRAGON_ROOT_DIR" in loaded_env + + +def test_load_env_cached_env(monkeypatch: pytest.MonkeyPatch, test_dir: str): + """Ensure repeated attempts to use dragon env don't hit file system""" + test_path = pathlib.Path(test_dir) + mock_dragon_root = pathlib.Path(test_dir) / "dragon" + + with monkeypatch.context() as ctx: + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", test_path) + create_dotenv(mock_dragon_root) + + # load config w/launcher + launcher = DragonLauncher() + loaded_env = launcher._load_persisted_env() + assert loaded_env + + # ensure attempting to reload would bomb + ctx.setattr(smartsim._core.config.CONFIG, "conf_dir", None) + + # attempt to load and if it doesn't blow up, it used the cached copy + loaded_env = launcher._load_persisted_env() + assert loaded_env