diff --git a/api/python/yr/common/utils.py b/api/python/yr/common/utils.py index e67a6a2bfac5ba69af595c42e601228c27b3d4ca..cc507966f7313d85b10ca8f88fdcfb5ad985aaf1 100644 --- a/api/python/yr/common/utils.py +++ b/api/python/yr/common/utils.py @@ -647,3 +647,122 @@ def is_static_method(base_cls, f_name): if f_name in base.__dict__: return isinstance(base.__dict__[f_name], staticmethod) return False + + +def _should_skip_env_line(line: str) -> bool: + """Check if a line should be skipped (empty or comment). + + Args: + line: The line to check (should already be stripped). + + Returns: + True if the line should be skipped, False otherwise. + """ + return not line or line.startswith('#') + + +def _strip_quotes(value: str) -> str: + """Remove surrounding quotes from a value if present. + + Args: + value: The value string that may have quotes. + + Returns: + The value with quotes removed if they were present. + """ + if len(value) >= 2: + if value.startswith('"') and value.endswith('"'): + return value[1:-1] + + if value.startswith("'") and value.endswith("'"): + return value[1:-1] + return value + + +def _parse_env_line(line: str, line_num: int, env_file_path: str): + """Parse a single line from .env file into key-value pair. + + Args: + line: The line to parse (should already be stripped). + line_num: Line number for error reporting. + env_file_path: File path for error reporting. + + Returns: + Tuple of (key, value) if parsing succeeds, None otherwise. + """ + if '=' not in line: + log.get_logger().warning( + f"Invalid format in {env_file_path} at line {line_num}: " + f"expected KEY=VALUE format, got: {line}") + return None + + # Split on first '=' to handle values that contain '=' + parts = line.split('=', 1) + if len(parts) != 2: + log.get_logger().warning( + f"Invalid format in {env_file_path} at line {line_num}: " + f"expected KEY=VALUE format, got: {line}") + return None + + key = parts[0].strip() + value = parts[1].strip() + + # Remove quotes if present + value = _strip_quotes(value) + + # Skip if key is empty + if not key: + log.get_logger().warning( + f"Empty key in {env_file_path} at line {line_num}") + return None + + return (key, value) + + +def load_env_from_file(env_file_path: str): + """Load environment variables from a .env format file. + This must be called before any code reads from os.environ. + + The file should contain environment variables in KEY=VALUE format, one per line, e.g.: + KEY1=VALUE1 + KEY2=VALUE2 + KEY3=value with spaces + + Lines starting with # are treated as comments and ignored. + Empty lines are ignored. + Leading and trailing whitespace around KEY and VALUE are stripped. + + Args: + env_file_path (str): Path to the environment variable file (.env format). + """ + if not env_file_path or env_file_path.strip() == "": + return + + if not os.path.exists(env_file_path): + log.get_logger().warning(f"Environment variable file not found: {env_file_path}") + return + + loaded_count = 0 + try: + with open(env_file_path, 'r', encoding='utf-8') as f: + for line_num, line in enumerate(f, 1): + line = line.strip() + + if _should_skip_env_line(line): + continue + + parsed = _parse_env_line(line, line_num, env_file_path) + if parsed is None: + continue + + key, value = parsed + os.environ[key] = value + loaded_count += 1 + except Exception as e: + log.get_logger().error( + f"Failed to load environment variables from {env_file_path}: {e}") + return + + if loaded_count > 0: + log.get_logger().debug( + f"Loaded {loaded_count} environment variables from {env_file_path}") diff --git a/api/python/yr/config.py b/api/python/yr/config.py index 85d6d6400c845c5a1424f5a98934e56c56d95d1b..0198ffb270897cc02d5c9dd8c6e1d9e8e6447989 100644 --- a/api/python/yr/config.py +++ b/api/python/yr/config.py @@ -123,6 +123,13 @@ class Config: #: Log directory, specifies the path where log files will be stored. #: Default is the current working directory ("./"). log_dir: str = "./" + #: Path to environment variable file (.env format) to load at startup. + #: The file should contain environment variables in KEY=VALUE format, one per line, e.g.: + #: "KEY1=VALUE1" + #: "KEY2=VALUE2" + #: Lines starting with # are treated as comments and ignored. Empty lines are ignored. + #: If specified, environment variables from this file will be loaded into os.environ. + env_file: str = "" #: Max size for log file, default is ``0`` (If the default value is ``0``, it will eventually be set to ``40``). log_file_size_max: int = 0 #: Max number for log file, default is ``0`` (If the default value is ``0``, it will eventually be set to ``20``). diff --git a/api/python/yr/config_manager.py b/api/python/yr/config_manager.py index 186c51764d8a7d485cbbde683f28355901c39cc6..882c818d7b66caedd0374c1ffe78b7bbc56dfe7f 100644 --- a/api/python/yr/config_manager.py +++ b/api/python/yr/config_manager.py @@ -96,6 +96,7 @@ class ConfigManager: self.namespace = "" self.log_to_driver = False self.dedup_logs = False + self.env_file = "" @property def deployment_config(self) -> DeploymentConfig: @@ -298,6 +299,7 @@ class ConfigManager: self.runtime_env = conf.runtime_env self.log_to_driver = conf.log_to_driver self.dedup_logs = conf.dedup_logs + self.env_file = conf.env_file def get_function_id_by_language(self, language): """ diff --git a/api/python/yr/fnruntime.pyx b/api/python/yr/fnruntime.pyx index 4c534b10eec824c81dc24ff64fd72d548ba0ce74..160233404376242ee546580c419bab9fb01b3008 100644 --- a/api/python/yr/fnruntime.pyx +++ b/api/python/yr/fnruntime.pyx @@ -1310,6 +1310,7 @@ cdef class Fnruntime: config.ns = ConfigManager().ns config.logToDriver = ConfigManager().log_to_driver config.dedupLogs = ConfigManager().dedup_logs + config.envFile = ConfigManager().env_file for key, value in ConfigManager().custom_envs.items(): config.customEnvs.insert(pair[string, string](key, value)) with nogil: diff --git a/api/python/yr/includes/libruntime.pxd b/api/python/yr/includes/libruntime.pxd index 6a87ffdca42623813277a472fe3a3c6aa877d9dc..936f633f3527bcdcc67553ace970267cf7794147 100644 --- a/api/python/yr/includes/libruntime.pxd +++ b/api/python/yr/includes/libruntime.pxd @@ -268,6 +268,7 @@ cdef extern from "src/libruntime/libruntime_config.h" nogil: string runtimePrivateKeyPath string dsPublicKeyPath bool encryptEnable + string envFile CLibruntimeConfig() void InitConfig(const CMetaConfig & config) void BuildMetaConfig(CMetaConfig & config) diff --git a/api/python/yr/main/yr_runtime_main.py b/api/python/yr/main/yr_runtime_main.py index df5d743a3fad748cf3a879abca178acbb5298091..3121f67d2844bea1dd760b4d3652fbd5578e72ae 100644 --- a/api/python/yr/main/yr_runtime_main.py +++ b/api/python/yr/main/yr_runtime_main.py @@ -24,7 +24,7 @@ import sys from yr import init, log from yr.apis import receive_request_loop from yr.config import Config -from yr.common.utils import try_install_uvloop +from yr.common.utils import load_env_from_file, try_install_uvloop DEFAULT_LOG_DIR = "/home/snuser/log/" _ENV_KEY_FUNCTION_LIB_PATH = "YR_FUNCTION_LIB_PATH" @@ -69,14 +69,16 @@ def get_runtime_config(): return config_json -def configure(): +def configure(args: argparse.Namespace = None): """configure""" - args = parse_args() + if args is None: + args = parse_args() config = Config() config.rt_server_address = args.rt_server_address config.runtime_id = args.runtime_id config.job_id = args.job_id config.log_level = args.log_level + config.env_file = os.environ.get("YR_ENV_FILE", "") config.ds_address = os.environ.get("DATASYSTEM_ADDR") config.is_driver = False log_dir = os.getenv("GLOG_log_dir") @@ -102,9 +104,23 @@ def insert_sys_path(): def main(): """main""" + # Parse arguments + args = parse_args() + + # If YR_SEED_FILE is set, read the specified file to block + seed_file = os.environ.get("YR_SEED_FILE", "") + if seed_file: + with open(seed_file, 'rb') as f: + f.read() + + # Get env_file from environment variable + # Load environment variables from file BEFORE any code reads from os.environ + # This must be done before insert_sys_path() and configure() which read env vars + load_env_from_file(os.environ.get("YR_ENV_FILE", "")) + # If args are invalid, the script automatically exits when calling 'parser.parse_args()'. insert_sys_path() - init(configure()) + init(configure(args)) try_install_uvloop() receive_request_loop() diff --git a/api/python/yr/tests/BUILD.bazel b/api/python/yr/tests/BUILD.bazel index 0d6ce2481165bcf273b2587800c27c090c262673..a4ab9ef6faab2448461241c313bc5b5a6bb6b83e 100644 --- a/api/python/yr/tests/BUILD.bazel +++ b/api/python/yr/tests/BUILD.bazel @@ -148,3 +148,12 @@ py_test( deps = ["//api/python:yr_lib"], env = SAN_ENV, ) +py_test( + name="test_load_env_from_file", + size="small", + srcs=["test_load_env_from_file.py"], + tags = ["smoke"], + imports = ["../../"], + deps = ["//api/python:yr_lib"], + env = SAN_ENV, +) diff --git a/api/python/yr/tests/test_load_env_from_file.py b/api/python/yr/tests/test_load_env_from_file.py new file mode 100644 index 0000000000000000000000000000000000000000..68b78bc7acff68755d908626c816d3e5203a961e --- /dev/null +++ b/api/python/yr/tests/test_load_env_from_file.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# coding=UTF-8 +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import tempfile +import unittest +import logging +from unittest.mock import patch + +from yr.common import utils + +logger = logging.getLogger(__name__) + + +class TestLoadEnvFromFile(unittest.TestCase): + """Test cases for load_env_from_file function.""" + + def setUp(self): + """Set up test fixtures.""" + # Save original environment variables that might be modified + self.original_env = os.environ.copy() + # Clear test environment variables + test_keys = ["TEST_KEY1", "TEST_KEY2", "TEST_KEY3", "TEST_NUM", "TEST_BOOL", "TEST_QUOTED"] + for key in test_keys: + if key in os.environ: + del os.environ[key] + + def tearDown(self): + """Clean up after tests.""" + # Restore original environment + os.environ.clear() + os.environ.update(self.original_env) + + @patch("yr.log.get_logger") + def test_load_env_from_file(self, mock_get_logger): + mock_get_logger.return_value = logger + + """Test loading basic environment variables from .env file.""" + env_content = """ + TEST_KEY1=value1 + TEST_KEY2=value2 + TEST_KEY3=value with spaces + TEST_QUOTED="quoted value" + """ + + with tempfile.NamedTemporaryFile(mode='w', suffix='.env', delete=False) as f: + f.write(env_content) + env_file_path = f.name + + try: + utils.load_env_from_file(env_file_path) + + self.assertEqual(os.environ.get("TEST_KEY1"), "value1") + self.assertEqual(os.environ.get("TEST_KEY2"), "value2") + self.assertEqual(os.environ.get("TEST_KEY3"), "value with spaces") + self.assertEqual(os.environ.get("TEST_QUOTED"), "quoted value") + finally: + os.unlink(env_file_path) + + @patch("yr.log.get_logger") + def test_load_env_from_file_with_comments(self, mock_get_logger): + mock_get_logger.return_value = logger + + """Test loading environment variables with comments.""" + env_content = """ + # This is a comment + TEST_KEY1=value1 + # Another comment + TEST_KEY2=value2 + """ + + with tempfile.NamedTemporaryFile(mode='w', suffix='.env', delete=False) as f: + f.write(env_content) + env_file_path = f.name + + try: + utils.load_env_from_file(env_file_path) + + self.assertEqual(os.environ.get("TEST_KEY1"), "value1") + self.assertEqual(os.environ.get("TEST_KEY2"), "value2") + finally: + os.unlink(env_file_path) + + @patch("yr.log.get_logger") + def test_load_env_from_file_empty_lines(self, mock_get_logger): + mock_get_logger.return_value = logger + + """Test loading environment variables with empty lines.""" + env_content = """ + TEST_KEY1=value1 + + TEST_KEY2=value2 + """ + + with tempfile.NamedTemporaryFile(mode='w', suffix='.env', delete=False) as f: + f.write(env_content) + env_file_path = f.name + + try: + utils.load_env_from_file(env_file_path) + + self.assertEqual(os.environ.get("TEST_KEY1"), "value1") + self.assertEqual(os.environ.get("TEST_KEY2"), "value2") + finally: + os.unlink(env_file_path) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/dto/config.cpp b/src/dto/config.cpp index 065260aae7cd10c06aa738085085e92df044f853..0966331bf36c960f1eb338fe7fe2e8b4c82ff166 100644 --- a/src/dto/config.cpp +++ b/src/dto/config.cpp @@ -23,5 +23,10 @@ Config &Config::Instance() { return Config::c; } + +void Config::Reset() +{ + Config::c = Config(); // Reconstruct to reload environment variables +} } // namespace Libruntime } // namespace YR \ No newline at end of file diff --git a/src/dto/config.h b/src/dto/config.h index c5648b4869444d7760bca96eaa711eab9c553641..af3f5aedb897a7ddbf7a49b8d7b7796b30fd9fa4 100644 --- a/src/dto/config.h +++ b/src/dto/config.h @@ -67,6 +67,7 @@ Type Cast(const char *key, const std::string &value) class Config { public: static Config &Instance(); + static void Reset(); #define CONFIG_DECLARE(type, name, default) \ private: \ type name##_ = ParseFromEnv(#name, default); \ diff --git a/src/libruntime/libruntime_config.h b/src/libruntime/libruntime_config.h index 831cc8259a7bf6d51fed0dd18db429372026215b..cf0b7b096d1dd859eba140bca0fed365ca65a1ba 100644 --- a/src/libruntime/libruntime_config.h +++ b/src/libruntime/libruntime_config.h @@ -265,6 +265,7 @@ struct LibruntimeConfig { uint32_t maxConnSize = 10000; std::string nodeId; std::string nodeIp; + std::string envFile = ""; // Path to environment variable file (.env format: KEY=VALUE, one per line) private: ErrorInfo ValidateKeyParams(); }; diff --git a/src/libruntime/libruntime_manager.cpp b/src/libruntime/libruntime_manager.cpp index 9768986055eec64e4073460b60adf52289c3b870..4eeedb6208ac3b282489d5c3c1d460a60dbe55da 100644 --- a/src/libruntime/libruntime_manager.cpp +++ b/src/libruntime/libruntime_manager.cpp @@ -108,6 +108,12 @@ LibruntimeManager::LibruntimeManager() ErrorInfo LibruntimeManager::Init(const LibruntimeConfig &config, const std::string &rtCtx) { + // Load environment variables from file BEFORE any code reads from environment variables + // This must be done first to ensure all subsequent getenv() calls can read the loaded variables + if (!config.envFile.empty()) { + YR::LoadEnvFromFile(config.envFile); + } + auto err = config.Check(); if (!err.OK()) { YRLOG_ERROR("config check failed, job id is {}, err code is {}, err msg is {}", config.jobId, diff --git a/src/libruntime/utils/utils.cpp b/src/libruntime/utils/utils.cpp index 2da6bc87e0c7f56b9df5215b7ec44ae8cb17cdc6..5789a2cf5a65b67c688bf4c3cf584fca33604bc9 100644 --- a/src/libruntime/utils/utils.cpp +++ b/src/libruntime/utils/utils.cpp @@ -17,8 +17,11 @@ #include #include +#include #include +#include "src/utility/logger/logger.h" + namespace YR { const int IP_INDEX = 0; const int PORT_INDEX = 1; @@ -72,7 +75,7 @@ long long GetCurrentTimestampNs() { auto now = std::chrono::system_clock::now(); auto duration = now.time_since_epoch(); - return std::chrono::duration_cast(duration).count(); + return std::chrono::duration_cast(duration).count(); } std::string GetCurrentUTCTime() @@ -294,4 +297,111 @@ std::string GetEnvValue(const std::string &key) } return std::string(""); } + +std::string StripWhitespace(const std::string &str) +{ + size_t start = str.find_first_not_of(" \t\r\n"); + if (start == std::string::npos) { + return ""; + } + size_t end = str.find_last_not_of(" \t\r\n"); + return str.substr(start, end - start + 1); +} + +bool ShouldSkipLine(const std::string &line) +{ + return line.empty() || line[0] == '#'; +} + +bool ParseKeyValue(const std::string &line, std::string &key, std::string &value) +{ + size_t eqPos = line.find('='); + if (eqPos == std::string::npos) { + return false; + } + + key = line.substr(0, eqPos); + value = line.substr(eqPos + 1); + return true; +} + +std::string StripKeyValueWhitespace(const std::string &str) +{ + size_t start = str.find_first_not_of(" \t"); + if (start == std::string::npos) { + return ""; + } + size_t end = str.find_last_not_of(" \t"); + return str.substr(start, end - start + 1); +} + +void RemoveQuotes(std::string &value) +{ + const uint32_t minSize = 2; + if (value.length() >= minSize) { + if ((value[0] == '"' && value[value.length() - 1] == '"') || + (value[0] == '\'' && value[value.length() - 1] == '\'')) { + value = value.substr(1, value.length() - minSize); + } + } +} + +void LoadEnvFromFile(const std::string &envFile) +{ + if (envFile.empty()) { + YRLOG_WARN("Environment variable file is empty"); + return; + } + + std::ifstream file(envFile); + if (!file.is_open()) { + YRLOG_WARN("Environment variable file not found: {}", envFile); + return; + } + + try { + std::string line; + size_t lineNum = 0; + size_t loadedCount = 0; + + while (std::getline(file, line)) { + lineNum++; + + line = StripWhitespace(line); + if (ShouldSkipLine(line)) { + continue; + } + + std::string key = ""; + std::string value = ""; + if (!ParseKeyValue(line, key, value)) { + YRLOG_WARN("Invalid format in {} at line {}: expected KEY=VALUE format, got: {}", envFile, lineNum, + line); + continue; + } + + key = StripKeyValueWhitespace(key); + if (key.empty()) { + YRLOG_WARN("Empty key in {} at line {}", envFile, lineNum); + continue; + } + + value = StripKeyValueWhitespace(value); + RemoveQuotes(value); + + if (setenv(key.c_str(), value.c_str(), 1) != 0) { + YRLOG_ERROR("Failed to set environment variable: {}", key); + continue; + } + loadedCount++; + } + + if (loadedCount > 0) { + YRLOG_DEBUG("Loaded {} environment variables from {}", loadedCount, envFile); + Libruntime::Config::Reset(); + } + } catch (const std::exception &e) { + YRLOG_ERROR("Failed to load environment variables from {}: {}", envFile, e.what()); + } +} } // namespace YR diff --git a/src/libruntime/utils/utils.h b/src/libruntime/utils/utils.h index 57cf81446da1ac21c85be08898b357fef8302e8c..f064212b3db730010ae68ef21517ffd3242d0ecf 100644 --- a/src/libruntime/utils/utils.h +++ b/src/libruntime/utils/utils.h @@ -66,4 +66,11 @@ std::string GetEnvValue(const std::string &key); int32_t ToMs(int32_t timeoutS); bool WillSizeOverFlow(size_t a, size_t b); GroupPolicy ConvertStrategyToPolicy(const std::string &stategy); + +// Load environment variables from a .env format file +// This must be called before any code reads from environment variables +// The file should contain environment variables in KEY=VALUE format, one per line +// Lines starting with # are treated as comments and ignored +// Empty lines are ignored +void LoadEnvFromFile(const std::string &envFile); } // namespace YR diff --git a/test/libruntime/utils_test.cpp b/test/libruntime/utils_test.cpp index f1d3bc6b2c3230175239843b86ec8787ddbcb3f3..767fd786fe77efbe653a2f0203bb3e245c7e1977 100644 --- a/test/libruntime/utils_test.cpp +++ b/test/libruntime/utils_test.cpp @@ -16,6 +16,10 @@ #include #include +#include +#include +#include +#include #include "src/libruntime/err_type.h" #include "src/libruntime/fsclient/protobuf/common.pb.h" @@ -106,7 +110,7 @@ TEST_F(UtilsTest, SensitiveDataTest) std::string str = "data"; std::unique_ptr charArray = std::make_unique(str.size() + 1); std::copy(str.begin(), str.end(), charArray.get()); - charArray[str.size()] = '\0'; + charArray[str.size()] = '\0'; SensitiveData d1(std::move(charArray), str.size() + 1); SensitiveData d2 = SensitiveData(d1); SensitiveData d3 = d1; @@ -185,5 +189,49 @@ TEST_F(UtilsTest, unhexlifyTest) ASSERT_EQ(res, -1); } +TEST_F(UtilsTest, LoadEnvFromFile) +{ + // Create temporary .env file with various scenarios + std::string tempFile = "/tmp/test_env_" + std::to_string(getpid()) + ".env"; + std::ofstream file(tempFile); + file << "# Comment line\n"; + file << "TEST_KEY1=value1\n"; + file << "TEST_KEY2=value with spaces\n"; + file << "TEST_SINGLE_QUOTE='quoted value'\n"; + file << "TEST_DOUBLE_QUOTE=\"quoted value\"\n"; + file << "TEST_WITH_EQUALS=key=value\n"; + file << "TEST_WHITESPACE= value with spaces \n"; + file << "\n"; // Empty line + file.close(); + + // Clear test environment variables + unsetenv("TEST_KEY1"); + unsetenv("TEST_KEY2"); + unsetenv("TEST_SINGLE_QUOTE"); + unsetenv("TEST_DOUBLE_QUOTE"); + unsetenv("TEST_WITH_EQUALS"); + unsetenv("TEST_WHITESPACE"); + + { + YR::LoadEnvFromFile(tempFile); + + ASSERT_STREQ(std::getenv("TEST_KEY1"), "value1"); + ASSERT_STREQ(std::getenv("TEST_KEY2"), "value with spaces"); + ASSERT_STREQ(std::getenv("TEST_SINGLE_QUOTE"), "quoted value"); + ASSERT_STREQ(std::getenv("TEST_DOUBLE_QUOTE"), "quoted value"); + ASSERT_STREQ(std::getenv("TEST_WITH_EQUALS"), "key=value"); + ASSERT_STREQ(std::getenv("TEST_WHITESPACE"), "value with spaces"); + } + + // Cleanup + unlink(tempFile.c_str()); + unsetenv("TEST_KEY1"); + unsetenv("TEST_KEY2"); + unsetenv("TEST_SINGLE_QUOTE"); + unsetenv("TEST_DOUBLE_QUOTE"); + unsetenv("TEST_WITH_EQUALS"); + unsetenv("TEST_WHITESPACE"); +} + } // namespace test } // namespace YR \ No newline at end of file