diff --git a/README.md b/README.md index 8473c08c5b..5b495c938e 100644 --- a/README.md +++ b/README.md @@ -115,7 +115,8 @@ The sample chat app to run is found as [model-chat.py](https://github.com/micros - [Recipes](https://github.com/microsoft/olive-recipes) ## Data/Telemetry -Distributions of this project may collect usage data and send it to Microsoft to help improve our products and services. See the [privacy statement](docs/Privacy.md) for more details. + +This project may collect usage data and send it to Microsoft to help improve our products and services. See the [privacy statement](docs/Privacy.md) for details. ## 🤝 Contributions and Feedback - We welcome contributions! Please read the [contribution guidelines](./CONTRIBUTING.md) for more details on how to contribute to the Olive project. @@ -133,4 +134,3 @@ Licensed under the [MIT](./LICENSE) License. [![Build Status](https://dev.azure.com/aiinfra/PublicPackages/_apis/build/status%2FOlive%20CI?label=Olive-CI)](https://dev.azure.com/aiinfra/PublicPackages/_build/latest?definitionId=1240) [![Build Status](https://dev.azure.com/aiinfra/PublicPackages/_apis/build/status%2FOlive-ORT-Nightly?label=Olive-ORT-Nightly)](https://dev.azure.com/aiinfra/PublicPackages/_build/latest?definitionId=1279) - diff --git a/docs/Privacy.md b/docs/Privacy.md index 95aee00b0b..f934c10714 100644 --- a/docs/Privacy.md +++ b/docs/Privacy.md @@ -6,11 +6,8 @@ The software may collect information about you and your use of the software and *** ## Technical Details -Olive uses the [OpenTelemetry](https://opentelemetry.io/) API for its implementation. Telemetry is turned ON by default. Based on user consent, this data may be periodically sent to Microsoft servers following GDPR and privacy regulations for anonymity and data access controls. Application, device, and version information is collected automatically. +Telemetry is turned ON by default. Based on user consent, this data may be periodically sent to Microsoft servers following GDPR and privacy regulations for anonymity and data access controls. -In addition, Olive may collect additional telemetry data such as: -- Invoked commands -- Performance data -- Exception information +You can fully disable telemetry by adding the `--disable_telemetry` flag to any Olive CLI command, setting `OLIVE_DISABLE_TELEMETRY=1` or `ORT_DISABLE_TELEMETRY=1` before running, or calling `olive.telemetry.disable_telemetry()`. Each option suppresses every subsequent Olive telemetry event for the remainder of the process, including Olive workflow containers started by that process. When the opt-out is active before first telemetry use, Olive does not construct the telemetry singleton or create the telemetry queue, uploader, or persistent device identifier. Disabling during runtime stops this process's uploader and retains already queued unsent rows unchanged for a later telemetry-enabled process. The environment variables accept `1`, `true`, `yes`, `on`, or `y` after trimming and without regard to case. -Collection of this additional telemetry can be disabled by adding the `--disable_telemetry` flag to any Olive CLI command, or by setting the `OLIVE_DISABLE_TELEMETRY` environment variable to `1` before running. Telemetry is also automatically disabled when a CI/CD environment is detected (e.g., GitHub Actions, Azure Pipelines, Jenkins). If telemetry is enabled, but cannot be sent to Microsoft, it will be stored locally and sent when a connection is available. You can override the default cache location by setting the `OLIVE_TELEMETRY_CACHE_DIR` environment variable to a valid directory path. +In CI/CD environments (e.g., GitHub Actions, Azure Pipelines, Jenkins), Olive only emits the `OliveRecipe` event with recipe metadata. Any full opt-out takes precedence and sends nothing. diff --git a/olive/cli/api.py b/olive/cli/api.py index fbbcdcb90f..5bac80a87e 100644 --- a/olive/cli/api.py +++ b/olive/cli/api.py @@ -19,6 +19,7 @@ from olive.cli.run import WorkflowRunCommand from olive.cli.session_params_tuning import SessionParamsTuningCommand from olive.engine.output import WorkflowOutput +from olive.telemetry import disable_telemetry # pylint: disable=W0212 @@ -127,6 +128,8 @@ def _run_unified_command(command_class, **kwargs) -> Any: known_kwargs = {k: v for k, v in kwargs.items() if k in args_schema} args = _create_unified_args(command_class, known_kwargs) + if getattr(args, "disable_telemetry", False): + disable_telemetry() # Check if command handles unknown_args (like FineTuneCommand) constructor_params = inspect.signature(command_class.__init__).parameters diff --git a/olive/cli/base.py b/olive/cli/base.py index e40e5d98c4..ef19fc8968 100644 --- a/olive/cli/base.py +++ b/olive/cli/base.py @@ -338,7 +338,9 @@ def _run_workflow(self): mark_test_output_path(self.args.output_path) print("Dry run mode enabled. Configuration file is generated but no optimization is performed.") return None - workflow_output = olive_run(run_config) + workflow_output = olive_run( + run_config, recipe_telemetry_metadata=self._get_recipe_telemetry_metadata(), emit_error_telemetry=False + ) if is_test: mark_test_output_path(self.args.output_path) save_discrepancy_check_results(workflow_output, self.args.output_path) @@ -349,6 +351,17 @@ def _run_workflow(self): print(f"Model is saved at {self.args.output_path}") return workflow_output + def _get_recipe_telemetry_metadata(self) -> dict[str, str]: + recipe_name = self.__class__.__name__ + if recipe_name.endswith("Command"): + recipe_name = recipe_name[: -len("Command")] + return { + "recipe_name": recipe_name, + "recipe_command": recipe_name, + "recipe_source": "generated_cli", + "recipe_format": "generated", + } + @staticmethod def _parse_extra_options(kv_items): from onnxruntime_genai import __version__ as OrtGenaiVersion @@ -1003,7 +1016,11 @@ def add_search_options(sub_parser: ArgumentParser): def add_telemetry_options(sub_parser: ArgumentParser): """Add telemetry options to the sub_parser.""" - sub_parser.add_argument("--disable_telemetry", action="store_true", help="Disable telemetry for this command.") + sub_parser.add_argument( + "--disable_telemetry", + action="store_true", + help="Disable all Olive telemetry for this process.", + ) return sub_parser diff --git a/olive/cli/init/__init__.py b/olive/cli/init/__init__.py index 26ae8d3562..213de7aa0b 100644 --- a/olive/cli/init/__init__.py +++ b/olive/cli/init/__init__.py @@ -4,7 +4,8 @@ # -------------------------------------------------------------------------- from argparse import ArgumentParser -from olive.cli.base import BaseOliveCLICommand +from olive.cli.base import BaseOliveCLICommand, add_telemetry_options +from olive.telemetry import action class InitCommand(BaseOliveCLICommand): @@ -21,8 +22,10 @@ def register_subcommand(parser: ArgumentParser): default="./olive-output", help="Default output directory for the generated command. Default is ./olive-output.", ) + add_telemetry_options(sub_parser) sub_parser.set_defaults(func=InitCommand) + @action def run(self): from olive.cli.init.wizard import InitWizard diff --git a/olive/cli/launcher.py b/olive/cli/launcher.py index 55e6ffdeb4..369c93092a 100644 --- a/olive/cli/launcher.py +++ b/olive/cli/launcher.py @@ -24,7 +24,7 @@ from olive.cli.run_pass import RunPassCommand from olive.cli.session_params_tuning import SessionParamsTuningCommand from olive.cli.shared_cache import SharedCacheCommand -from olive.telemetry import Telemetry +from olive.telemetry import Telemetry, disable_telemetry def get_cli_parser(called_as_console_script: bool = True) -> ArgumentParser: @@ -66,18 +66,20 @@ def main(raw_args=None, called_as_console_script: bool = True): args, unknown_args = parser.parse_known_args(raw_args) - telemetry = Telemetry() - if args.disable_telemetry: - telemetry.disable_telemetry() - if not hasattr(args, "func"): parser.print_help() sys.exit(1) - # Run the command - service = args.func(parser, args, unknown_args) - service.run() - telemetry.shutdown() + if getattr(args, "disable_telemetry", False): + disable_telemetry() + telemetry = None + try: + telemetry = Telemetry.get_or_create_if_enabled() + service = args.func(parser, args, unknown_args) + service.run() + finally: + if telemetry is not None: + telemetry.shutdown() def legacy_call(deprecated_module: str, command_name: str, *args): diff --git a/olive/cli/run.py b/olive/cli/run.py index 9fe8999f0c..825f172ed7 100644 --- a/olive/cli/run.py +++ b/olive/cli/run.py @@ -17,6 +17,7 @@ mark_test_output_path, save_discrepancy_check_results, validate_test_output_path, + warn_unused_test_metrics, ) from olive.telemetry import action @@ -56,29 +57,38 @@ def register_subcommand(parser: ArgumentParser): @action def run(self): + from copy import deepcopy + from olive.common.config_utils import load_config_file from olive.workflows import run as olive_run # allow the run_config to be a dict already (for api use) - run_config = self.args.run_config - if not isinstance(run_config, dict): - run_config = load_config_file(run_config) - if "builds" in run_config and self.args.test not in (None, False): + run_config_input = self.args.run_config + run_config = ( + deepcopy(run_config_input) if isinstance(run_config_input, dict) else load_config_file(run_config_input) + ) + config_overrides = {} + test = getattr(self.args, "test", None) + test_metrics = _flatten_test_metrics(getattr(self.args, "test_metrics", None)) + test_llama_path = getattr(self.args, "test_llama_path", None) + warn_unused_test_metrics(test, test_metrics, test_llama_path) + + if "builds" in run_config and test not in (None, False): raise ValueError("--test is not supported with multi-build run configurations.") if input_model_config := get_input_model_config(self.args, required=False): print("Replacing input model config in run config") run_config["input_model"] = input_model_config - elif self.args.test not in (None, False): + config_overrides["input_model"] = input_model_config + elif test not in (None, False): input_model = run_config.get("input_model") if not isinstance(input_model, dict) or input_model.get("type", "").lower() != "hfmodel": raise ValueError("--test for olive run requires a Hugging Face input_model in the run config.") output_path = ( self.args.output_path or run_config.get("output_dir") or run_config.get("engine", {}).get("output_dir") ) - validate_test_output_path(output_path, self.args.test) - run_config["input_model"] = add_hf_test_model_config(input_model, self.args.test, output_path) - test_metrics = _flatten_test_metrics(getattr(self.args, "test_metrics", None)) - run_config = add_discrepancy_check_pass(run_config, test_metrics) + validate_test_output_path(output_path, test) + run_config["input_model"] = add_hf_test_model_config(input_model, test, output_path) + run_config = add_discrepancy_check_pass(run_config, test_metrics, test_llama_path) for arg_key, rc_key in [("output_path", "output_dir"), ("log_level", "log_severity_level")]: if (arg_value := getattr(self.args, arg_key)) is not None: @@ -87,6 +97,19 @@ def run(self): run_config.get("engine", {}).pop(rc_key, None) # add value to run config directly run_config[rc_key] = arg_value + config_overrides[rc_key] = arg_value + + recipe_telemetry_metadata = { + "recipe_command": "WorkflowRun", + "recipe_source": "config_dict" if isinstance(run_config_input, dict) else "config_file", + "recipe_format": "dict" + if isinstance(run_config_input, dict) + else Path(run_config_input).suffix.lstrip(".").lower() or "unknown", + "execution_mode": "list_required_packages" if self.args.list_required_packages else "run", + "package_config_provided": bool(self.args.package_config), + } + if config_overrides: + recipe_telemetry_metadata["config_overrides"] = config_overrides output_path = run_config.get("output_dir") or run_config.get("engine", {}).get("output_dir") workflow_output = olive_run( @@ -94,8 +117,10 @@ def run(self): list_required_packages=self.args.list_required_packages, tempdir=self.args.tempdir, package_config=self.args.package_config, + recipe_telemetry_metadata=recipe_telemetry_metadata, + emit_error_telemetry=False, ) - if self.args.test not in (None, False): + if test not in (None, False): mark_test_output_path(output_path) save_discrepancy_check_results(workflow_output, output_path) diff --git a/olive/cli/run_pass.py b/olive/cli/run_pass.py index 3ed269185f..3c7f8239ff 100644 --- a/olive/cli/run_pass.py +++ b/olive/cli/run_pass.py @@ -19,7 +19,6 @@ from olive.telemetry import action -@action class RunPassCommand(BaseOliveCLICommand): @staticmethod def register_subcommand(parser: ArgumentParser): @@ -123,6 +122,7 @@ def _get_run_config(self, tempdir: str) -> dict[str, Any]: return config + @action def run(self): # Check if user wants to list passes if hasattr(self.args, "list_passes") and self.args.list_passes: diff --git a/olive/systems/docker/docker_system.py b/olive/systems/docker/docker_system.py index 8371cccd44..1fff188c2a 100644 --- a/olive/systems/docker/docker_system.py +++ b/olive/systems/docker/docker_system.py @@ -232,6 +232,8 @@ def _prepare_run_params(self) -> dict: def _prepare_environment(self, base_env) -> dict: """Prepare environment variables for container.""" + from olive.telemetry.telemetry import Telemetry, is_ci_environment + # Convert list to dict if needed if isinstance(base_env, list): environment = {env.split("=")[0]: env.split("=")[1] for env in base_env} @@ -241,6 +243,10 @@ def _prepare_environment(self, base_env) -> dict: # Add default environment variables environment.setdefault("PYTHONPYCACHEPREFIX", "/tmp") environment["OLIVE_LOG_LEVEL"] = logging.getLevelName(logger.getEffectiveLevel()) + if is_ci_environment(): + environment["CI"] = "1" + if Telemetry.is_process_telemetry_disabled(): + environment["OLIVE_DISABLE_TELEMETRY"] = "1" # Add HuggingFace token if needed if self.hf_token: diff --git a/olive/systems/docker/workflow_runner.py b/olive/systems/docker/workflow_runner.py index 5842d0bd49..e2f1bedd47 100644 --- a/olive/systems/docker/workflow_runner.py +++ b/olive/systems/docker/workflow_runner.py @@ -5,6 +5,7 @@ from olive.common.hf.login import huggingface_login from olive.logging import get_olive_logger, set_verbosity_from_env +from olive.telemetry.telemetry import Telemetry from olive.workflows import run as olive_run logger = get_olive_logger() @@ -20,7 +21,12 @@ def runner_entry(config): config = json.load(f) logger.info("Running workflow with config: %s", config) - olive_run(config) + try: + olive_run(config, emit_error_telemetry=False, emit_recipe_telemetry=False) + finally: + telemetry = Telemetry.get_existing_instance() + if telemetry is not None: + telemetry.shutdown(flush=True) if __name__ == "__main__": diff --git a/olive/telemetry/__init__.py b/olive/telemetry/__init__.py index 0ecbbc7056..f022f6177d 100644 --- a/olive/telemetry/__init__.py +++ b/olive/telemetry/__init__.py @@ -2,7 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -from olive.telemetry.telemetry import Telemetry +from olive.telemetry.telemetry import Telemetry, disable_telemetry from olive.telemetry.telemetry_extensions import action -__all__ = ["Telemetry", "action"] +__all__ = ["Telemetry", "action", "disable_telemetry"] diff --git a/olive/telemetry/deviceid/__init__.py b/olive/telemetry/deviceid/__init__.py index 50698c12d6..b467e8973b 100644 --- a/olive/telemetry/deviceid/__init__.py +++ b/olive/telemetry/deviceid/__init__.py @@ -1,3 +1,3 @@ -from olive.telemetry.deviceid.deviceid import get_encrypted_device_id_and_status +from olive.telemetry.deviceid.deviceid import get_hashed_device_id_and_status -__all__ = ["get_encrypted_device_id_and_status"] +__all__ = ["get_hashed_device_id_and_status"] diff --git a/olive/telemetry/deviceid/_store.py b/olive/telemetry/deviceid/_store.py index 97846ed051..dd62644417 100644 --- a/olive/telemetry/deviceid/_store.py +++ b/olive/telemetry/deviceid/_store.py @@ -1,9 +1,22 @@ +import os +import stat +import tempfile +from contextlib import suppress from pathlib import Path from olive.telemetry.utils import get_telemetry_base_dir REGISTRY_PATH = r"SOFTWARE\Microsoft\DeveloperTools\.onnxruntime" REGISTRY_KEY = "deviceid" +MAX_DEVICE_ID_FILE_SIZE = 256 + + +def _chmod_best_effort(path: Path, mode: int) -> None: + try: + path.chmod(mode) + except OSError: + # Permission tightening is best-effort on filesystems that do not support chmod. + pass class Store: @@ -21,23 +34,64 @@ def retrieve_id(self) -> str: :return: The device id. :rtype: str """ - # check if file doesnt exist and raise an Exception - if not self._file_path.is_file(): - raise FileExistsError(f"File {self._file_path.stem} does not exist") - - return self._file_path.read_text(encoding="utf-8").strip() - - def store_id(self, device_id: str) -> None: + flags = os.O_RDONLY + for optional_flag in ("O_CLOEXEC", "O_NOFOLLOW", "O_NONBLOCK"): + flags |= getattr(os, optional_flag, 0) + try: + fd = os.open(self._file_path, flags) + except FileNotFoundError: + raise FileNotFoundError(f"File {self._file_path.stem} does not exist") from None + try: + if not stat.S_ISREG(os.fstat(fd).st_mode): + raise PermissionError(f"File {self._file_path.stem} is not a regular file") + chunks = [] + remaining = MAX_DEVICE_ID_FILE_SIZE + 1 + while remaining: + chunk = os.read(fd, remaining) + if not chunk: + break + chunks.append(chunk) + remaining -= len(chunk) + content = b"".join(chunks) + if len(content) > MAX_DEVICE_ID_FILE_SIZE: + raise ValueError(f"File {self._file_path.stem} is too large") + return content.decode("utf-8").strip() + except UnicodeDecodeError: + raise ValueError(f"File {self._file_path.stem} is not valid UTF-8") from None + finally: + os.close(fd) + + def store_id(self, device_id: str, replace_existing: bool = False) -> bool: """Store the device id in the store location. :param str device_id: The device id to store. :type device_id: str """ - # create the folder location if it does not exist - self._file_path.parent.mkdir(parents=True, exist_ok=True) - - self._file_path.touch() - self._file_path.write_text(device_id, encoding="utf-8") + # create the folder location if it does not exist, owner-only (0700) so other users on the + # machine cannot traverse into it to reach the device id. + self._file_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + _chmod_best_effort(self._file_path.parent, 0o700) + + fd, temp_path = tempfile.mkstemp(prefix="deviceid.tmp.", dir=self._file_path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as temp_file: + temp_file.write(device_id) + temp_file.flush() + os.fsync(temp_file.fileno()) + _chmod_best_effort(Path(temp_path), 0o600) + if replace_existing: + Path(temp_path).replace(self._file_path) + temp_path = "" + return True + try: + os.link(temp_path, self._file_path) + except FileExistsError: + return False + return True + finally: + if temp_path: + with suppress(OSError): + Path(temp_path).unlink() class WindowsStore: @@ -46,15 +100,15 @@ def retrieve_id(self) -> str: """Retrieve the device id from the Windows registry.""" import winreg - device_id: str - with winreg.OpenKeyEx( winreg.HKEY_CURRENT_USER, REGISTRY_PATH, reserved=0, access=winreg.KEY_READ | winreg.KEY_WOW64_64KEY ) as key_handle: - device_id = winreg.QueryValueEx(key_handle, REGISTRY_KEY) - return device_id[0].strip() + device_id, value_type = winreg.QueryValueEx(key_handle, REGISTRY_KEY) + if value_type != winreg.REG_SZ or not isinstance(device_id, str): + raise ValueError(f"Registry value {REGISTRY_KEY} is not a string") + return device_id.strip() - def store_id(self, device_id: str) -> None: + def store_id(self, device_id: str, replace_existing: bool = False) -> bool: """Store the device id in the windows registry. :param str device_id: The device id to store. @@ -65,6 +119,13 @@ def store_id(self, device_id: str) -> None: winreg.HKEY_CURRENT_USER, REGISTRY_PATH, reserved=0, - access=winreg.KEY_ALL_ACCESS | winreg.KEY_WOW64_64KEY, + access=winreg.KEY_QUERY_VALUE | winreg.KEY_SET_VALUE | winreg.KEY_CREATE_SUB_KEY | winreg.KEY_WOW64_64KEY, ) as key_handle: + if not replace_existing: + try: + winreg.QueryValueEx(key_handle, REGISTRY_KEY) + return False + except FileNotFoundError: + pass winreg.SetValueEx(key_handle, REGISTRY_KEY, 0, winreg.REG_SZ, device_id) + return True diff --git a/olive/telemetry/deviceid/deviceid.py b/olive/telemetry/deviceid/deviceid.py index 09087f33c3..a357ec6749 100644 --- a/olive/telemetry/deviceid/deviceid.py +++ b/olive/telemetry/deviceid/deviceid.py @@ -1,101 +1,245 @@ import hashlib +import os import platform +import threading import uuid +from contextlib import suppress from enum import Enum -from typing import Union +from typing import ClassVar from olive.telemetry.deviceid._store import Store, WindowsStore +from olive.telemetry.process_lock import ProcessDrainLock +from olive.telemetry.utils import get_telemetry_base_dir class DeviceIdStatus(Enum): - NEW = "new" - EXISTING = "existing" - CORRUPTED = "corrupted" - FAILED = "failed" + NEW = "New" + EXISTING = "Existing" + CORRUPTED = "Corrupted" + FAILED = "Failed" _device_id_state = {"device_id": None, "status": DeviceIdStatus.NEW} +_device_id_lock = threading.RLock() -def get_device_id() -> str: +def _fnv1a_hex_bytes(value: bytes) -> str: + """Hash the Windows SID only for the native-compatible mutex name.""" + hash_value = 14695981039346656037 + for byte in value: + hash_value ^= byte + hash_value = (hash_value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return f"{hash_value:016x}" + + +class _WindowsDeviceIdMutex: + """Named mutex compatible with the native device-id publication protocol.""" + + def __init__(self) -> None: + self._handle = None + self._acquired = False + self._kernel32 = None + + def acquire(self) -> bool: + try: + import ctypes + import ctypes.wintypes as wintypes + + class SidAndAttributes(ctypes.Structure): + _fields_: ClassVar = [("sid", ctypes.c_void_p), ("attributes", wintypes.DWORD)] + + class TokenUser(ctypes.Structure): + _fields_: ClassVar = [("user", SidAndAttributes)] + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + advapi32 = ctypes.WinDLL("advapi32", use_last_error=True) + kernel32.GetCurrentProcess.restype = wintypes.HANDLE + kernel32.CreateMutexW.argtypes = [ctypes.c_void_p, wintypes.BOOL, wintypes.LPCWSTR] + kernel32.CreateMutexW.restype = wintypes.HANDLE + kernel32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD] + kernel32.WaitForSingleObject.restype = wintypes.DWORD + kernel32.ReleaseMutex.argtypes = [wintypes.HANDLE] + kernel32.CloseHandle.argtypes = [wintypes.HANDLE] + advapi32.OpenProcessToken.argtypes = [ + wintypes.HANDLE, + wintypes.DWORD, + ctypes.POINTER(wintypes.HANDLE), + ] + advapi32.GetTokenInformation.argtypes = [ + wintypes.HANDLE, + ctypes.c_int, + ctypes.c_void_p, + wintypes.DWORD, + ctypes.POINTER(wintypes.DWORD), + ] + advapi32.IsValidSid.argtypes = [ctypes.c_void_p] + advapi32.GetLengthSid.argtypes = [ctypes.c_void_p] + advapi32.GetLengthSid.restype = wintypes.DWORD + + token = wintypes.HANDLE() + if not advapi32.OpenProcessToken(kernel32.GetCurrentProcess(), 0x0008, ctypes.byref(token)): + return False + try: + size = wintypes.DWORD() + advapi32.GetTokenInformation(token, 1, None, 0, ctypes.byref(size)) + if not size.value: + return False + token_info = ctypes.create_string_buffer(size.value) + if not advapi32.GetTokenInformation(token, 1, token_info, size.value, ctypes.byref(size)): + return False + sid = ctypes.cast(token_info, ctypes.POINTER(TokenUser)).contents.user.sid + if not sid or not advapi32.IsValidSid(sid): + return False + sid_size = advapi32.GetLengthSid(sid) + sid_hash = _fnv1a_hex_bytes(ctypes.string_at(sid, sid_size)) + finally: + kernel32.CloseHandle(token) + + handle = kernel32.CreateMutexW( + None, + False, + f"Global\\Microsoft.DeveloperTools.OnnxRuntime.DeviceId.{sid_hash}", + ) + if not handle: + return False + self._handle = handle + self._kernel32 = kernel32 + wait_result = kernel32.WaitForSingleObject(handle, 1000) + self._acquired = wait_result in (0x00000000, 0x00000080) + return self._acquired + except Exception: + self.release() + return False + + def release(self) -> None: + if self._handle is None or self._kernel32 is None: + return + if self._acquired: + with suppress(Exception): + self._kernel32.ReleaseMutex(self._handle) + with suppress(Exception): + self._kernel32.CloseHandle(self._handle) + self._handle = None + self._acquired = False + self._kernel32 = None + + +def _is_valid_device_id(value: str) -> bool: + if not isinstance(value, str) or len(value) != 36: + return False + hyphens = {8, 13, 18, 23} + return all( + char == "-" if index in hyphens else char.lower() in "0123456789abcdef" for index, char in enumerate(value) + ) + + +def _initialize_device_id() -> str: r"""Get the device id from the store or create one if it does not exist. An empty string is returned if an error occurs during saving or retrieval of the device id. - Linux id location: $XDG_CACHE_HOME/Microsoft/DeveloperTools/.onnxruntime/deviceid if defined + POSIX id location: $XDG_CACHE_HOME/Microsoft/DeveloperTools/.onnxruntime/deviceid if defined else $HOME/.cache/Microsoft/DeveloperTools/.onnxruntime/deviceid MacOS id location: $HOME/Library/Application Support/Microsoft/DeveloperTools/.onnxruntime/deviceid - Windows id location: HKEY_CURRENT_USER\SOFTWARE\Microsoft\.onnxruntime\deviceid + Windows id location: HKEY_CURRENT_USER\SOFTWARE\Microsoft\DeveloperTools\.onnxruntime\deviceid :return: The device id. :rtype: str """ - device_id: str = "" - store: Union[Store, WindowsStore] - create_new_id = False - - try: - if platform.system() == "Windows": - store = WindowsStore() - elif platform.system() in ("Linux", "Darwin"): + system = platform.system() + if system == "Windows": + store = WindowsStore() + elif system in ("Linux", "Darwin") or os.name == "posix": + try: store = Store() - else: - _device_id_state["status"] = DeviceIdStatus.FAILED - _device_id_state["device_id"] = device_id - return device_id - - device_id = store.retrieve_id - if len(device_id) > 256: - _device_id_state["status"] = DeviceIdStatus.CORRUPTED - _device_id_state["device_id"] = "" - create_new_id = True - else: - try: - uuid.UUID(device_id) - except ValueError: - _device_id_state["status"] = DeviceIdStatus.CORRUPTED - _device_id_state["device_id"] = "" - create_new_id = True - else: - _device_id_state["status"] = DeviceIdStatus.EXISTING - _device_id_state["device_id"] = device_id - return device_id - except (FileExistsError, FileNotFoundError): - _device_id_state["status"] = DeviceIdStatus.NEW - _device_id_state["device_id"] = "" - create_new_id = True - except (PermissionError, ValueError, NotImplementedError): - _device_id_state["status"] = DeviceIdStatus.FAILED - _device_id_state["device_id"] = device_id - return device_id - except Exception: - _device_id_state["status"] = DeviceIdStatus.FAILED - _device_id_state["device_id"] = device_id - return device_id - - if create_new_id: - device_id = str(uuid.uuid4()).lower() - + except Exception: + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + else: + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + + def read_existing() -> tuple[str, str]: try: - store.store_id(device_id) + existing = store.retrieve_id + except (FileExistsError, FileNotFoundError): + return ("missing", "") + except ValueError: + return ("invalid", "") except Exception: - _device_id_state["status"] = DeviceIdStatus.FAILED - device_id = "" - _device_id_state["device_id"] = device_id - - return device_id - - -def get_encrypted_device_id_and_status() -> tuple[str, DeviceIdStatus]: - """Generate a FIPS-compliant encrypted device ID using SHA256 and returns the deviceIdStatus. + return ("failed", "") + return ("valid", existing) if _is_valid_device_id(existing) else ("invalid", "") + + initial_state, existing = read_existing() + if initial_state == "valid": + _device_id_state.update({"status": DeviceIdStatus.EXISTING, "device_id": existing}) + return existing + if initial_state == "failed": + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + + lock = None + acquired = True + if system == "Windows": + lock = _WindowsDeviceIdMutex() + acquired = lock.acquire() + elif initial_state == "invalid": + lock = ProcessDrainLock(str(get_telemetry_base_dir() / "deviceid.lock")) + acquired = lock.acquire(1.0) - This method uses SHA256 which is FIPS 140-2 approved for cryptographic operations. - The device ID is hashed to ensure deterministic but secure device identification. + try: + if not acquired: + winner_state, winner = read_existing() + generated = winner if winner_state == "valid" else str(uuid.uuid4()).lower() + status = DeviceIdStatus.EXISTING if winner_state == "valid" else DeviceIdStatus.FAILED + _device_id_state.update({"status": status, "device_id": generated}) + return generated + + current_state, current = read_existing() + if current_state == "valid": + _device_id_state.update({"status": DeviceIdStatus.EXISTING, "device_id": current}) + return current + if current_state == "failed": + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + + corrupted = initial_state == "invalid" or current_state == "invalid" + generated = str(uuid.uuid4()).lower() + try: + stored = store.store_id(generated, replace_existing=corrupted) + except Exception: + stored = False + if stored: + status = DeviceIdStatus.CORRUPTED if corrupted else DeviceIdStatus.NEW + _device_id_state.update({"status": status, "device_id": generated}) + return generated + + winner_state, winner = read_existing() + if winner_state == "valid": + _device_id_state.update({"status": DeviceIdStatus.EXISTING, "device_id": winner}) + return winner + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + finally: + if lock is not None: + lock.release() - Returns: - str: FIPS-compliant encrypted device ID (base64-encoded) - """ - device_id = _device_id_state["device_id"] if _device_id_state["device_id"] is not None else get_device_id() - encrypted_device_id = hashlib.sha256(device_id.encode("utf-8")).digest().hex().upper() if device_id else "" - return encrypted_device_id, _device_id_state["status"] +def get_device_id() -> str: + """Get the process-cached persistent device ID, initializing it once.""" + with _device_id_lock: + if _device_id_state["device_id"] is None: + return _initialize_device_id() + return _device_id_state["device_id"] + + +def get_hashed_device_id_and_status() -> tuple[str, DeviceIdStatus]: + """Get the canonical shared SHA-256 device ID and its status.""" + with _device_id_lock: + device_id = get_device_id() + hashed = hashlib.sha256(device_id.encode("utf-8")).hexdigest() if device_id else "" + return f"c:{hashed}" if hashed else "", _device_id_state["status"] diff --git a/olive/telemetry/library/__init__.py b/olive/telemetry/library/__init__.py index 39831da66e..dcf07e101d 100644 --- a/olive/telemetry/library/__init__.py +++ b/olive/telemetry/library/__init__.py @@ -3,44 +3,14 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""OneCollector Exporter for OpenTelemetry Python. +"""OneCollector building blocks (standard library only). -This package provides an OpenTelemetry exporter that sends telemetry data -to Microsoft OneCollector using the Common Schema JSON format. - -Example usage: - - from onecollector_exporter import ( - OneCollectorLogExporter, - OneCollectorExporterOptions, - get_telemetry_logger, - ) - - # Option 1: Use with OpenTelemetry SDK directly - options = OneCollectorExporterOptions( - connection_string="InstrumentationKey=your-key-here" - ) - exporter = OneCollectorLogExporter(options=options) - - # Add to logger provider - from opentelemetry.sdk._logs import LoggerProvider - from opentelemetry.sdk._logs.export import BatchLogRecordProcessor - - provider = LoggerProvider() - provider.add_log_record_processor(BatchLogRecordProcessor(exporter)) - - # Option 2: Use the simplified telemetry logger - logger = get_telemetry_logger( - connection_string="InstrumentationKey=your-key-here" - ) - logger.log("MyEvent", {"key": "value"}) - logger.shutdown() +Helpers for serializing telemetry to Common Schema JSON and posting it to the +Microsoft OneCollector endpoint. These modules have no third-party dependency +and are driven directly by the SQLite-backed uploader. """ -from olive.telemetry.library.callback_manager import CallbackManager, PayloadTransmittedCallbackArgs from olive.telemetry.library.connection_string_parser import ConnectionStringParser -from olive.telemetry.library.event_source import OneCollectorEventId, OneCollectorEventSource, event_source -from olive.telemetry.library.exporter import OneCollectorLogExporter from olive.telemetry.library.options import ( CompressionType, OneCollectorExporterOptions, @@ -48,37 +18,16 @@ OneCollectorTransportOptions, ) from olive.telemetry.library.payload_builder import PayloadBuilder -from olive.telemetry.library.retry import RetryHandler from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper -from olive.telemetry.library.telemetry_logger import ( - TelemetryLogger, - get_telemetry_logger, - log_event, - shutdown_telemetry, -) -from olive.telemetry.library.transport import HttpJsonPostTransport, ITransport - -__version__ = "0.0.1" +from olive.telemetry.library.transport import HttpJsonPostTransport __all__ = [ - "CallbackManager", "CommonSchemaJsonSerializationHelper", "CompressionType", "ConnectionStringParser", "HttpJsonPostTransport", - "ITransport", - "OneCollectorEventId", - "OneCollectorEventSource", "OneCollectorExporterOptions", "OneCollectorExporterValidationError", - "OneCollectorLogExporter", "OneCollectorTransportOptions", "PayloadBuilder", - "PayloadTransmittedCallbackArgs", - "RetryHandler", - "TelemetryLogger", - "event_source", - "get_telemetry_logger", - "log_event", - "shutdown_telemetry", ] diff --git a/olive/telemetry/library/callback_manager.py b/olive/telemetry/library/callback_manager.py deleted file mode 100644 index ee62553163..0000000000 --- a/olive/telemetry/library/callback_manager.py +++ /dev/null @@ -1,110 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""Callback manager for payload transmission events.""" - -import threading -from dataclasses import dataclass -from typing import Callable, Optional - -from olive.telemetry.library.event_source import event_source - - -@dataclass -class PayloadTransmittedCallbackArgs: - """Arguments passed to payload transmitted callbacks.""" - - succeeded: bool - """Whether the transmission succeeded.""" - - status_code: Optional[int] - """HTTP status code, if available.""" - - payload_size_bytes: int - """Size of the transmitted payload in bytes.""" - - item_count: int - """Number of items in the payload.""" - - payload_bytes: Optional[bytes] = None - """Raw payload bytes (uncompressed), if available.""" - - -class CallbackManager: - """Manages callbacks for payload transmission events. - - Allows registration of callbacks that are invoked when payloads - are successfully transmitted or fail. - """ - - def __init__(self): - """Initialize the callback manager.""" - self._callbacks: list[tuple[Callable[[PayloadTransmittedCallbackArgs], None], bool]] = [] - self._lock = threading.Lock() - self._closed = False - - def register( - self, callback: Callable[[PayloadTransmittedCallbackArgs], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback to be invoked on payload transmission. - - Args: - callback: Function to call when payload is transmitted - include_failures: Whether to invoke callback on transmission failures - - Returns: - Function to call to unregister the callback - - """ - with self._lock: - if self._closed: - return lambda: None # No-op unregister if disposed - entry = (callback, include_failures) - self._callbacks.append(entry) - - def unregister(): - """Unregister this callback.""" - with self._lock: - try: - self._callbacks.remove(entry) - except ValueError: - # The callback was already removed. - pass - - return unregister - - def notify(self, args: PayloadTransmittedCallbackArgs) -> None: - """Notify all registered callbacks. - - Args: - args: Callback arguments - - """ - # Get snapshot of callbacks to avoid holding lock during invocation - with self._lock: - if self._closed: - return - callbacks_snapshot = self._callbacks.copy() - - # Invoke callbacks - for callback, include_failures in callbacks_snapshot: - # Check if we should invoke this callback - if not args.succeeded and not include_failures: - continue - - try: - callback(args) - except Exception as ex: - # Log but don't propagate exceptions from user code - event_source.exception_thrown_from_user_code("PayloadTransmittedCallback", ex) - - def close(self) -> None: - """Close the callback manager and prevent further registrations. - - This method is idempotent and can be called multiple times. - """ - with self._lock: - self._callbacks.clear() - self._closed = True diff --git a/olive/telemetry/library/event_source.py b/olive/telemetry/library/event_source.py deleted file mode 100644 index e65d9d546a..0000000000 --- a/olive/telemetry/library/event_source.py +++ /dev/null @@ -1,257 +0,0 @@ -"""EventSource-style logging for OneCollector exporter. - -Provides structured logging similar to .NET EventSource for diagnostics and monitoring. -""" - -import logging -from enum import IntEnum - - -class OneCollectorEventId(IntEnum): - """Event IDs matching .NET EventSource implementation.""" - - EXPORT_EXCEPTION = 1 - TRANSPORT_DATA_SENT = 2 - SINK_DATA_WRITTEN = 3 - DATA_DROPPED = 4 - TRANSPORT_EXCEPTION = 5 - HTTP_ERROR_RESPONSE = 6 - EVENT_FULL_NAME_DISCARDED = 7 - EVENT_NAMESPACE_INVALID = 8 - EVENT_NAME_INVALID = 9 - USER_CODE_EXCEPTION = 10 - ATTRIBUTE_DROPPED = 11 - - -class OneCollectorEventSource: - """EventSource for OneCollector exporter diagnostics. - - Provides structured logging matching the .NET EventSource implementation. - """ - - def __init__(self): - self.logger = logging.getLogger("OpenTelemetry.Exporter.OneCollector") - # Set default level to INFO to match .NET behavior - if not self.logger.handlers: - handler = logging.StreamHandler() - formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") - handler.setFormatter(formatter) - self.logger.addHandler(handler) - self.logger.setLevel(logging.INFO) - - @property - def is_informational_logging_enabled(self) -> bool: - """Check if informational level logging is enabled.""" - return self.logger.isEnabledFor(logging.INFO) - - @property - def is_warning_logging_enabled(self) -> bool: - """Check if warning level logging is enabled.""" - return self.logger.isEnabledFor(logging.WARNING) - - @property - def is_error_logging_enabled(self) -> bool: - """Check if error level logging is enabled.""" - return self.logger.isEnabledFor(logging.ERROR) - - def export_exception_thrown(self, item_type: str, exception: Exception) -> None: - """Log an exception thrown during export. - - Args: - item_type: Type of item being exported (e.g., 'LogData') - exception: The exception that was thrown - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Exception thrown exporting '%s' batch: %s", - item_type, - exception, - exc_info=exception, - extra={"event_id": OneCollectorEventId.EXPORT_EXCEPTION}, - ) - - def transport_data_sent(self, item_type: str, num_records: int, transport_description: str) -> None: - """Log successful data transmission. - - Args: - item_type: Type of items sent - num_records: Number of records sent - transport_description: Description of transport used - - """ - if self.is_informational_logging_enabled: - self.logger.info( - "Sent '%s' batch of %s item(s) to '%s' transport", - item_type, - num_records, - transport_description, - extra={"event_id": OneCollectorEventId.TRANSPORT_DATA_SENT}, - ) - - def sink_data_written(self, item_type: str, num_records: int, sink_description: str) -> None: - """Log data written to sink. - - Args: - item_type: Type of items written - num_records: Number of records written - sink_description: Description of sink used - - """ - if self.is_informational_logging_enabled: - self.logger.info( - "Wrote '%s' batch of %s item(s) to '%s' sink", - item_type, - num_records, - sink_description, - extra={"event_id": OneCollectorEventId.SINK_DATA_WRITTEN}, - ) - - def data_dropped( - self, item_type: str, num_records: int, during_serialization: int, during_transmission: int - ) -> None: - """Log dropped data. - - Args: - item_type: Type of items dropped - num_records: Total number of records dropped - during_serialization: Number dropped during serialization - during_transmission: Number dropped during transmission - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Dropped %s '%s' item(s). %s item(s) dropped during serialization. %s item(s) dropped due to " - "transmission failure", - num_records, - item_type, - during_serialization, - during_transmission, - extra={"event_id": OneCollectorEventId.DATA_DROPPED}, - ) - - def transport_exception_thrown(self, transport_type: str, exception: Exception) -> None: - """Log transport exception. - - Args: - transport_type: Type of transport - exception: The exception that was thrown - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Exception thrown by '%s' transport: %s", - transport_type, - exception, - exc_info=exception, - extra={"event_id": OneCollectorEventId.TRANSPORT_EXCEPTION}, - ) - - def http_transport_error_response( - self, transport_type: str, status_code: int, error_message: str, error_details: str - ) -> None: - """Log HTTP error response. - - Args: - transport_type: Type of transport - status_code: HTTP status code - error_message: Error message from response - error_details: Additional error details - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Error response received by '%s' transport. StatusCode: %s, ErrorMessage: '%s', ErrorDetails: '%s'", - transport_type, - status_code, - error_message, - error_details, - extra={"event_id": OneCollectorEventId.HTTP_ERROR_RESPONSE}, - ) - - def event_full_name_discarded(self, event_namespace: str, event_name: str) -> None: - """Log event full name discarded. - - Args: - event_namespace: Event namespace - event_name: Event name - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Event full name discarded. EventNamespace: '%s', EventName: '%s'", - event_namespace, - event_name, - extra={"event_id": OneCollectorEventId.EVENT_FULL_NAME_DISCARDED}, - ) - - def event_namespace_invalid(self, event_namespace: str) -> None: - """Log invalid event namespace. - - Args: - event_namespace: The invalid namespace - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Event namespace invalid. EventNamespace: '%s'", - event_namespace, - extra={"event_id": OneCollectorEventId.EVENT_NAMESPACE_INVALID}, - ) - - def event_name_invalid(self, event_name: str) -> None: - """Log invalid event name. - - Args: - event_name: The invalid event name - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Event name invalid. EventName: '%s'", - event_name, - extra={"event_id": OneCollectorEventId.EVENT_NAME_INVALID}, - ) - - def exception_thrown_from_user_code(self, user_code_type: str, exception: Exception) -> None: - """Log exception from user code (e.g., callbacks). - - Args: - user_code_type: Type of user code that threw exception - exception: The exception that was thrown - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Exception thrown by '%s' user code: %s", - user_code_type, - exception, - exc_info=exception, - extra={"event_id": OneCollectorEventId.USER_CODE_EXCEPTION}, - ) - - def attribute_dropped(self, item_type: str, attribute_name: str, reason: str) -> None: - """Log dropped attribute. - - Args: - item_type: Type of item - attribute_name: Name of dropped attribute - reason: Reason for dropping - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Dropped %s attribute '%s': %s", - item_type, - attribute_name, - reason, - extra={"event_id": OneCollectorEventId.ATTRIBUTE_DROPPED}, - ) - - def disable(self) -> None: - """Disable telemetry logging.""" - self.logger.disabled = True - - -# Global event source instance -event_source = OneCollectorEventSource() diff --git a/olive/telemetry/library/exporter.py b/olive/telemetry/library/exporter.py deleted file mode 100644 index 045647fe9d..0000000000 --- a/olive/telemetry/library/exporter.py +++ /dev/null @@ -1,335 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""Main OneCollector log exporter implementation.""" - -import threading -from collections.abc import Sequence -from datetime import datetime, timezone -from time import time -from typing import TYPE_CHECKING, Any, Callable, Optional - -import requests -from opentelemetry.sdk._logs import ReadableLogRecord -from opentelemetry.sdk._logs.export import LogExportResult, LogRecordExporter -from opentelemetry.sdk.resources import Resource - -from olive.telemetry.library.callback_manager import CallbackManager -from olive.telemetry.library.event_source import event_source -from olive.telemetry.library.options import OneCollectorExporterOptions -from olive.telemetry.library.payload_builder import PayloadBuilder -from olive.telemetry.library.retry import RetryHandler -from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper -from olive.telemetry.library.transport import HttpJsonPostTransport - -if TYPE_CHECKING: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - - -class OneCollectorLogExporter(LogRecordExporter): - """OpenTelemetry log exporter for Microsoft OneCollector. - - Implements the OpenTelemetry LogRecordExporter interface and sends logs - to OneCollector using the Common Schema JSON format. - """ - - def __init__( - self, - options: Optional[OneCollectorExporterOptions] = None, - excluded_attributes: Optional[set[str]] = None, - ): - """Initialize the OneCollector log exporter. - - Args: - options: Exporter configuration options - excluded_attributes: Attribute keys to exclude from log attributes - - """ - # Validate options - if options is None: - raise ValueError("OneCollectorExporterOptions is required") - options.validate() - - self._options = options - self._shutdown_lock = threading.Lock() - self._shutdown = False - self._shutdown_event = threading.Event() - if excluded_attributes is None: - self._excluded_attributes = { - "code.filepath", - "code.function", - "code.lineno", - "code.file.path", - "code.function.name", - "code.line.number", - } - else: - self._excluded_attributes = set(excluded_attributes) - - # Initialize transport - transport_opts = options.transport_options - - # Create or get HTTP session - if transport_opts.http_client_factory: - self._session = transport_opts.http_client_factory() - self._owns_session = False - else: - self._session = requests.Session() - self._owns_session = True - - try: - # Build iKey with tenant prefix - self._ikey = f"{CommonSchemaJsonSerializationHelper.ONE_COLLECTOR_TENANCY_SYMBOL}:{options.tenant_token}" - - # Initialize callback manager - self._callback_manager = CallbackManager() - - # Initialize transport with callback manager - self._transport = HttpJsonPostTransport( - endpoint=transport_opts.endpoint, - ikey=options.instrumentation_key, - compression=transport_opts.compression, - session=self._session, - callback_manager=self._callback_manager, - ) - - # Initialize payload builder - self._payload_builder = PayloadBuilder( - max_size_bytes=transport_opts.max_payload_size_bytes, max_items=transport_opts.max_items_per_payload - ) - - # Initialize retry handler - self._retry_handler = RetryHandler(max_retries=6) - - # Initialize metadata - self._metadata: dict[str, Any] = {} - - # Cache for resource (populated on first export) - self._resource: Optional[Resource] = None - except Exception: - if self._owns_session: - self._session.close() - raise - - def add_metadata(self, metadata: dict[str, Any]) -> None: - """Add custom metadata fields to all exported logs. - - Args: - metadata: Dictionary of metadata fields to add - - """ - self._metadata.update(metadata) - - def register_payload_transmitted_callback( - self, callback: Callable[["PayloadTransmittedCallbackArgs"], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback that will be invoked on payload transmission. - - Callbacks are invoked after each HTTP request completes. If retries are - enabled, callbacks will be invoked for each retry attempt. - - Args: - callback: Function to call when payload is transmitted. - Receives PayloadTransmittedCallbackArgs with transmission details. - include_failures: If True, callback is invoked on both success and failure. - If False, callback is only invoked on success. - - Returns: - Function to call to unregister the callback. - - Example: - >>> def on_transmitted(args): - ... if args.succeeded: - ... print(f"✅ Sent {args.item_count} items ({args.payload_size_bytes} bytes)") - ... else: - ... print(f"❌ Failed: status={args.status_code}") - >>> - >>> unregister = exporter.register_payload_transmitted_callback( - ... on_transmitted, - ... include_failures=True - ... ) - >>> # Later: unregister() - - """ - return self._transport.register_payload_transmitted_callback(callback, include_failures) - - def export(self, batch: Sequence[ReadableLogRecord]) -> LogExportResult: - """Export a batch of log records. - - Args: - batch: Sequence of log data records to export - - Returns: - LogExportResult indicating success or failure - - """ - if self._shutdown: - return LogExportResult.FAILURE - - try: - # Get resource (cache for subsequent calls) - if self._resource is None: - first_item = batch[0] if batch else None - resource = getattr(first_item, "resource", None) - if resource is None and first_item is not None: - resource = getattr(first_item.log_record, "resource", None) - self._resource = resource or Resource.create() - - # Serialize log records to JSON - serialized_items = [] - for log_data in batch: - try: - item_bytes = self._serialize_log_data(log_data) - serialized_items.append(item_bytes) - except Exception as ex: - event_source.export_exception_thrown("ReadableLogRecord", ex) - # Continue with other items - - if not serialized_items: - return LogExportResult.FAILURE - - # Build payloads respecting size/count limits - payloads = self._build_payloads(serialized_items) - - # Send each payload with retry logic - deadline_sec = time() + self._options.transport_options.timeout_seconds - - for payload in payloads: - # Count items in this payload (approximation based on newlines) - item_count = payload.count(b"\n") + 1 if payload else 0 - success = self._retry_handler.execute_with_retry( - operation=lambda payload=payload, item_count=item_count: self._transport.send( - payload, max(0.1, deadline_sec - time()), item_count=item_count - ), - deadline_sec=deadline_sec, - shutdown_event=self._shutdown_event, - ) - - if not success: - return LogExportResult.FAILURE - - # Check if shutdown occurred - if self._shutdown: - return LogExportResult.FAILURE - - # Log success - event_source.sink_data_written("ReadableLogRecord", len(batch), "OneCollector") - - return LogExportResult.SUCCESS - - except Exception as ex: - event_source.export_exception_thrown("ReadableLogRecord", ex) - return LogExportResult.FAILURE - - def _serialize_log_data(self, log_data: ReadableLogRecord) -> bytes: - """Serialize a single log record to JSON bytes. - - Args: - log_data: Log data to serialize - - Returns: - UTF-8 encoded JSON bytes - - """ - log_record = log_data.log_record - - # Build data dictionary - data = {} - - # Add resource attributes (if available) - if self._resource and self._resource.attributes: - for key, value in self._resource.attributes.items(): - # Map common resource attributes - if key == "service.name" and "app_name" not in data: - data["app_name"] = value - elif key == "service.version" and "app_version" not in data: - data["app_version"] = value - elif key == "service.instance.id" and "app_instance_id" not in data: - data["app_instance_id"] = value - else: - data[key] = value - - # Add log record attributes (override resource attributes) - if log_record.attributes: - data.update( - {key: value for key, value in log_record.attributes.items() if key not in self._excluded_attributes} - ) - - # Add custom metadata - data.update(self._metadata) - - # Format timestamp - if log_record.timestamp: - timestamp = datetime.fromtimestamp(log_record.timestamp / 1e9, tz=timezone.utc) - else: - timestamp = datetime.now(timezone.utc) - - # Create event envelope - event_name = str(log_record.body) if log_record.body else "UnnamedEvent" - - envelope = CommonSchemaJsonSerializationHelper.create_event_envelope( - event_name=event_name, timestamp=timestamp, ikey=self._ikey, data=data - ) - - # Serialize to JSON bytes - return CommonSchemaJsonSerializationHelper.serialize_to_json_bytes(envelope) - - def _build_payloads(self, serialized_items: list[bytes]) -> list[bytes]: - """Build payloads from serialized items respecting size and count limits. - - Args: - serialized_items: List of serialized item bytes - - Returns: - List of payload bytes - - """ - payloads = [] - self._payload_builder.reset() - - for item_bytes in serialized_items: - if not self._payload_builder.can_add(item_bytes) and not self._payload_builder.is_empty: - # Current payload is full, build it and start a new one - payloads.append(self._payload_builder.build()) - self._payload_builder.reset() - - self._payload_builder.add(item_bytes) - - # Build final payload - if not self._payload_builder.is_empty: - payloads.append(self._payload_builder.build()) - - return payloads - - def force_flush(self, timeout_millis: float = 10_000) -> bool: - """Force flush any buffered data. - - Note: This exporter doesn't buffer data internally, so this is a no-op. - - Args: - timeout_millis: Timeout in milliseconds - - Returns: - True (always succeeds) - - """ - return True - - def shutdown(self) -> None: - """Shutdown the exporter and release resources.""" - with self._shutdown_lock: - if self._shutdown: - return - - self._shutdown = True - self._shutdown_event.set() - - # Close HTTP session (only if we own it) - if hasattr(self, "_session") and getattr(self, "_owns_session", True): - self._session.close() - - # Close callback manager - if hasattr(self, "_callback_manager"): - self._callback_manager.close() diff --git a/olive/telemetry/library/options.py b/olive/telemetry/library/options.py index dd934cad2d..7367c0a062 100644 --- a/olive/telemetry/library/options.py +++ b/olive/telemetry/library/options.py @@ -7,9 +7,7 @@ from dataclasses import dataclass, field from enum import Enum -from typing import Callable, Optional - -import requests +from typing import Optional from olive.telemetry.library.connection_string_parser import ConnectionStringParser @@ -35,7 +33,6 @@ class OneCollectorTransportOptions: max_items_per_payload: int = DEFAULT_MAX_ITEMS_PER_PAYLOAD compression: CompressionType = CompressionType.DEFLATE timeout_seconds: float = 10.0 - http_client_factory: Optional[Callable[[], requests.Session]] = None def validate(self) -> None: """Validate the transport options. @@ -62,6 +59,7 @@ class OneCollectorExporterOptions: """Configuration options for OneCollector exporter.""" connection_string: Optional[str] = None + service_name: Optional[str] = None transport_options: OneCollectorTransportOptions = field(default_factory=OneCollectorTransportOptions) # Internal fields populated during validation diff --git a/olive/telemetry/library/retry.py b/olive/telemetry/library/retry.py deleted file mode 100644 index 9f0cc7cfd8..0000000000 --- a/olive/telemetry/library/retry.py +++ /dev/null @@ -1,98 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""Retry logic with exponential backoff for OneCollector exporter.""" - -import random -import threading -from time import time -from typing import Callable, Optional - -from olive.telemetry.library.event_source import event_source -from olive.telemetry.library.transport import HttpJsonPostTransport - - -class RetryHandler: - """Handles retry logic with exponential backoff and jitter. - - Implements retry strategy matching the .NET implementation. - """ - - def __init__(self, max_retries: int = 6, base_delay: float = 1.0, max_delay: float = 60.0): - """Initialize retry handler. - - Args: - max_retries: Maximum number of retry attempts - base_delay: Base delay for exponential backoff (seconds) - max_delay: Maximum delay between retries (seconds) - - """ - self.max_retries = max_retries - self.base_delay = base_delay - self.max_delay = max_delay - - def execute_with_retry( - self, - operation: Callable[[], tuple[bool, Optional[int]]], - deadline_sec: float, - shutdown_event: threading.Event, - ) -> bool: - """Execute an operation with retry logic. - - Args: - operation: Function that returns (success, status_code) - deadline_sec: Absolute deadline timestamp - shutdown_event: Event to signal shutdown - - Returns: - True if operation succeeded, False otherwise - - """ - for retry_num in range(self.max_retries): - # Check if we've exceeded the deadline - remaining_time = deadline_sec - time() - if remaining_time <= 0: - return False - - try: - # Execute the operation - success, status_code = operation() - - if success: - return True - - # Check if response is retryable - if not HttpJsonPostTransport.is_retryable(status_code): - return False - - except Exception as ex: - event_source.export_exception_thrown("RetryHandler", ex) - - # Last retry - don't wait - if retry_num + 1 == self.max_retries: - return False - - # Last retry - failed - if retry_num + 1 == self.max_retries: - return False - - # Calculate backoff with exponential increase and jitter - backoff = min(self.base_delay * (2**retry_num), self.max_delay) - # Add +/-20% jitter - backoff *= random.uniform(0.8, 1.2) - - # Don't wait longer than remaining time - remaining_time = deadline_sec - time() - wait_time = min(backoff, remaining_time) - - if wait_time <= 0: - return False - - # Wait with ability to interrupt on shutdown - if shutdown_event.wait(wait_time): - # Shutdown occurred - return False - - return False diff --git a/olive/telemetry/library/serialization.py b/olive/telemetry/library/serialization.py index 069f85d7e1..8bc0a1a6bd 100644 --- a/olive/telemetry/library/serialization.py +++ b/olive/telemetry/library/serialization.py @@ -7,10 +7,13 @@ import base64 import json +import math from datetime import date, datetime, time, timedelta, timezone from typing import Any from uuid import UUID +from olive.telemetry.telemetry_redaction import scrub_string_for_telemetry + class CommonSchemaJsonSerializationHelper: """Helper class for serializing values to Common Schema JSON format. @@ -41,7 +44,11 @@ def serialize_value(value: Any) -> Any: return value # Numeric types - if isinstance(value, (int, float)): + if isinstance(value, int): + return value + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError("Telemetry numeric values must be finite") return value # String @@ -80,22 +87,30 @@ def serialize_value(value: Any) -> Any: return base64.b64encode(bytes(value)).decode("ascii") # Arrays/Lists - if isinstance(value, (list, tuple)): - return [CommonSchemaJsonSerializationHelper.serialize_value(item) for item in value] + if isinstance(value, (list, tuple, set, frozenset)): + items = [CommonSchemaJsonSerializationHelper.serialize_value(item) for item in value] + if isinstance(value, (set, frozenset)): + items.sort(key=lambda item: (type(item).__name__, repr(item))) + return items # Dictionary/Map if isinstance(value, dict): result = {} - for k, v in value.items(): - if k: # Skip empty keys - result[str(k)] = CommonSchemaJsonSerializationHelper.serialize_value(v) - return result + collisions = set() + for k, v in sorted(value.items(), key=lambda item: str(item[0])): + key = scrub_string_for_telemetry(str(CommonSchemaJsonSerializationHelper.serialize_value(k))) + if key: + if key in result: + collisions.add(key) + else: + result[key] = CommonSchemaJsonSerializationHelper.serialize_value(v) + return {key: value for key, value in result.items() if key not in collisions} # Default: convert to string try: - return str(value) + return scrub_string_for_telemetry(str(value)) except Exception: - return f"ERROR: type {type(value).__name__} is not supported" + return f"[unsupported:{type(value).__name__}]" @staticmethod def create_event_envelope( @@ -138,4 +153,6 @@ def serialize_to_json_bytes(envelope: dict[str, Any]) -> bytes: UTF-8 encoded JSON bytes """ - return json.dumps(envelope, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + return json.dumps(envelope, ensure_ascii=False, separators=(",", ":"), allow_nan=False, sort_keys=True).encode( + "utf-8" + ) diff --git a/olive/telemetry/library/telemetry_logger.py b/olive/telemetry/library/telemetry_logger.py deleted file mode 100644 index 7eb236e759..0000000000 --- a/olive/telemetry/library/telemetry_logger.py +++ /dev/null @@ -1,197 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""High-level telemetry logger facade for easy usage.""" - -import logging -import uuid -from typing import Any, Callable, Optional - -from opentelemetry._logs import set_logger_provider -from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler -from opentelemetry.sdk._logs.export import BatchLogRecordProcessor -from opentelemetry.sdk.resources import Resource - -from olive.telemetry.library.exporter import OneCollectorLogExporter -from olive.telemetry.library.options import OneCollectorExporterOptions -from olive.version import __version__ as VERSION - - -class TelemetryLogger: - """Singleton telemetry logger for simplified OneCollector integration. - - Provides a simple interface for logging telemetry events without - needing to configure OpenTelemetry directly. - """ - - _instance: Optional["TelemetryLogger"] = None - _default_logger: Optional["TelemetryLogger"] = None - _logger: Optional[logging.Logger] = None - _logger_exporter: Optional[OneCollectorLogExporter] = None - _logger_provider: Optional[LoggerProvider] = None - - def __new__(cls, options: Optional[OneCollectorExporterOptions] = None): - """Create or return the singleton instance. - - Args: - options: Exporter options (only used on first instantiation) - - """ - if cls._instance is None: - cls._instance = super().__new__(cls) - cls._instance._initialize(options) - - return cls._instance - - def _initialize(self, options: Optional[OneCollectorExporterOptions]) -> None: - """Initialize the logger (called only once). - - Args: - options: Exporter configuration options - - """ - try: - # Create exporter - self._logger_exporter = OneCollectorLogExporter(options=options) - - # Create logger provider - self._logger_provider = LoggerProvider( - resource=Resource.create( - { - "service.name": __name__.split(".", maxsplit=1)[0], - "service.version": VERSION, - "service.instance.id": str(uuid.uuid4()), # Unique instance ID; can double as session ID - } - ) - ) - - # Set as global logger provider - set_logger_provider(self._logger_provider) - - # Add batch processor - self._logger_provider.add_log_record_processor( - BatchLogRecordProcessor( - self._logger_exporter, - schedule_delay_millis=1000, - ) - ) - - # Create logging handler - handler = LoggingHandler(level=logging.INFO, logger_provider=self._logger_provider) - - # Set up Python logger - logger = logging.getLogger(__name__) - logger.propagate = False - logger.setLevel(logging.INFO) - logger.addHandler(handler) - - self._logger = logger - - except Exception: - # Silently fail initialization - logger will be None - self._logger = None - self._logger_provider = None - self._logger_exporter = None - - def add_global_metadata(self, metadata: dict[str, Any]) -> None: - """Add metadata fields to all telemetry events. - - Args: - metadata: Dictionary of metadata to add - - """ - if self._logger_exporter: - self._logger_exporter.add_metadata(metadata) - - def register_payload_transmitted_callback( - self, callback, include_failures: bool = False - ) -> Optional[Callable[[], None]]: - """Register a callback for payload transmission events.""" - if self._logger_exporter: - return self._logger_exporter.register_payload_transmitted_callback(callback, include_failures) - return None - - def log(self, event_name: str, attributes: Optional[dict[str, Any]] = None) -> None: - """Log a telemetry event. - - Args: - event_name: Name of the event - attributes: Optional event attributes - - """ - if self._logger: - extra = attributes if attributes else {} - self._logger.info(event_name, extra=extra) - - def disable_telemetry(self) -> None: - """Disable telemetry logging.""" - if self._logger: - self._logger.disabled = True - - def enable_telemetry(self) -> None: - """Enable telemetry logging.""" - if self._logger: - self._logger.disabled = False - - def shutdown(self) -> None: - """Shutdown the telemetry logger and flush pending data.""" - if self._logger_provider: - self._logger_provider.shutdown() - - @classmethod - def get_default_logger(cls, connection_string: Optional[str] = None) -> "TelemetryLogger": - """Get or create the default telemetry logger. - - Args: - connection_string: OneCollector connection string (only used on first call) - - Returns: - TelemetryLogger instance - - """ - if cls._default_logger is None: - options = None - if connection_string: - options = OneCollectorExporterOptions(connection_string=connection_string) - cls._default_logger = cls(options=options) - - return cls._default_logger - - @classmethod - def shutdown_default_logger(cls) -> None: - """Shutdown the default telemetry logger.""" - if cls._default_logger: - cls._default_logger.shutdown() - cls._default_logger = None - - -def get_telemetry_logger(connection_string: Optional[str] = None) -> TelemetryLogger: - """Get or create the default telemetry logger. - - Args: - connection_string: OneCollector connection string (only used on first call) - - Returns: - TelemetryLogger instance - - """ - return TelemetryLogger.get_default_logger(connection_string=connection_string) - - -def log_event(event_name: str, attributes: Optional[dict[str, Any]] = None) -> None: - """Log a telemetry event using the default logger. - - Args: - event_name: Name of the event - attributes: Optional event attributes - - """ - logger = get_telemetry_logger() - logger.log(event_name, attributes) - - -def shutdown_telemetry() -> None: - """Shutdown the default telemetry logger.""" - TelemetryLogger.shutdown_default_logger() diff --git a/olive/telemetry/library/transport.py b/olive/telemetry/library/transport.py index 93db8ab27f..cc679f4473 100644 --- a/olive/telemetry/library/transport.py +++ b/olive/telemetry/library/transport.py @@ -3,259 +3,172 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""HTTP transport implementation for OneCollector exporter.""" +"""HTTP transport for the OneCollector exporter (standard library only). + +Posts Common Schema JSON to the OneCollector endpoint using ``urllib`` so the +telemetry pipeline has no third-party dependency. +""" import gzip +import queue +import threading +import time +import urllib.error +import urllib.request import zlib -from abc import ABC, abstractmethod from io import BytesIO -from typing import TYPE_CHECKING, Callable, Optional - -import requests +from typing import Callable, Optional -from olive.telemetry.library.event_source import event_source from olive.telemetry.library.options import CompressionType -if TYPE_CHECKING: - from olive.telemetry.library.callback_manager import CallbackManager, PayloadTransmittedCallbackArgs - - -class ITransport(ABC): - """Abstract base class for transports.""" - - @abstractmethod - def send(self, payload: bytes, timeout_sec: float, item_count: int = 1) -> tuple[bool, Optional[int]]: - """Send a payload. - - Args: - payload: The data to send - timeout_sec: Timeout in seconds - item_count: Number of items in the payload (for callbacks) - - Returns: - Tuple of (success, status_code) - - """ - - @abstractmethod - def register_payload_transmitted_callback( - self, callback: Callable[["PayloadTransmittedCallbackArgs"], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback for payload transmission events. - - Args: - callback: Function to call when payload is transmitted - include_failures: Whether to invoke callback on failures - - Returns: - Function to call to unregister the callback - - """ - -class HttpJsonPostTransport(ITransport): - """HTTP JSON POST transport implementation. - - Sends telemetry data to OneCollector via HTTP POST with JSON payload. - """ +class HttpJsonPostTransport: + """HTTP JSON POST transport using ``urllib`` (no third-party dependency).""" def __init__( self, endpoint: str, ikey: str, compression: CompressionType, - session: requests.Session, - callback_manager: Optional["CallbackManager"] = None, - sdk_version: str = "OTel-python-1.0.0", + sdk_version: str = "py-olive-1.0.0", ): - """Initialize the HTTP transport. - - Args: - endpoint: OneCollector endpoint URL - ikey: Instrumentation key - compression: Compression type to use - session: Requests session for connection pooling - callback_manager: Optional callback manager for payload events - sdk_version: SDK version string - - """ self.endpoint = endpoint self.ikey = ikey self.compression = compression - self.session = session self.sdk_version = sdk_version - self.callback_manager = callback_manager - # Build base headers self.headers = { "x-apikey": ikey, - "User-Agent": "Python/3 HttpClient", - "Host": "mobile.events.data.microsoft.com", + "User-Agent": "Python/3 urllib", "Content-Type": "application/x-json-stream; charset=utf-8", "sdk-version": sdk_version, "NoResponseBody": "true", } - if compression != CompressionType.NO_COMPRESSION: self.headers["Content-Encoding"] = compression.value + self._worker_lock = threading.Lock() + self._inflight_worker: Optional[threading.Thread] = None + self._inflight_results = None + self._inflight_request_key = None - def register_payload_transmitted_callback( - self, callback: Callable[["PayloadTransmittedCallbackArgs"], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback for payload transmission events. - - Args: - callback: Function to call when payload is transmitted - include_failures: Whether to invoke callback on failures - - Returns: - Function to call to unregister the callback - - """ - if self.callback_manager is None: - # Import here to avoid circular dependency - from olive.telemetry.library.callback_manager import CallbackManager - - self.callback_manager = CallbackManager() - - return self.callback_manager.register(callback, include_failures) - - def send(self, payload: bytes, timeout_sec: float, item_count: int = 1) -> tuple[bool, Optional[int]]: - """Send payload via HTTP POST. - - Args: - payload: Uncompressed payload bytes - timeout_sec: Request timeout in seconds - item_count: Number of items in the payload (for callbacks) - - Returns: - Tuple of (success, status_code) - - """ - payload_size_bytes = len(payload) - + def send( + self, + payload: bytes, + timeout_sec: float, + item_count: int = 1, + on_send_admitted: Optional[Callable[[], None]] = None, + ) -> tuple[bool, Optional[int]]: + """Send payload via HTTP POST. Returns (success, status_code).""" try: - # Compress payload compressed_payload = self._compress(payload) - - # Update headers with content length headers = {**self.headers, "Content-Length": str(len(compressed_payload))} + request = urllib.request.Request(url=self.endpoint, data=compressed_payload, headers=headers, method="POST") - # Send request - try: - response = self.session.post( - url=self.endpoint, data=compressed_payload, headers=headers, timeout=timeout_sec - ) - except (requests.exceptions.ConnectionError, requests.exceptions.Timeout): - # Retry once on transient transport errors - response = self.session.post( - url=self.endpoint, data=compressed_payload, headers=headers, timeout=timeout_sec - ) - - # Check response - success = response.ok - status_code = response.status_code - - # Invoke callbacks - if self.callback_manager: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - - self.callback_manager.notify( - PayloadTransmittedCallbackArgs( - succeeded=success, - status_code=status_code, - payload_size_bytes=payload_size_bytes, - item_count=item_count, - payload_bytes=payload, - ) - ) - - if success: - return True, status_code - else: - # Log error response - if event_source.is_error_logging_enabled: - collector_error = response.headers.get("Collector-Error", "") - error_details = response.text[:100] if response.text else "" - event_source.http_transport_error_response( - "HttpJsonPost", status_code, collector_error, error_details - ) - return False, status_code - - except requests.exceptions.Timeout: - # Invoke failure callbacks - if self.callback_manager: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - - self.callback_manager.notify( - PayloadTransmittedCallbackArgs( - succeeded=False, - status_code=None, - payload_size_bytes=payload_size_bytes, - item_count=item_count, - payload_bytes=payload, - ) - ) - - event_source.transport_exception_thrown("HttpJsonPost", Exception("Request timeout")) + if on_send_admitted is not None: + on_send_admitted() + success, status_code = self._do_request(request, timeout_sec) + return success, status_code + except Exception: return False, None - except Exception as ex: - # Invoke failure callbacks - if self.callback_manager: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - self.callback_manager.notify( - PayloadTransmittedCallbackArgs( - succeeded=False, - status_code=None, - payload_size_bytes=payload_size_bytes, - item_count=item_count, - payload_bytes=payload, - ) - ) + def _do_request(self, request: "urllib.request.Request", timeout_sec: float) -> tuple[bool, Optional[int]]: + """Run the request behind a wall-clock deadline, including DNS resolution.""" + request_key = (request.full_url, bytes(request.data or b"")) + with self._worker_lock: + worker = self._inflight_worker + if worker is not None: + if worker.is_alive(): + return (False, None) + result = self._consume_inflight_result() + if self._inflight_request_key == request_key: + self._clear_inflight() + return result + self._clear_inflight() + + results = queue.Queue(maxsize=1) + worker = threading.Thread( + target=self._run_request_worker, + args=(request, timeout_sec, results), + name="olive-telemetry-http", + daemon=True, + ) + self._inflight_worker = worker + self._inflight_results = results + self._inflight_request_key = request_key + worker.start() + + worker.join(max(0.0, timeout_sec)) + with self._worker_lock: + if worker.is_alive(): + return (False, None) + result = self._consume_inflight_result() + self._clear_inflight() + return result - event_source.transport_exception_thrown("HttpJsonPost", ex) - return False, None + @staticmethod + def _run_request_worker(request, timeout_sec: float, results) -> None: + try: + result = HttpJsonPostTransport._do_request_blocking(request, timeout_sec) + except Exception: + result = (False, None) + try: + results.put_nowait(result) + except queue.Full: + # The one-result queue may already contain this worker's terminal result. + pass - def _compress(self, data: bytes) -> bytes: - """Compress data according to configured compression type. + def _consume_inflight_result(self) -> tuple[bool, Optional[int]]: + try: + return self._inflight_results.get_nowait() + except (AttributeError, queue.Empty): + return (False, None) - Args: - data: Uncompressed data + def _clear_inflight(self) -> None: + self._inflight_worker = None + self._inflight_results = None + self._inflight_request_key = None - Returns: - Compressed data + @staticmethod + def _do_request_blocking(request: "urllib.request.Request", timeout_sec: float) -> tuple[bool, Optional[int]]: + """Perform the request, retrying once on a transient connection error.""" + deadline = time.monotonic() + max(0.0, timeout_sec) + for attempt in range(2): + remaining = deadline - time.monotonic() + if remaining <= 0.0: + return (False, None) + try: + with urllib.request.urlopen(request, timeout=remaining) as response: + status = getattr(response, "status", response.getcode()) + return (200 <= status < 300, status) + except urllib.error.HTTPError as http_err: + # Server responded with a non-2xx status (4xx/5xx): not retried here. + try: + http_err.close() + except Exception: + # The HTTP status remains authoritative if response cleanup fails. + pass + return (False, http_err.code) + except (urllib.error.URLError, TimeoutError, OSError): + # Connection-level failure: retry once, then give up. + if attempt == 0: + continue + return (False, None) + return (False, None) - """ + def _compress(self, data: bytes) -> bytes: if self.compression == CompressionType.DEFLATE: - # Raw deflate (no zlib header) compressor = zlib.compressobj(wbits=-zlib.MAX_WBITS) - compressed = compressor.compress(data) - compressed += compressor.flush() - return compressed - + return compressor.compress(data) + compressor.flush() elif self.compression == CompressionType.GZIP: gzip_buffer = BytesIO() with gzip.GzipFile(fileobj=gzip_buffer, mode="w") as gzip_file: gzip_file.write(data) return gzip_buffer.getvalue() - - else: # NO_COMPRESSION - return data + return data @staticmethod def is_retryable(status_code: Optional[int]) -> bool: - """Check if a response status code indicates the request should be retried. - - Args: - status_code: HTTP status code, or None if request failed - - Returns: - True if request should be retried - - """ + """Whether a response status indicates the request should be retried.""" if status_code is None: return True # Network errors are retryable - - # Retryable status codes - return status_code in {408, 429, 500, 502, 503, 504} + return status_code in {408, 429} or 500 <= status_code <= 599 diff --git a/olive/telemetry/offline_store.py b/olive/telemetry/offline_store.py new file mode 100644 index 0000000000..16603d4602 --- /dev/null +++ b/olive/telemetry/offline_store.py @@ -0,0 +1,274 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""SQLite-backed durable queue for telemetry events. + +A deliberately small subset of the Microsoft 1DS C++ SDK offline store +(cpp_client_telemetry/lib/offline/OfflineStorage_SQLite.cpp): a single FIFO +table of serialized event payloads. An uploader drains it, deleting rows on +success, dropping them on a permanent (non-retryable) send result, and leaving +them for the next attempt on a transient failure. Because every event is +written to disk before any network call, the process can exit at any time +without losing data and without an exit-time flush. + +Uses only the Python standard library (``sqlite3``), so it adds no dependency. + +Intentionally omitted from the full 1DS store (not needed for low-volume CLI +telemetry): per-event priority (``latency``), persistence classes, leases, +per-row retry counters, tenant multiplexing, and the ``settings`` table. The +schema version is tracked with SQLite's built-in ``PRAGMA user_version``. +""" + +import os +import sqlite3 +import threading +import time +from contextlib import contextmanager, suppress +from pathlib import Path +from typing import Optional + +SCHEMA_VERSION = 3 + + +def _chmod_best_effort(path: str, mode: int) -> None: + if os.name == "nt" or not path: + return + try: + Path(path).chmod(mode) + except OSError: + # Permission tightening is best-effort on filesystems that do not support chmod. + pass + + +class OfflineEventStore: + """Durable FIFO queue of serialized telemetry event payloads. + + All methods are best-effort and swallow storage errors: telemetry must + never crash the host application. Thread-safe via a per-instance lock; + tolerant of concurrent processes via WAL mode + ``busy_timeout``. + """ + + def __init__(self, db_path: str, max_records: int = 2048, busy_timeout_ms: int = 3000): + self._db_path = db_path + self._max_records = max_records + # When full, trim back to this watermark so we don't trim on every insert. + self._trim_target = max(1, (max_records * 3) // 4) + self._busy_timeout_ms = busy_timeout_ms + self._lock = threading.Lock() + self._conn: Optional[sqlite3.Connection] = None + self._initialize() + + def _initialize(self) -> None: + parent = os.path.dirname(self._db_path) + try: + os.makedirs(parent, mode=0o700, exist_ok=True) + _chmod_best_effort(parent, 0o700) + except Exception: + # sqlite3.connect below reports whether storage can actually be opened. + pass + conn = None + try: + conn = sqlite3.connect( + self._db_path, + timeout=self._busy_timeout_ms / 1000.0, + check_same_thread=False, + ) + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute(f"PRAGMA busy_timeout={self._busy_timeout_ms}") + conn.execute("BEGIN IMMEDIATE") + conn.execute( + "CREATE TABLE IF NOT EXISTS events " + "(id INTEGER PRIMARY KEY AUTOINCREMENT, payload BLOB NOT NULL, " + "available_at REAL NOT NULL DEFAULT 0, acknowledged INTEGER NOT NULL DEFAULT 0)" + ) + columns = {row[1] for row in conn.execute("PRAGMA table_info(events)")} + if "available_at" not in columns: + conn.execute("ALTER TABLE events ADD COLUMN available_at REAL NOT NULL DEFAULT 0") + if "acknowledged" not in columns: + conn.execute("ALTER TABLE events ADD COLUMN acknowledged INTEGER NOT NULL DEFAULT 0") + conn.execute(f"PRAGMA user_version={SCHEMA_VERSION}") + conn.commit() + conn.execute(f"PRAGMA busy_timeout={self._busy_timeout_ms}") + self._conn = conn + self._harden_permissions() + except Exception: + if conn is not None: + with suppress(Exception): + conn.close() + self._conn = None + + def _ensure_open(self) -> bool: + return self._conn is not None + + def _harden_permissions(self) -> None: + _chmod_best_effort(os.path.dirname(self._db_path), 0o700) + for path in (self._db_path, self._db_path + "-wal", self._db_path + "-shm"): + if os.path.exists(path): + _chmod_best_effort(path, 0o600) + + @contextmanager + def _bounded_busy_timeout(self, deadline: Optional[float]): + if deadline is None or self._conn is None: + yield + return + remaining_ms = max(0, int((deadline - time.monotonic()) * 1000)) + self._conn.execute(f"PRAGMA busy_timeout={min(self._busy_timeout_ms, remaining_ms)}") + try: + yield + finally: + with suppress(Exception): + self._conn.execute(f"PRAGMA busy_timeout={self._busy_timeout_ms}") + + @property + def is_open(self) -> bool: + with self._lock: + return self._ensure_open() + + @property + def db_path(self) -> str: + return self._db_path + + def store(self, payload: bytes) -> bool: + """Append one serialized event; trims the oldest rows if over capacity.""" + return self._store(payload, 0.0) is not None + + def reserve(self, payload: bytes, available_after_seconds: float) -> Optional[int]: + """Persist an event but defer draining until it is released or the delay expires.""" + return self._store(payload, time.time() + max(0.0, available_after_seconds)) + + def _store(self, payload: bytes, available_at: float) -> Optional[int]: + if not payload: + return None + with self._lock: + if not self._ensure_open(): + return None + try: + cursor = self._conn.execute( + "INSERT INTO events (payload, available_at) VALUES (?, ?)", + (sqlite3.Binary(payload), available_at), + ) + count = self._conn.execute("SELECT COUNT(*) FROM events").fetchone()[0] + if count > self._max_records: + self._conn.execute( + "DELETE FROM events WHERE id IN " + "(SELECT id FROM events WHERE available_at <= ? ORDER BY id ASC LIMIT ?)", + (time.time(), count - self._trim_target), + ) + self._conn.commit() + self._harden_permissions() + return int(cursor.lastrowid) + except Exception: + with suppress(Exception): + self._conn.rollback() + return None + + def get_batch(self, max_count: int) -> list[tuple[int, bytes]]: + """Return up to ``max_count`` oldest events as (id, payload) pairs.""" + return self.get_batch_for_upload(max_count) or [] + + def get_batch_for_upload( + self, max_count: int, deadline: Optional[float] = None + ) -> Optional[list[tuple[int, bytes]]]: + """Return an uploadable batch, or None when local storage failed.""" + with self._lock: + if not self._ensure_open(): + return None + try: + with self._bounded_busy_timeout(deadline): + rows = self._conn.execute( + "SELECT id, payload FROM events " + "WHERE available_at <= ? AND acknowledged=0 ORDER BY id ASC LIMIT ?", + (time.time(), max_count if max_count > 0 else -1), + ).fetchall() + return [(r[0], bytes(r[1])) for r in rows] + except Exception: + return None + + def get_acknowledged_ids(self, max_count: int, deadline: Optional[float] = None) -> Optional[list[int]]: + """Return terminally handled rows that only need local deletion.""" + with self._lock: + if not self._ensure_open(): + return None + try: + with self._bounded_busy_timeout(deadline): + rows = self._conn.execute( + "SELECT id FROM events WHERE acknowledged=1 ORDER BY id ASC LIMIT ?", + (max_count if max_count > 0 else -1,), + ).fetchall() + return [int(row[0]) for row in rows] + except Exception: + return None + + def release(self, row_id: int, payload: Optional[bytes] = None) -> bool: + """Make a reserved event drainable, optionally replacing its payload.""" + with self._lock: + if not self._ensure_open(): + return False + try: + if payload is None: + cursor = self._conn.execute("UPDATE events SET available_at=0 WHERE id=?", (row_id,)) + else: + cursor = self._conn.execute( + "UPDATE events SET payload=?, available_at=0 WHERE id=?", + (sqlite3.Binary(payload), row_id), + ) + self._conn.commit() + return cursor.rowcount == 1 + except Exception: + with suppress(Exception): + self._conn.rollback() + return False + + def acknowledge(self, ids: list[int], deadline: Optional[float] = None) -> bool: + """Persist that rows were delivered or permanently rejected.""" + if not ids: + return True + with self._lock: + if not self._ensure_open(): + return False + try: + with self._bounded_busy_timeout(deadline): + self._conn.executemany("UPDATE events SET acknowledged=1 WHERE id=?", [(i,) for i in ids]) + self._conn.commit() + return True + except Exception: + with suppress(Exception): + self._conn.rollback() + return False + + def delete(self, ids: list[int], deadline: Optional[float] = None) -> bool: + """Remove rows by id (after a successful upload or a permanent drop).""" + if not ids: + return True + with self._lock: + if not self._ensure_open(): + return False + try: + with self._bounded_busy_timeout(deadline): + self._conn.executemany("DELETE FROM events WHERE id=?", [(i,) for i in ids]) + self._conn.commit() + return True + except Exception: + # Failed deletes leave rows durable for a later drain attempt. + with suppress(Exception): + self._conn.rollback() + return False + + def count(self) -> int: + with self._lock: + if not self._ensure_open(): + return 0 + try: + return int(self._conn.execute("SELECT COUNT(*) FROM events").fetchone()[0]) + except Exception: + return 0 + + def close(self) -> None: + with self._lock: + if self._conn is not None: + with suppress(Exception): + self._conn.close() + self._conn = None diff --git a/olive/telemetry/process_lock.py b/olive/telemetry/process_lock.py new file mode 100644 index 0000000000..14ce23fcb4 --- /dev/null +++ b/olive/telemetry/process_lock.py @@ -0,0 +1,117 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Cross-platform single-holder advisory lock (standard library only). + +Used so that, when several processes on one device share a telemetry database, +only one of them runs the uploader's drain loop at a time. Other processes keep +writing events durably to the store; the lock holder drains everyone's rows. +This avoids the same event being uploaded twice by concurrent drainers without +needing per-row reservation bookkeeping. + +The lock is an OS advisory lock on a sidecar file (``msvcrt`` on Windows, +``fcntl`` on POSIX). It is released explicitly and also by the OS when the +process exits, so a crashed holder never blocks other processes permanently. +""" + +import errno +import os +import time +from contextlib import suppress + + +class ProcessDrainLock: + """Non-blocking exclusive advisory lock backed by a sidecar file.""" + + def __init__(self, lock_path: str): + self._lock_path = lock_path + self._fh = None + self._posix_lock_api = None + + @property + def held(self) -> bool: + return self._fh is not None + + def acquire(self, timeout_seconds: float = 0.0) -> bool: + """Try to acquire the lock without blocking. Returns True if held.""" + if self._fh is not None: + return True + deadline = time.monotonic() + max(0.0, timeout_seconds) + while True: + fh = None + try: + with suppress(Exception): + os.makedirs(os.path.dirname(self._lock_path), exist_ok=True) + # The handle must remain open while the advisory lock is held. + fh = open(self._lock_path, "a+b") # noqa: SIM115 # pylint: disable=consider-using-with + if os.name == "nt": + import msvcrt + + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_NBLCK, 1) + else: + import fcntl + + try: + fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + self._posix_lock_api = "flock" + except AttributeError: + fcntl.lockf(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + self._posix_lock_api = "lockf" + except OSError as exc: + unsupported_errors = {errno.ENOSYS} + if hasattr(errno, "ENOTSUP"): + unsupported_errors.add(errno.ENOTSUP) + if hasattr(errno, "EOPNOTSUPP"): + unsupported_errors.add(errno.EOPNOTSUPP) + if exc.errno not in unsupported_errors: + raise + fcntl.lockf(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + self._posix_lock_api = "lockf" + self._fh = fh + return True + except Exception: + self._posix_lock_api = None + if fh is not None: + with suppress(Exception): + fh.close() + if time.monotonic() >= deadline: + return False + time.sleep(min(0.01, max(0.0, deadline - time.monotonic()))) + + def release(self) -> None: + if self._fh is None: + return + fh = self._fh + posix_lock_api = self._posix_lock_api + self._fh = None + self._posix_lock_api = None + try: + if os.name == "nt": + import msvcrt + + try: + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1) + except Exception: + # The OS releases the lock when the handle closes below. + pass + else: + import fcntl + + try: + if posix_lock_api == "lockf": + fcntl.lockf(fh.fileno(), fcntl.LOCK_UN) + else: + fcntl.flock(fh.fileno(), fcntl.LOCK_UN) + except Exception: + # The OS releases the lock when the handle closes below. + pass + finally: + try: + fh.close() + except Exception: + # Lock cleanup must never fail telemetry callers. + pass diff --git a/olive/telemetry/recipe_telemetry.py b/olive/telemetry/recipe_telemetry.py new file mode 100644 index 0000000000..18f42bf048 --- /dev/null +++ b/olive/telemetry/recipe_telemetry.py @@ -0,0 +1,457 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import functools +import json +import re +from copy import deepcopy +from os import PathLike +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import TYPE_CHECKING, Any, Optional, Union + +from olive.common.config_utils import load_config_file +from olive.common.utils import hash_dict +from olive.package_config import OlivePackageConfig +from olive.systems.common import SystemType +from olive.telemetry.telemetry import is_ci_environment +from olive.telemetry.telemetry_redaction import ( + is_environment_config_key_for_telemetry, + is_path_like_config_key_for_telemetry, + is_sensitive_config_key_for_telemetry, + normalize_config_key_for_telemetry, + scrub_string_for_telemetry, +) + +if TYPE_CHECKING: + from olive.workflows.run.config import RunConfig + +RECIPE_HASH_REDACTED_VALUE = "" +CONFIG_REFERENCE_REDACTED_VALUE = "" +CONFIG_CALLABLE_REDACTED_VALUE = "" +CONFIG_UNKNOWN_REDACTED_VALUE = "" +CONFIG_SECRET_REDACTED_VALUE = "" +RECIPE_HASH_REDACTED_KEYS = { + "output_dir", + "cache_dir", + "tempdir", + "additional_files", + "dockerfile", + "build_context_path", + "python_environment_path", + "prepend_to_path", + "script_dir", + "model_script", + # package_config is tracked separately via package_config_provided and + # package_config_overrides, but excluded from recipe_hash because it is an + # environment/infrastructure path. + "package_config", + "work_dir", +} +CONFIG_SNAPSHOT_REDACTED_KEYS = RECIPE_HASH_REDACTED_KEYS | { + "model_path", + "_name_or_path", + "adapter_path", + "user_script", +} +HF_MODEL_IDENTIFIER_KEYS = {"model_path", "_name_or_path"} +CONFIG_REFERENCE_KEYS = {"host", "target", "evaluator"} +LOCAL_MODEL_FILE_SUFFIXES = {".bin", ".model", ".onnx", ".pb", ".pt", ".pth", ".safetensors", ".tflite"} +HF_CACHE_MODEL_PATTERN = re.compile(r"(?:^|[\\/])models--([^\\/]+)--([^\\/]+)(?:[\\/]|$)") +HF_REPO_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*(/[A-Za-z0-9][A-Za-z0-9._-]*)?$") +_NO_OVERRIDE = object() + + +def _build_recipe_result_metadata( + run_config_input: Union[str, Path, dict], + run_config_telemetry_input: Optional[Any], + run_config: Optional["RunConfig"], + recipe_telemetry_metadata: Optional[dict[str, Any]], + *, + list_required_packages: bool, + package_config_input: Optional[Union[str, Path, dict]], + package_config_provided: bool, +) -> dict[str, Any]: + metadata = dict(recipe_telemetry_metadata or {}) + default_source, default_format = _classify_run_config_source(run_config_input) + metadata.setdefault("recipe_source", default_source) + metadata.setdefault("recipe_format", default_format) + metadata.setdefault("execution_mode", "list_required_packages" if list_required_packages else "run") + metadata.setdefault("package_config_provided", package_config_provided) + config_overrides = metadata.pop("config_overrides", _NO_OVERRIDE) + if config_overrides is _NO_OVERRIDE: + config_overrides = _build_config_overrides(run_config_telemetry_input) + else: + config_overrides = _sanitize_provided_config_snapshot(config_overrides) + if config_overrides is not None: + metadata["config_overrides"] = config_overrides + package_config_overrides = metadata.pop("package_config_overrides", _NO_OVERRIDE) + if package_config_overrides is not _NO_OVERRIDE: + package_config_overrides = _sanitize_provided_config_snapshot(package_config_overrides) + elif package_config_provided: + package_config_overrides = _build_package_config_overrides(package_config_input) + else: + package_config_overrides = None + if package_config_overrides is not None: + metadata["package_config_overrides"] = package_config_overrides + metadata["is_ci"] = is_ci_environment() + + if run_config is None: + metadata.setdefault("recipe_name", metadata.get("recipe_command") or "WorkflowRun") + return metadata + + run_config_json = run_config.to_json(make_absolute=False) + model_metadata = _extract_input_model_metadata(run_config_json["input_model"]) + target_metadata = _extract_target_metadata(run_config) + host_metadata = _extract_host_metadata(run_config) + pass_types = _get_used_pass_types(run_config) + + metadata.setdefault("recipe_name", metadata.get("recipe_command") or run_config.workflow_id) + metadata.setdefault("workflow_id", run_config.workflow_id) + metadata.setdefault("recipe_hash", _build_recipe_hash(run_config_json)) + metadata.setdefault("input_model_type", run_config.input_model.type) + metadata.setdefault("input_model_source", model_metadata["input_model_source"]) + metadata.setdefault("model_task", model_metadata["model_task"]) + _set_metadata_if_present(metadata, target_metadata) + _set_metadata_if_present(metadata, host_metadata) + metadata.setdefault("pass_types", ";".join(pass_types)) + metadata.setdefault("pass_count", len(pass_types)) + metadata.setdefault("data_config_count", len(run_config.data_configs)) + metadata.setdefault("search_enabled", bool(run_config.engine.search_strategy)) + return metadata + + +def _classify_run_config_source(run_config_input: Any) -> tuple[str, str]: + if isinstance(run_config_input, dict): + return "config_dict", "dict" + + if isinstance(run_config_input, (str, PathLike)): + suffix = Path(run_config_input).suffix.lstrip(".").lower() + return "config_file", suffix or "unknown" + + return "config_object", "object" + + +def _build_config_overrides(config_input: Any) -> Optional[str]: + try: + config_data = _load_config_input_for_telemetry(config_input) + if config_data is None: + return None + + snapshot = _sanitize_config_snapshot(config_data) + if snapshot in (None, {}, []): + return None + + return json.dumps(snapshot, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + except Exception: + return None + + +def _sanitize_provided_config_snapshot(value: Any) -> Optional[str]: + if isinstance(value, str): + try: + value = json.loads(value) + except (TypeError, ValueError): + return None + return _build_config_overrides(value) + + +def _build_package_config_overrides(config_input: Any) -> Optional[str]: + try: + config_data = _load_config_input_for_telemetry(config_input) + if not isinstance(config_data, dict): + return None + + default_config = _load_default_package_config_for_telemetry() + baseline = ( + _normalize_package_config_snapshot(default_config) if isinstance(default_config, dict) else _NO_OVERRIDE + ) + overrides = _extract_config_overrides(_normalize_package_config_snapshot(config_data), baseline) + if overrides is _NO_OVERRIDE: + return None + + snapshot = _sanitize_config_snapshot(overrides) + if not isinstance(snapshot, dict): + return None + + return json.dumps(snapshot, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + except Exception: + return None + + +@functools.lru_cache +def _load_default_package_config_for_telemetry() -> Optional[dict[str, Any]]: + try: + default_config = load_config_file(OlivePackageConfig.get_default_config_path()) + except Exception: + return None + + return default_config if isinstance(default_config, dict) else None + + +def _normalize_package_config_snapshot(config_data: Any) -> Any: + if not isinstance(config_data, dict): + return config_data + + normalized = deepcopy(config_data) + passes = normalized.get("passes") + if isinstance(passes, dict): + normalized["passes"] = {str(pass_name).lower(): pass_config for pass_name, pass_config in passes.items()} + return normalized + + +def _extract_config_overrides(value: Any, baseline: Any = _NO_OVERRIDE) -> Any: + if baseline is _NO_OVERRIDE: + return deepcopy(value) + + if isinstance(value, dict) and isinstance(baseline, dict): + overrides = {} + for key, child_value in value.items(): + child_override = _extract_config_overrides(child_value, baseline.get(key, _NO_OVERRIDE)) + if child_override is not _NO_OVERRIDE: + overrides[key] = child_override + if overrides: + return overrides + return _NO_OVERRIDE + + if isinstance(value, list): + if isinstance(baseline, list) and value == baseline: + return _NO_OVERRIDE + return deepcopy(value) + + if isinstance(value, tuple): + value_list = list(value) + baseline_list = list(baseline) if isinstance(baseline, tuple) else baseline + if isinstance(baseline_list, list) and value_list == baseline_list: + return _NO_OVERRIDE + return value_list + + return deepcopy(value) if value != baseline else _NO_OVERRIDE + + +def _load_config_input_for_telemetry(config_input: Any) -> Optional[Any]: + if config_input is None: + return None + if isinstance(config_input, dict): + return deepcopy(config_input) + if isinstance(config_input, (str, PathLike)): + return load_config_file(config_input) + + model_dump = getattr(config_input, "model_dump", None) + if callable(model_dump): + return model_dump(exclude_defaults=True, exclude_none=True, by_alias=True) + return None + + +def _sanitize_config_snapshot(value: Any, key: Optional[str] = None, model_type: Optional[str] = None) -> Any: + normalized_key = normalize_config_key_for_telemetry(key) + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return CONFIG_SECRET_REDACTED_VALUE + if normalized_key in HF_MODEL_IDENTIFIER_KEYS: + if str(model_type).lower() == "hfmodel": + hf_model_id = _extract_huggingface_model_id(value) + if hf_model_id: + return hf_model_id + return RECIPE_HASH_REDACTED_VALUE + if normalized_key in CONFIG_SNAPSHOT_REDACTED_KEYS or is_path_like_config_key_for_telemetry(normalized_key): + return RECIPE_HASH_REDACTED_VALUE + if normalized_key in CONFIG_REFERENCE_KEYS and isinstance(value, str): + return CONFIG_REFERENCE_REDACTED_VALUE + + if isinstance(value, dict): + child_model_type = _get_model_type(value) or model_type + if normalized_key == "systems": + return [_sanitize_config_snapshot(system, "system", child_model_type) for system in value.values()] + if normalized_key == "passes": + passes = [] + for pass_configs in value.values(): + if isinstance(pass_configs, list): + passes.extend(pass_configs) + else: + passes.append(pass_configs) + return [_sanitize_config_snapshot(pass_config, "pass", child_model_type) for pass_config in passes] + if normalized_key == "evaluators": + return [ + _sanitize_config_snapshot(evaluator, "evaluator_config", child_model_type) + for evaluator in value.values() + ] + return { + child_key: _sanitize_config_snapshot(child_value, child_key, child_model_type) + for child_key, child_value in value.items() + if child_value is not None + } + if isinstance(value, list): + return [_sanitize_config_snapshot(item, key, model_type) for item in value] + if isinstance(value, tuple): + return [_sanitize_config_snapshot(item, key, model_type) for item in value] + if isinstance(value, Path): + return RECIPE_HASH_REDACTED_VALUE + if callable(value): + return CONFIG_CALLABLE_REDACTED_VALUE + if isinstance(value, str): + return scrub_string_for_telemetry(value) + if isinstance(value, (int, float, bool)) or value is None: + return value + if hasattr(value, "value") and isinstance(value.value, (str, int, float, bool)): + return scrub_string_for_telemetry(value.value) if isinstance(value.value, str) else value.value + return CONFIG_UNKNOWN_REDACTED_VALUE + + +def _get_model_type(config: dict[str, Any]) -> Optional[str]: + model_type = config.get("type") + return str(model_type).lower() if model_type is not None else None + + +def _extract_huggingface_model_id(model_identifier: Any) -> Optional[str]: + if not isinstance(model_identifier, str): + return None + + identifier = model_identifier.strip() + if not identifier: + return None + + if identifier.startswith("https://huggingface.co/"): + parts = identifier.removeprefix("https://huggingface.co/").strip("/").split("/") + if len(parts) >= 2: + return f"{parts[0]}/{parts[1]}" + if parts and parts[0]: + return parts[0] + + if match := HF_CACHE_MODEL_PATTERN.search(identifier): + return f"{match.group(1)}/{match.group(2)}" + + if HF_REPO_ID_PATTERN.match(identifier) and not _has_local_model_file_suffix(identifier): + return identifier + + return None + + +def _extract_input_model_metadata(input_model_config: dict[str, Any]) -> dict[str, Optional[str]]: + model_config = input_model_config.get("config", {}) + model_attributes = model_config.get("model_attributes", {}) + model_task = model_attributes.get("hf_task") or model_config.get("task") + raw_identifier = model_attributes.get("_name_or_path") or model_config.get("model_path") + return { + "input_model_source": _classify_input_model_source(raw_identifier), + "model_task": str(model_task) if model_task is not None else None, + } + + +def _classify_input_model_source(model_identifier: Any) -> str: + if model_identifier is None: + return "unknown" + if isinstance(model_identifier, dict): + resource_type = model_identifier.get("type") + if resource_type == "azureml_registry_model": + return "azureml" + return "structured_resource" + + identifier = str(model_identifier) + if identifier.startswith("azureml://"): + return "azureml" + if identifier.startswith("https://huggingface.co/"): + return "huggingface_url" + if identifier.startswith(("http://", "https://")): + return "url" + + if _is_explicit_local_model_path(identifier): + suffix = PureWindowsPath(identifier).suffix or PurePosixPath(identifier).suffix + return "local_file" if suffix else "local_folder" + return "string_name" + + +def _is_explicit_local_model_path(identifier: str) -> bool: + if _has_local_model_file_suffix(identifier): + return True + return ( + identifier.startswith(("./", "../", ".\\", "..\\", "~/", "~\\", "/", "\\\\")) + or PureWindowsPath(identifier).is_absolute() + or PurePosixPath(identifier).is_absolute() + ) + + +def _has_local_model_file_suffix(identifier: str) -> bool: + suffix = PureWindowsPath(identifier).suffix or PurePosixPath(identifier).suffix + return suffix.lower() in LOCAL_MODEL_FILE_SUFFIXES + + +def _extract_target_metadata(run_config: "RunConfig") -> dict[str, Optional[str]]: + target_system = run_config.engine.target + return _extract_system_metadata(target_system, "target") + + +def _extract_host_metadata(run_config: "RunConfig") -> dict[str, Optional[str]]: + host_system = run_config.engine.host + if host_system is None: + return { + "host_system_type": SystemType.Local.value, + } + return _extract_system_metadata(host_system, "host") + + +def _extract_system_metadata(system_config: Optional[Any], field_prefix: str) -> dict[str, Optional[str]]: + system_type = system_config.type.value if system_config is not None else None + device = None + execution_provider = None + execution_providers = None + + accelerators = system_config.config.accelerators if system_config and system_config.config else None + if accelerators: + accelerator = accelerators[0] + device = str(accelerator.device) if accelerator.device is not None else None + ep_values = accelerator.get_ep_strs() or [] + if ep_values: + execution_provider = ep_values[0] + execution_providers = ";".join(ep_values) + + return { + f"{field_prefix}_system_type": system_type, + f"{field_prefix}_device": device, + f"{field_prefix}_execution_provider": execution_provider, + f"{field_prefix}_execution_providers": execution_providers, + } + + +def _set_metadata_if_present(metadata: dict[str, Any], values: dict[str, Optional[str]]) -> None: + for key, value in values.items(): + if value is not None: + metadata.setdefault(key, value) + + +def _get_used_pass_types(run_config: "RunConfig") -> list[str]: + return ( + [pass_config.type for _, pass_configs in run_config.passes.items() for pass_config in pass_configs] + if run_config.passes + else [] + ) + + +def _build_recipe_hash(run_config_json: dict[str, Any]) -> str: + sanitized = deepcopy(run_config_json) + _redact_recipe_hash_keys(sanitized) + return hash_dict(sanitized)[:16] + + +def _redact_recipe_hash_keys(value: Any, key: Optional[str] = None) -> Any: + normalized_key = normalize_config_key_for_telemetry(key) + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return CONFIG_SECRET_REDACTED_VALUE + if normalized_key in RECIPE_HASH_REDACTED_KEYS or is_path_like_config_key_for_telemetry(normalized_key): + return RECIPE_HASH_REDACTED_VALUE + if isinstance(value, dict): + for child_key in list(value): + value[child_key] = _redact_recipe_hash_keys(value[child_key], child_key) + elif isinstance(value, list): + for index, item in enumerate(value): + value[index] = _redact_recipe_hash_keys(item, key) + elif isinstance(value, tuple): + return [_redact_recipe_hash_keys(item, key) for item in value] + elif isinstance(value, Path): + return RECIPE_HASH_REDACTED_VALUE + elif isinstance(value, str): + return scrub_string_for_telemetry(value) + elif callable(value): + return CONFIG_CALLABLE_REDACTED_VALUE + elif hasattr(value, "value") and isinstance(value.value, (str, int, float, bool)): + return value.value + return value diff --git a/olive/telemetry/telemetry.py b/olive/telemetry/telemetry.py index 0ddb690e2a..83cf316b39 100644 --- a/olive/telemetry/telemetry.py +++ b/olive/telemetry/telemetry.py @@ -2,58 +2,73 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Thin wrapper around the OneCollector telemetry logger with event helpers.""" +"""Telemetry singleton backed by a durable SQLite event queue. + +Detailed events are serialized to Common Schema JSON and written to a per-app +SQLite store; a background uploader drains the store to Microsoft OneCollector. +The one-per-process Heartbeat uses the same durable queue. The pipeline uses only +the Python standard library (no OpenTelemetry, no requests). +""" import base64 -import errno import json import os import platform import threading import time -from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional - -from olive.telemetry.constants import CONNECTION_STRING -from olive.telemetry.deviceid import get_encrypted_device_id_and_status -from olive.telemetry.library.event_source import event_source -from olive.telemetry.library.telemetry_logger import TelemetryLogger, get_telemetry_logger -from olive.telemetry.utils import ( - _decode_cache_line, - _encode_cache_line, - _exclusive_file_lock, - get_telemetry_base_dir, +import uuid +from datetime import datetime, timezone +from typing import Any, Optional + +from olive.telemetry.deviceid import get_hashed_device_id_and_status +from olive.telemetry.library.options import OneCollectorExporterOptions +from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper +from olive.telemetry.offline_store import OfflineEventStore +from olive.telemetry.telemetry_redaction import ( + MAX_TELEMETRY_STRING_LENGTH, + scrub_config_snapshot_for_telemetry, + scrub_error_message_for_telemetry, + scrub_value_for_telemetry, ) +from olive.telemetry.uploader import EventUploader +from olive.telemetry.utils import get_telemetry_base_dir -if TYPE_CHECKING: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs +try: + from olive.version import __version__ as VERSION +except Exception: + VERSION = "unknown" # Default event names used by the high-level telemetry helpers. HEARTBEAT_EVENT_NAME = "OliveHeartbeat" +RECIPE_EVENT_NAME = "OliveRecipe" +ACTION_EVENT_NAME = "OliveAction" +ERROR_EVENT_NAME = "OliveError" # CI/CD environment variables whose presence indicates an automated pipeline. _CI_ENV_VARS = ( "CI", # GitHub Actions, GitLab CI, Travis CI, CircleCI, generic "TF_BUILD", # Azure Pipelines "GITHUB_ACTIONS", # GitHub Actions + "GITLAB_CI", # GitLab CI + "CIRCLECI", # CircleCI + "TRAVIS", # Travis CI "JENKINS_URL", # Jenkins "CODEBUILD_BUILD_ID", # AWS CodeBuild "BUILDKITE", # Buildkite + "TEAMCITY_VERSION", # TeamCity + "APPVEYOR", # AppVeyor + "BITBUCKET_BUILD_NUMBER", # Bitbucket Pipelines "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", # Azure DevOps ) -ACTION_EVENT_NAME = "OliveAction" -ERROR_EVENT_NAME = "OliveError" ALLOWED_KEYS = { HEARTBEAT_EVENT_NAME: { "device_id", - "id_status", - "os.name", - "os.version", - "os.release", - "os.arch", - "app_version", - "app_instance_id", + "device_id_status", + "os", + "os_version", + "os_release", + "os_arch", "initTs", }, ACTION_EVENT_NAME: { @@ -61,656 +76,482 @@ "action_name", "duration_ms", "success", - "app_version", - "app_instance_id", "initTs", }, ERROR_EVENT_NAME: { "exception_type", "exception_message", - "app_version", - "app_instance_id", + "initTs", + }, + RECIPE_EVENT_NAME: { + "recipe_name", + "recipe_hash", + "recipe_source", + "recipe_format", + "recipe_command", + "execution_mode", + "workflow_id", + "config_overrides", + "success", + "input_model_type", + "input_model_source", + "model_task", + "target_system_type", + "target_device", + "target_execution_provider", + "target_execution_providers", + "host_system_type", + "host_device", + "host_execution_provider", + "host_execution_providers", + "pass_types", + "pass_count", + "data_config_count", + "search_enabled", + "package_config_provided", + "package_config_overrides", + "is_ci", "initTs", }, } -CRITICAL_EVENTS = {HEARTBEAT_EVENT_NAME} -MAX_CACHE_SIZE_BYTES = 5 * 1024 * 1024 -HARD_MAX_CACHE_SIZE_BYTES = 10 * 1024 * 1024 -CACHE_FILE_NAME = "olive.json" +FIELD_NAMES = { + "device_id": "deviceId", + "device_id_status": "deviceIdStatus", + "os_version": "osVersion", + "os_release": "osRelease", + "os_arch": "osArchitecture", + "invoked_from": "invokedFrom", + "action_name": "actionName", + "duration_ms": "durationMs", + "exception_type": "exceptionType", + "exception_message": "exceptionMessage", + "recipe_name": "recipeName", + "recipe_hash": "recipeHash", + "recipe_source": "recipeSource", + "recipe_format": "recipeFormat", + "recipe_command": "recipeCommand", + "execution_mode": "executionMode", + "workflow_id": "workflowId", + "config_overrides": "configOverrides", + "input_model_type": "inputModelType", + "input_model_source": "inputModelSource", + "model_task": "modelTask", + "target_system_type": "targetSystemType", + "target_device": "targetDevice", + "target_execution_provider": "targetExecutionProvider", + "target_execution_providers": "targetExecutionProviders", + "host_system_type": "hostSystemType", + "host_device": "hostDevice", + "host_execution_provider": "hostExecutionProvider", + "host_execution_providers": "hostExecutionProviders", + "pass_types": "passTypes", + "pass_count": "passCount", + "data_config_count": "dataConfigCount", + "search_enabled": "searchEnabled", + "package_config_provided": "packageConfigProvided", + "package_config_overrides": "packageConfigOverrides", + "is_ci": "isCI", + "app_version": "LibraryVersion", + "app_instance_id": "AppSessionGuid", +} +# Per-app database file. Olive and other apps use separate files so a process +# never drains another app's events (which carry a different tenant key). +DB_FILE_NAME = "olive_telemetry.db" +CI_DB_FILE_NAME = "olive_recipe_telemetry.db" +_HEARTBEAT_RELEASE_SECONDS = 60.0 +_SHUTDOWN_TIMEOUT_SECONDS = 2.0 -class TelemetryCacheHandler: - """Handles caching of failed telemetry events for offline resilience. - Design decisions: - - Single shared cache file (olive.json) for simplicity - - Cache writes are synchronous (fast JSON operations don't need async) - - Cache flush runs in a separate thread (slow network I/O) - - Flush triggered on success when cached events exist - - All critical sections protected by lock to prevent race conditions - - Newline-delimited JSON format for human readability and partial corruption recovery +def _is_environment_signal_truthy(value: str) -> bool: + return value.strip().lower() not in {"", "0", "false", "no", "off"} - Assumptions: - - File I/O (JSON lines) is fast enough for synchronous execution (~microseconds) - - Network I/O is slow and should not block the callback thread - - Successful send indicates network is available to retry cached events - - Cache persists across sessions for offline resilience - """ - def __init__(self, telemetry: "Telemetry") -> None: - self._telemetry = telemetry - # Single shared cache file for all processes - self._cache_file_name = CACHE_FILE_NAME - self._shutdown = False - # Protects all shared state to prevent race conditions - self._lock = threading.Lock() - self._callback_condition = threading.Condition() - self._callbacks_item_count = 0 - self._events_logged = 0 - # Prevents concurrent flush operations - self._is_flushing = False - - def shutdown(self) -> None: - """Signal shutdown to prevent new operations. - - Note: Does NOT flush the cache. Cache persists across sessions for - offline resilience. If network is working, success callbacks already - flushed. If network is down, flushing would fail anyway. - """ - with self._lock: - self._shutdown = True +def is_ci_environment() -> bool: + """Detect CI/CD environments by checking well-known environment variables.""" + return any(_is_environment_signal_truthy(os.environ.get(var, "")) for var in _CI_ENV_VARS) - def __del__(self): - """Cleanup cache handler resources on garbage collection. - Safety net to ensure shutdown is called even if not done explicitly. - """ - try: - self.shutdown() - except Exception: - # Silently ignore errors during cleanup - pass +def is_telemetry_disabled_by_environment() -> bool: + """Return whether any supported environment variable requests full suppression.""" + return any( + os.environ.get(variable, "").strip().lower() in {"1", "true", "yes", "on", "y"} + for variable in ("ORT_DISABLE_TELEMETRY", "OLIVE_DISABLE_TELEMETRY") + ) - def on_payload_transmitted(self, args: "PayloadTransmittedCallbackArgs") -> None: - """Telemetry payload transmission callback. - Design decisions: - - Ignore callbacks during flush (unlikely to fail during successful flush) - - On success: flush cache if any cached events exist - - On failure: write to cache immediately (synchronous for simplicity) +class Telemetry: + """Per-process singleton that persists events to SQLite and uploads them. - Assumptions: - - Successful transmission indicates network is available to retry cached events - - If flush is in progress, we already successfully sent an event, so unlikely an event would suddenly fail - - Multiple concurrent successes don't need multiple flush operations - - Failed payloads should be cached immediately to avoid loss - """ - try: - payload = None - should_flush = False + Separate processes get separate in-memory singletons and coordinate only + through the shared SQLite store and its single-drainer file lock. + Use Telemetry() to get the singleton instance. + """ - with self._lock: - if self._shutdown: - return + _instance: Optional["Telemetry"] = None + _lock = threading.RLock() + _process_disabled = False + _heartbeat_enqueued = False + # Declared here (not assigned) so Pylint recognizes the per-instance attribute set in + # _new_unpublished_instance; every instance still gets its own dict there. + _global_metadata: dict[str, Any] + + @classmethod + def get_existing_instance(cls) -> Optional["Telemetry"]: + """Return the current singleton without creating telemetry.""" + return cls._instance - # Skip callbacks from replayed events during flush - # If a flush is in progress it means we successfully sent an event, - # so it's unlikely that an event would suddenly fail and need to be cached - # and we don't need to flush again. - if self._is_flushing: - with self._callback_condition: - self._callbacks_item_count += args.item_count - self._callback_condition.notify_all() - return + @classmethod + def get_or_create_if_enabled(cls) -> Optional["Telemetry"]: + """Return the singleton only when process telemetry is not fully disabled.""" + with cls._lock: + if cls._process_disabled or is_telemetry_disabled_by_environment(): + cls._process_disabled = True + if cls._instance is not None: + cls._instance.disable_telemetry() + return None + return cls() - if args.succeeded: - # Only flush if cache exists and no flush is in progress - cache_path = self.cache_path - if cache_path and cache_path.exists(): - should_flush = True - else: - payload = args.payload_bytes - - if should_flush: - # Release lock before scheduling (flush runs in separate thread) - self._schedule_flush() - elif payload: - # Write synchronously - JSON operations are fast enough - self._write_payload_to_cache(payload) - except Exception: - # Fail silently - telemetry should never crash the application - pass - finally: - with self._callback_condition: - self._callbacks_item_count += args.item_count - self._callback_condition.notify_all() - - def wait_for_callbacks(self, timeout_sec: float, during_flush: bool = False) -> bool: - deadline = time.time() + timeout_sec - while True: - with self._callback_condition: - callbacks_item_count = self._callbacks_item_count - expected_items = self._events_logged - if (during_flush or not self.is_flushing) and callbacks_item_count >= expected_items: - return True - remaining = deadline - time.time() - if remaining <= 0: - return False - with self._callback_condition: - self._callback_condition.wait(timeout=remaining) - - def record_event_logged(self, count: int = 1) -> None: - with self._callback_condition: - self._events_logged += count - - def _schedule_flush(self) -> None: - """Schedule cache flush in a separate thread to avoid blocking the callback. - - Design decisions: - - Check _is_flushing before spawning thread to avoid unnecessary threads - - Run flush in daemon thread (don't block process exit) - - Acquire lock at start to set _is_flushing flag atomically - - Always clear _is_flushing flag even if flush fails - - Assumptions: - - Flush operations are slow (network I/O) and should not block callbacks - - Daemon thread is acceptable (flush is best-effort) - """ - # Check before spawning thread to avoid unnecessary thread creation + def __new__(cls): + """Create or return the singleton instance.""" + with cls._lock: + if cls._process_disabled or is_telemetry_disabled_by_environment(): + cls._process_disabled = True + if cls._instance is not None: + cls._instance.disable_telemetry() + return cls._instance + return cls._new_unpublished_instance(disabled=True, initialized=True) + if cls._instance is None: + cls._instance = cls._new_unpublished_instance(disabled=False, initialized=False) + return cls._instance + + @classmethod + def _new_unpublished_instance(cls, *, disabled: bool, initialized: bool) -> "Telemetry": + instance = super().__new__(cls) + instance._initialized = initialized + instance._disabled = disabled + instance._store = None + instance._uploader = None + instance._enabled = not disabled + instance._recipe_only_ci_telemetry = False + instance._global_metadata = {} + instance._instrumentation_key = "" + instance._envelope_ikey = "" + instance._app_session_guid = "" + return instance + + def __init__(self): + """Initialize the telemetry store and uploader (runs once).""" with self._lock: - if self._shutdown or self._is_flushing: + if self._initialized: return - self._is_flushing = True - - def flush_task(): - try: - self._flush_cache() - except Exception: - # Fail silently - pass - finally: - # Always clear flag, even on exception - with self._lock: - self._is_flushing = False - - thread = threading.Thread(target=flush_task, daemon=True) - thread.start() - - @property - def cache_path(self) -> Optional[Path]: - """Get the path to the telemetry cache file. + # Mark initialized under the lock before doing any work, so two + # threads whose first Telemetry() calls interleave cannot both run + # the body (which would create two uploaders and two heartbeats). + self._initialized = True - Returns: - Optional[Path]: Path to cache file, or None if base directory unavailable. + self._enabled = True + self._recipe_only_ci_telemetry = False - """ - telemetry_cache_dir = None - if "OLIVE_TELEMETRY_CACHE_DIR" in os.environ: - telemetry_cache_dir = os.environ["OLIVE_TELEMETRY_CACHE_DIR"] - if not telemetry_cache_dir: - telemetry_cache_dir = get_telemetry_base_dir() / "cache" - return telemetry_cache_dir / self._cache_file_name - - def _write_payload_to_cache(self, payload: bytes) -> None: - """Write failed telemetry payload to cache for later retry. - - Design decisions: - - Parse payload to extract individual events (allows filtering) - - Filter to only critical events near size limit (preserves important data) - - Use file locking for multi-process safety (prevents corruption) - - Use exponential backoff for file contention (avoids spinning) - - Fail silently on errors (telemetry should never crash app) - - Assumptions: - - JSON operations are fast enough for synchronous execution - - File contention is rare and transient (retry a few times) - - Cache size limits prevent unbounded growth - - Critical events (heartbeat) are more important than others - """ - try: - cache_path = self.cache_path - if cache_path is None: + if self._disabled or is_telemetry_disabled_by_environment(): + type(self)._process_disabled = True + self._disabled = True + self._enabled = False return - # Parse payload into individual events for filtering - entries = _parse_payload(payload) - if not entries: - return + self._recipe_only_ci_telemetry = is_ci_environment() + self._app_session_guid = str(uuid.uuid4()) - cache_path.parent.mkdir(parents=True, exist_ok=True) - - max_retries = 3 - for attempt in range(max_retries + 1): - try: - cache_size = cache_path.stat().st_size if cache_path.exists() else 0 - - # Hard limit: stop caching entirely to prevent unbounded growth - if cache_size >= HARD_MAX_CACHE_SIZE_BYTES: - return - - # Soft limit: keep only critical events to preserve space - if cache_size >= MAX_CACHE_SIZE_BYTES: - entries = [entry for entry in entries if entry["event_name"] in CRITICAL_EVENTS] - if not entries: - return - - # Append base64-encoded newline-delimited entries - # Use exclusive file lock for multi-process safety - with _exclusive_file_lock(cache_path, mode="a") as cache_file: - for entry in entries: - plain = json.dumps(entry, ensure_ascii=False, separators=(",", ":")) - cache_file.write(_encode_cache_line(plain) + "\n") - return - except OSError as exc: - # Retry only on transient access errors (file locked by another process) - if exc.errno not in {errno.EACCES, errno.EAGAIN, errno.EWOULDBLOCK, errno.EBUSY}: - return - if attempt >= max_retries: - return - # Exponential backoff: 50ms, 100ms, 200ms (aligned with C# implementation) - time.sleep(0.05 * (2**attempt)) - except Exception: - # Fail silently - telemetry errors should not crash the application - return - - def _flush_cache(self) -> None: - """Flush this process's cached events back to telemetry service.""" - cache_path = self.cache_path - if cache_path is None or not cache_path.exists(): - return - - self._flush_cache_file(cache_path) - - def _flush_cache_file(self, cache_path: Path) -> None: - """Flush cached events back to telemetry service. - - Approach: - 1. Atomically rename cache → .flush (claims ownership, prevents concurrent flushes) - 2. Read all events from .flush file - 3. Queue all events for sending via telemetry logger - 4. Force flush with 2-second timeout - 5. On success: delete .flush file - 6. On failure: restore .flush → cache for retry - - Multi-process coordination: - - `replace()` is atomic; only one process can successfully rename the cache file - - If another process already renamed it, we get FileNotFoundError and abort - - Stale .flush files from crashes are overwritten by the atomic rename - - Shutdown handling: - - If shutdown flag set during flush, restore cache before returning - - This preserves events even if callbacks don't fire during shutdown - - Callback behavior: - - Queued events trigger callbacks with success/failure - - Failed events are automatically re-cached via callbacks (unless shutting down) - - The _is_flushing flag prevents re-caching of replayed events during flush - """ - flush_path = None - try: - # Check shutdown before starting (under lock to prevent race) - with self._lock: - if self._shutdown: - return - - if not cache_path.exists(): - return - - # Atomically rename to .flush file to claim ownership - # Overwrite any stale .flush file from crashed process (C# pattern) - flush_path = cache_path.with_name(f"{cache_path.name}.flush") try: - # On Windows/POSIX, replace() overwrites existing files atomically - cache_path.replace(flush_path) - except FileNotFoundError: - # Cache already claimed by another flush or doesn't exist - return - - # Read all cached entries (base64-decoded) - entries = _read_cache_entries(flush_path) - - if not entries: - # Empty cache, just delete the flush file - flush_path.unlink(missing_ok=True) - return + options = OneCollectorExporterOptions( + connection_string=base64.b64decode( + "SW5zdHJ1bWVudGF0aW9uS2V5PTYyMTUwOTExZGMwMDRmYzliYjY3YmE5NjA2NDI3ZTU2LWVjNjFmOWFmLTVkN2EtNGQxOS1hZjMxLWI5Y2Q2OWU5ODdmMS02OTE1" + ).decode() + ) + options.validate() + self._instrumentation_key = options.instrumentation_key + self._envelope_ikey = ( + f"{CommonSchemaJsonSerializationHelper.ONE_COLLECTOR_TENANCY_SYMBOL}:{options.tenant_token}" + ) - # Replay all events through telemetry logger - # Note: _is_flushing flag (set by caller) prevents these callbacks from re-caching or triggering nested flushes - # (unlikely since we just successfully sent an event, indicating network is available) - for entry in entries: - try: - event_name = entry["event_name"] - event_data = entry["event_data"] - if not event_name or not event_data: - continue - attributes = json.loads(event_data) - if not isinstance(attributes, dict): - continue - # Preserve original timestamp - attributes["initTs"] = entry.get("initTs", entry["ts"]) - self._telemetry.log(event_name, attributes, None) - except Exception: - # Skip malformed entries - continue - - # Check if shutdown happened during flush - with self._lock: - if self._shutdown: - # Restore cache to avoid data loss during shutdown - if flush_path and flush_path.exists(): - try: - cache_path.parent.mkdir(parents=True, exist_ok=True) - flush_path.replace(cache_path) - except Exception: - # Silently ignore errors during cleanup - pass + # Durable on-disk queue + background uploader. The uploader + # retries detailed events until delivery. CI has a separate + # recipe-only queue so it cannot drain local action/error rows. + db_file_name = CI_DB_FILE_NAME if self._recipe_only_ci_telemetry else DB_FILE_NAME + db_path = os.path.join(get_telemetry_base_dir(), db_file_name) + self._store = OfflineEventStore(db_path) + if not self._store.is_open: + self._store = None + self._enabled = False + self._initialized = False return - - # Wait for in-flight callbacks to complete before deciding success/failure - flush_success = self.wait_for_callbacks(timeout_sec=5.0, during_flush=True) - if flush_success: - # Success: delete the flush file (events were sent) - if flush_path: - flush_path.unlink(missing_ok=True) - elif flush_path and flush_path.exists(): - # Failure: restore cache for retry later - cache_path.parent.mkdir(parents=True, exist_ok=True) - flush_path.replace(cache_path) - except Exception: - # Best-effort restore on any exception to prevent data loss - if flush_path and flush_path.exists(): - try: - cache_path.parent.mkdir(parents=True, exist_ok=True) - flush_path.replace(cache_path) - except Exception: - # If restore fails, we lose the data (acceptable for telemetry) - pass - return - - @property - def is_flushing(self) -> bool: - with self._lock: - return self._is_flushing - - -class Telemetry: - """Wrapper that wires environment configuration into the library logger. - - This is a singleton class - all instances share the same state. - Use Telemetry() to get the singleton instance. - """ - - _instance: Optional["Telemetry"] = None - _lock = threading.Lock() - - def __new__(cls): - """Create or return the singleton instance. - - Thread-safe singleton implementation using double-checked locking. - """ - if cls._instance is None: - with cls._lock: - # Double-check pattern to prevent race conditions - if cls._instance is None: - instance = super().__new__(cls) - instance._initialized = False - cls._instance = instance - return cls._instance - - def __init__(self): - """Initialize the telemetry logger (only runs once for singleton).""" - # Prevent re-initialization - if self._initialized: + self._uploader = EventUploader(self._store, instrumentation_key=self._instrumentation_key) + if not self._recipe_only_ci_telemetry: + self._enqueue_heartbeat_once() + self._uploader.start() + except Exception: + # Fail silently — telemetry must never crash the host application + if self._store is not None: + self._store.close() + self._store = None + self._uploader = None + self._enabled = False + self._initialized = False + + def _enqueue_heartbeat_once(self) -> None: + """Reserve, enrich, and release this process's durable Heartbeat.""" + if type(self)._heartbeat_enqueued or self._store is None: return - - self._logger = None - self._cache_handler = None - try: - self._logger = self._create_logger() - event_source.disable() - - self._cache_handler = TelemetryCacheHandler(self) - self._setup_payload_callbacks() - if self._is_ci_environment(): - self.disable_telemetry() - self._initialized = True + device_id, device_id_status = get_hashed_device_id_and_status() + minimal_payload = self._build_payload( + HEARTBEAT_EVENT_NAME, + { + "device_id": device_id, + "device_id_status": device_id_status.value, + }, + ) + if minimal_payload is None: return - self._log_heartbeat() - if os.environ.get("OLIVE_DISABLE_TELEMETRY") == "1": - self.disable_telemetry() - self._initialized = True - except Exception: - # Fail silently — telemetry must never crash the host application - self._initialized = True - - @staticmethod - def _is_ci_environment() -> bool: - """Detect CI/CD environments by checking well-known environment variables.""" - return any(os.environ.get(var) for var in _CI_ENV_VARS) - - def _create_logger(self) -> Optional[TelemetryLogger]: - try: - return get_telemetry_logger(base64.b64decode(CONNECTION_STRING).decode()) + row_id = self._store.reserve(minimal_payload, _HEARTBEAT_RELEASE_SECONDS) + if row_id is None: + return + type(self)._heartbeat_enqueued = True + try: + full_payload = self._build_payload( + HEARTBEAT_EVENT_NAME, + { + "device_id": device_id, + "device_id_status": device_id_status.value, + "os": platform.system(), + "os_version": platform.version(), + "os_release": platform.release(), + "os_arch": platform.machine(), + }, + ) + if full_payload is not None and self._store.release(row_id, full_payload): + return + except Exception: + # The durable minimal heartbeat remains valid when enrichment fails. + pass + # If enrichment fails, release the already-durable minimal event. + self._store.release(row_id) except Exception: - return None - - def _setup_payload_callbacks(self) -> None: - # Register callback for payload transmission events - # No need to store unregister function - logger shutdown will clean up callbacks - if self._logger is None: - return - self._logger.register_payload_transmitted_callback( - self._cache_handler.on_payload_transmitted, - include_failures=True, - ) + # Heartbeat collection is best-effort and must not affect the host. + pass def add_global_metadata(self, metadata: dict[str, Any]) -> None: - """Add metadata to all telemetry events. - - Args: - metadata: Dictionary of metadata key-value pairs to add to all events. - These will be included in every telemetry event sent. - - Example: - >>> telemetry = Telemetry() - >>> telemetry.add_global_metadata({"user_id": "12345", "environment": "production"}) - - """ + """Merge metadata into every subsequent telemetry event.""" try: - if self._logger is None: - return - self._logger.add_global_metadata(metadata) + if metadata: + self._global_metadata = {**self._global_metadata, **metadata} except Exception: - # Fail silently — telemetry must never crash the host application pass + @property + def accepts_detailed_events(self) -> bool: + """Whether action and error events can currently be persisted.""" + return bool( + self._enabled and not self._recipe_only_ci_telemetry and self._store is not None and self._store.is_open + ) + def log( self, event_name: str, attributes: Optional[dict[str, Any]] = None, metadata: Optional[dict[str, Any]] = None, ) -> None: - """Log a telemetry event. - - Args: - event_name: Name of the event to log (e.g., "UserLogin", "ModelTrained"). - attributes: Optional dictionary of event-specific attributes. - metadata: Optional dictionary of additional metadata to merge with attributes. - - Example: - >>> telemetry = Telemetry() - >>> telemetry.log("ModelOptimized", {"model_type": "bert", "duration_ms": 1500}) - - """ + """Log a telemetry event (persisted durably, uploaded in the background).""" try: - attrs = _merge_metadata(attributes, metadata) - if self._logger is None: - return - self._logger.log(event_name, attrs) - if self._cache_handler: - self._cache_handler.record_event_logged() + with self._lock: + if not self._enabled or self._store is None: + return + if self._recipe_only_ci_telemetry and event_name != RECIPE_EVENT_NAME: + return + payload = self._build_payload(event_name, attributes, metadata) + if payload is None: + return + self._store.store(payload) + if self._uploader is not None: + self._uploader.request_drain() except Exception: # Fail silently — telemetry must never crash the host application pass - def _log_heartbeat( + def _build_payload( self, + event_name: str, + attributes: Optional[dict[str, Any]], metadata: Optional[dict[str, Any]] = None, - ) -> None: - """Log a heartbeat event with system information. - - Args: - metadata: Optional additional metadata to include. + ) -> Optional[bytes]: + """Merge metadata, filter to whitelisted keys, and serialize one event. + Returns the Common Schema JSON bytes, or None if the event is not + whitelisted or filters to nothing. """ - try: - encrypted_device_id, device_id_status = get_encrypted_device_id_and_status() - attributes = { - "device_id": encrypted_device_id, - "id_status": device_id_status.value, - "os": { - "name": platform.system().lower(), - "version": platform.version(), - "release": platform.release(), - "arch": platform.machine(), - }, + attrs = _merge_metadata(attributes, metadata) + if self._global_metadata: + attrs = {**self._global_metadata, **attrs} + filtered = _filter_event_data(event_name, attrs) + if filtered is None or not filtered: + # Unknown/empty event: not whitelisted. + return None + event_data = dict(filtered) + event_data.update( + { + "appName": "Olive", + "LibraryVersion": VERSION, + "AppSessionGuid": self._app_session_guid, } - self.log(HEARTBEAT_EVENT_NAME, attributes, metadata) - except Exception: - # Fail silently — telemetry must never crash the host application - pass + ) + serialized_snapshots = {} + for snapshot_field in ("configOverrides", "packageConfigOverrides"): + snapshot = event_data.get(snapshot_field) + if not isinstance(snapshot, str): + continue + event_data.pop(snapshot_field) + try: + parsed_snapshot = json.loads(snapshot) + except (TypeError, ValueError): + continue + scrubbed_snapshot = scrub_config_snapshot_for_telemetry(parsed_snapshot) + try: + serialized_snapshot = json.dumps( + scrubbed_snapshot, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + except ValueError: + serialized_snapshot = '{"truncated":"[truncated]"}' + if len(serialized_snapshot.encode("utf-8")) > MAX_TELEMETRY_STRING_LENGTH: + serialized_snapshot = '{"truncated":"[truncated]"}' + serialized_snapshots[snapshot_field] = serialized_snapshot + event_data = { + key: value + for key, value in event_data.items() + if value is not None and isinstance(value, (str, bytes, bytearray, bool, int, float, datetime, os.PathLike)) + } + scrubbed = scrub_value_for_telemetry(event_data) + if not isinstance(scrubbed, dict): + return None + scrubbed.update(serialized_snapshots) + exception_message = event_data.get("exceptionMessage") + if event_name == ERROR_EVENT_NAME and isinstance(exception_message, str): + scrubbed["exceptionMessage"] = scrub_error_message_for_telemetry(exception_message) + envelope = CommonSchemaJsonSerializationHelper.create_event_envelope( + event_name=event_name, + timestamp=datetime.now(timezone.utc), + ikey=self._envelope_ikey, + data=scrubbed, + ) + return CommonSchemaJsonSerializationHelper.serialize_to_json_bytes(envelope) def disable_telemetry(self) -> None: - """Disable all telemetry logging. - - After calling this method, no telemetry events will be sent until - telemetry is explicitly re-enabled. + """Fully disable telemetry for the remainder of this process.""" + with self._lock: + type(self)._process_disabled = True + self._disabled = True + self._enabled = False + uploader_stopped = True + if self._uploader is not None: + self._uploader.retain_queued_rows() + uploader_stopped = self._uploader.stop_loop(0) + if self._uploader is not None and uploader_stopped: + self._uploader.close() + self._uploader = None + if self._uploader is None and self._store is not None: + self._store.close() + self._store = None + + @classmethod + def disable_process_telemetry(cls) -> None: + """Latch full suppression without constructing telemetry resources.""" + with cls._lock: + cls._process_disabled = True + if cls._instance is not None: + cls._instance.disable_telemetry() + + @classmethod + def is_process_telemetry_disabled(cls) -> bool: + """Return whether environment or API state fully disables this process.""" + with cls._lock: + return cls._process_disabled or is_telemetry_disabled_by_environment() + + def shutdown(self, flush: bool = False) -> None: + """Stop the background uploader within a two-second total budget. + + Durable local events do not need an exit-time flush. Callers whose + telemetry store is ephemeral can request a best-effort flush. CI recipe + telemetry is always flushed because its uploader does not run in the + background. """ - try: - if self._logger is None: - return - self._logger.disable_telemetry() - except Exception: - # Fail silently — telemetry must never crash the host application - pass - - def shutdown(self, timeout_millis: float = 10_000, callback_timeout_millis: float = 2_000) -> None: - """Shutdown telemetry and flush pending events. + self._shutdown(flush, _SHUTDOWN_TIMEOUT_SECONDS) - Shutdown sequence: - 1. Wait for in-flight flush to complete (up to 1 second) - 2. Wait for callbacks + signal shutdown to cache handler - 3. Shutdown logger (cleans up callbacks automatically) - """ + def _shutdown(self, flush: bool, timeout_seconds: float) -> None: try: - # Step 1: Wait for pending flush to complete (matches C# 1-second timeout) - start_time = time.time() - while time.time() - start_time < 1.0: - if not self._cache_handler or not self._cache_handler.is_flushing: - break - time.sleep(0.05) - - # Step 2: Wait for callbacks/flush to complete before shutting down cache handler - if self._cache_handler: - # Nothing can be done if callbacks don't complete in time, so we ignore the result - _ = self._cache_handler.wait_for_callbacks(callback_timeout_millis / 1000) - self._cache_handler.shutdown() - - # Step 3: Shutdown logger (callbacks cleaned up automatically) - if self._logger is not None: - self._logger.shutdown() + timeout_seconds = max(0.0, timeout_seconds) + disabled = bool(getattr(self, "_disabled", False)) + flush = flush or bool(getattr(self, "_recipe_only_ci_telemetry", False)) + deadline = time.monotonic() + timeout_seconds + + def remaining_seconds() -> float: + return max(0.0, deadline - time.monotonic()) + + uploader_stopped = True + if self._uploader is not None: + uploader_stopped = self._uploader.stop_loop( + join_timeout_seconds=0 if disabled else min(timeout_seconds, remaining_seconds()) + ) + if uploader_stopped: + if flush and not disabled: + flush_timeout = remaining_seconds() + if flush_timeout > 0: + self._uploader.flush(flush_timeout) + self._uploader.close() + self._uploader = None + if self._store is not None and uploader_stopped: + self._store.close() + self._store = None + if self._uploader is None and self._store is None: + self._initialized = False except Exception: # Fail silently — telemetry must never crash the host application pass def __del__(self): - """Cleanup telemetry resources on garbage collection. - - This is a safety net to ensure resources are cleaned up even if - shutdown() is not explicitly called. However, relying on __del__ - is not recommended - always call shutdown() explicitly when done. - """ + """Safety-net cleanup on garbage collection.""" try: - self.shutdown() + self._shutdown(flush=False, timeout_seconds=0) except Exception: - # Silently ignore errors during cleanup pass -def _get_logger() -> Telemetry: - """Get or create the singleton Telemetry instance.""" - return Telemetry() - - -def _merge_metadata(attributes: Optional[dict[str, Any]], metadata: Optional[dict[str, Any]]) -> dict[str, Any]: - merged = dict(attributes or {}) - if metadata: - merged.update(metadata) - return merged +def _get_logger() -> Optional[Telemetry]: + """Get or create telemetry without publishing a disabled singleton.""" + return Telemetry.get_or_create_if_enabled() -def _parse_payload(payload: bytes) -> list[dict[str, Any]]: - """Parse telemetry payload into individual event entries. +def disable_telemetry() -> None: + """Fully disable telemetry for the remainder of this process.""" + Telemetry.disable_process_telemetry() - Design decisions: - - Filter events to only allowed keys (privacy/security) - - Store as minimal JSON (reduces cache size) - - Fail silently on malformed data (telemetry should be robust) - - Assumptions: - - Payload is newline-delimited JSON (OneCollector format) - - Events have "name", "time", and "data" fields - - Only whitelisted events and fields should be cached - """ - entries = [] - try: - payload_text = payload.decode("utf-8") - lines = payload_text.splitlines() - - for raw_line in lines: - line = raw_line.strip() - if not line: - continue - try: - event = json.loads(line) - event_name = event["name"] - if not event_name: - continue - # Filter to only allowed keys for privacy/security - filtered_data = _filter_event_data(event_name, event["data"]) - if not filtered_data: - continue - entries.append( - { - "ts": event["time"] or time.time(), - "event_name": event_name, - # Compact JSON to reduce cache size - "event_data": json.dumps(filtered_data, ensure_ascii=False, separators=(",", ":")), - } - ) - except Exception: - # Skip malformed lines - continue - except Exception: - # If entire payload is malformed, return empty list - return [] - return entries +def _merge_metadata(attributes: Optional[dict[str, Any]], metadata: Optional[dict[str, Any]]) -> dict[str, Any]: + merged = dict(metadata or {}) + if attributes: + merged.update(attributes) + return merged def _filter_event_data(event_name: str, data: dict[str, Any]) -> Optional[dict[str, Any]]: """Filter event data to only allowed keys for privacy/security. - Design decisions: - - Whitelist approach (only explicitly allowed keys are included) - - Support nested keys with dot notation (e.g., "os.name") - - Return None if no allowed keys found (filters out unknown events) - - Assumptions: - - ALLOWED_KEYS dict defines all cacheable events and their fields - - Unknown events should not be cached (privacy/security) + Whitelist approach: only explicitly allowed keys (with dot-notation support + for nested values, e.g. "os.name") are kept. Returns None for unknown events + so they are neither persisted nor sent. """ if event_name not in ALLOWED_KEYS: return None @@ -721,7 +562,7 @@ def _filter_event_data(event_name: str, data: dict[str, Any]) -> Optional[dict[s value = _get_nested_value(data, key) if value is None: continue - _set_nested_value(filtered, key, value) + _set_nested_value(filtered, FIELD_NAMES.get(key, key), value) return filtered or None @@ -740,37 +581,3 @@ def _set_nested_value(data: dict[str, Any], key: str, value: Any) -> None: for part in parts[:-1]: current = current.setdefault(part, {}) current[parts[-1]] = value - - -def _read_cache_entries(cache_path: Path) -> list[dict[str, Any]]: - """Read all entries from a cache file, decoding each line. - - Design decisions: - - Use file locking for multi-process safety - - Continue reading past malformed entries (partial data recovery) - - Return empty list on complete read failure (fail gracefully) - - Each line is base64-decoded before JSON parsing. - - Assumptions: - - Cache file contains newline-delimited base64-encoded entries (one per line) - - Each line is independent (one malformed line doesn't affect others) - - Empty or whitespace-only lines are skipped - """ - entries = [] - try: - with _exclusive_file_lock(cache_path, mode="r") as cache_file: - for raw_line in cache_file: - line = raw_line.strip() - if not line: - continue - try: - line = json.loads(_decode_cache_line(line)) - if isinstance(line, dict): - entries.append(line) - except Exception: - # Malformed line, skip and continue - continue - except Exception: - # If file cannot be opened or read, return empty list - return [] - return entries diff --git a/olive/telemetry/telemetry_extensions.py b/olive/telemetry/telemetry_extensions.py index e5b13395d0..e3e1e120b4 100644 --- a/olive/telemetry/telemetry_extensions.py +++ b/olive/telemetry/telemetry_extensions.py @@ -6,13 +6,16 @@ import functools import inspect import time +import traceback +from contextlib import suppress from types import TracebackType -from typing import Any, Callable, Optional, TypeVar +from typing import Any, Callable, Optional, TypeVar, cast -from olive.telemetry.telemetry import ACTION_EVENT_NAME, ERROR_EVENT_NAME, _get_logger -from olive.telemetry.utils import _format_exception_message +from olive.telemetry.telemetry import ACTION_EVENT_NAME, ERROR_EVENT_NAME, RECIPE_EVENT_NAME, _get_logger +from olive.telemetry.telemetry_redaction import scrub_error_message_for_telemetry _TFunc = TypeVar("_TFunc", bound=Callable[..., Any]) +_ERROR_LOGGED_ATTR = "_olive_telemetry_logged" def log_action( @@ -22,14 +25,17 @@ def log_action( success: bool, metadata: Optional[dict[str, Any]] = None, ) -> None: - telemetry = _get_logger() - attributes = { - "invoked_from": invoked_from, - "action_name": action_name, - "duration_ms": duration_ms, - "success": success, - } - telemetry.log(ACTION_EVENT_NAME, attributes, metadata) + with suppress(Exception): + telemetry = _get_logger() + if telemetry is None: + return + attributes = { + "invoked_from": invoked_from, + "action_name": action_name, + "duration_ms": duration_ms, + "success": success, + } + telemetry.log(ACTION_EVENT_NAME, attributes, metadata) def log_error( @@ -37,12 +43,63 @@ def log_error( exception_message: str, metadata: Optional[dict[str, Any]] = None, ) -> None: - telemetry = _get_logger() - attributes = { - "exception_type": exception_type, - "exception_message": exception_message, - } - telemetry.log(ERROR_EVENT_NAME, attributes, metadata) + with suppress(Exception): + telemetry = _get_logger() + if telemetry is None: + return + attributes = { + "exception_type": exception_type, + "exception_message": _redact_error_message(exception_message), + } + telemetry.log(ERROR_EVENT_NAME, attributes, metadata) + + +def log_recipe_result( + recipe_name: str, + success: bool, + metadata: Optional[dict[str, Any]] = None, +) -> None: + with suppress(Exception): + telemetry = _get_logger() + if telemetry is None: + return + attributes = { + "recipe_name": recipe_name, + "success": success, + } + telemetry.log(RECIPE_EVENT_NAME, attributes, metadata) + + +def _redact_error_message(text: str) -> str: + return scrub_error_message_for_telemetry(text) + + +def _is_exception_logged(exc: BaseException) -> bool: + return bool(getattr(exc, _ERROR_LOGGED_ATTR, False)) + + +def _mark_exception_logged(exc: BaseException) -> None: + try: + setattr(exc, _ERROR_LOGGED_ATTR, True) + except Exception: + # Some exception implementations do not allow custom attributes. + pass + + +def _format_exception_message(ex: BaseException, tb: Optional[TracebackType] = None) -> str: + """Format exception and frame metadata without collecting source-code lines.""" + lines = [] + if tb is not None: + lines.append("Traceback (most recent call last):") + for frame in traceback.extract_tb(tb, limit=5): + frame_name = _redact_error_message(frame.name) + lines.append(f'File "[path]", line {frame.lineno}, in {frame_name}') + try: + exception_message = str(ex) + except Exception: + exception_message = "" + lines.append(f"{type(ex).__name__}: {_redact_error_message(exception_message)}") + return _redact_error_message("\n".join(lines)) def _resolve_invoked_from(skip_frames: int = 0) -> str: @@ -69,6 +126,18 @@ def _resolve_invoked_from(skip_frames: int = 0) -> str: return "Interactive" +def _resolve_action_name(func: Callable[..., Any]) -> str: + action_name = getattr(func, "__name__", "unknown") + qualname = getattr(func, "__qualname__", action_name) + if "." not in qualname: + return action_name + owner = qualname.rsplit(".", 1)[0].rsplit(".", 1)[-1] + if owner == "": + return action_name + owner = owner[: -len("Command")] if owner.endswith("Command") else owner + return owner if action_name == "run" else f"{owner}.{action_name}" + + class ActionContext: """Context manager for recording telemetry around a block of work.""" @@ -79,7 +148,18 @@ def __init__( metadata: Optional[dict[str, Any]] = None, ): self.action_name = action_name - self.invoked_from = invoked_from if invoked_from is not None else _resolve_invoked_from() + try: + telemetry = _get_logger() + self._telemetry_enabled = bool(telemetry is not None and telemetry.accepts_detailed_events) + except Exception: + self._telemetry_enabled = False + self.invoked_from = ( + invoked_from + if invoked_from is not None + else _resolve_invoked_from() + if self._telemetry_enabled + else "disabled" + ) self.metadata = metadata or {} self._start_time: Optional[float] = None @@ -96,7 +176,11 @@ def __exit__( exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> bool: - duration_ms = int((time.perf_counter() - (self._start_time or time.perf_counter())) * 1000) + if not self._telemetry_enabled: + return False + end_time = time.perf_counter() + start_time = self._start_time if self._start_time is not None else end_time + duration_ms = int((end_time - start_time) * 1000) success = exc_type is None log_action( @@ -107,12 +191,13 @@ def __exit__( metadata=self.metadata, ) - if exc_type is not None and exc_val is not None: + if exc_type is not None and exc_val is not None and not _is_exception_logged(exc_val): log_error( exception_type=exc_type.__name__, exception_message=_format_exception_message(exc_val, exc_tb), metadata=self.metadata, ) + _mark_exception_logged(exc_val) # Do not suppress exceptions return False @@ -123,13 +208,21 @@ def action(func: _TFunc) -> _TFunc: @functools.wraps(func) def wrapper(*args: Any, **kwargs: Any): - invoked_from = _resolve_invoked_from() - action_name = func.__name__ - if args and hasattr(args[0], "__class__"): - cls_name = args[0].__class__.__name__ - cls_name = cls_name[: -len("Command")] if cls_name.endswith("Command") else cls_name - if cls_name: - action_name = cls_name if action_name == "run" else f"{cls_name}.{action_name}" + try: + telemetry = _get_logger() + if telemetry is None or not telemetry.accepts_detailed_events: + return func(*args, **kwargs) + except Exception: + return func(*args, **kwargs) + + # Resolve telemetry context defensively: instrumentation (including + # inspect.stack()) must never propagate into the wrapped call. + try: + invoked_from = _resolve_invoked_from() + action_name = _resolve_action_name(func) + except Exception: + invoked_from = "unknown" + action_name = getattr(func, "__name__", "unknown") start_time = time.perf_counter() success = True @@ -137,10 +230,12 @@ def wrapper(*args: Any, **kwargs: Any): return func(*args, **kwargs) except Exception as exc: success = False - log_error( - exception_type=type(exc).__name__, - exception_message=_format_exception_message(exc, exc.__traceback__), - ) + if not _is_exception_logged(exc): + log_error( + exception_type=type(exc).__name__, + exception_message=_format_exception_message(exc, exc.__traceback__), + ) + _mark_exception_logged(exc) raise finally: duration_ms = int((time.perf_counter() - start_time) * 1000) @@ -151,4 +246,4 @@ def wrapper(*args: Any, **kwargs: Any): success=success, ) - return wrapper # type: ignore[return-value] + return cast("_TFunc", wrapper) diff --git a/olive/telemetry/telemetry_redaction.py b/olive/telemetry/telemetry_redaction.py new file mode 100644 index 0000000000..afe54fe7bc --- /dev/null +++ b/olive/telemetry/telemetry_redaction.py @@ -0,0 +1,434 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Free-text telemetry redaction.""" + +import os +import re +from collections.abc import Mapping +from collections.abc import Set as AbstractSet +from datetime import date, datetime, time, timedelta +from uuid import UUID + +MAX_TELEMETRY_STRING_LENGTH = 40_960 +MAX_ERROR_MESSAGE_LENGTH = MAX_TELEMETRY_STRING_LENGTH + +_SENSITIVE_COMPACT_KEYS = { + "accountkey", + "accesskey", + "accesstoken", + "apikey", + "auth", + "authorization", + "authtoken", + "clientsecret", + "connectionstring", + "credential", + "credentials", + "passwd", + "password", + "privatekey", + "pwd", + "sastoken", + "secret", + "secretkey", + "servicecredential", + "sig", + "signature", + "subscriptionkey", + "token", +} +_ENVIRONMENT_COMPACT_KEYS = {"env", "environment", "environmentvariables", "environmentvars", "envvariables", "envvars"} +_PATH_KEY_SUFFIXES = ("dir", "dirs", "file", "files", "path", "paths") + + +def normalize_config_key_for_telemetry(key) -> str: + if key is None: + return "" + value = str(key) + value = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", value) + value = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", value) + return re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_") + + +def is_sensitive_config_key_for_telemetry(key) -> bool: + normalized = normalize_config_key_for_telemetry(key) + compact = normalized.replace("_", "") + parts = set(normalized.split("_")) + if compact in _SENSITIVE_COMPACT_KEYS: + return True + if parts & {"auth", "authorization", "credential", "credentials", "passwd", "password", "pwd", "secret", "token"}: + return True + if "key" in parts and parts & {"access", "account", "api", "private", "secret", "subscription"}: + return True + if {"connection", "string"} <= parts: + return True + if compact.endswith( + ("authorization", "credential", "credentials", "passwd", "password", "pwd", "secret", "signature", "token") + ): + return True + return any( + marker in compact + for marker in ( + "accesskey", + "accountkey", + "apikey", + "connectionstring", + "privatekey", + "secretkey", + "subscriptionkey", + ) + ) + + +def is_environment_config_key_for_telemetry(key) -> bool: + normalized = normalize_config_key_for_telemetry(key) + return normalized.replace("_", "") in _ENVIRONMENT_COMPACT_KEYS or bool( + set(normalized.split("_")) & {"env", "environment"} + ) + + +def is_path_like_config_key_for_telemetry(key) -> bool: + compact = normalize_config_key_for_telemetry(key).replace("_", "") + return bool(compact) and compact.endswith(_PATH_KEY_SUFFIXES) + + +def scrub_config_snapshot_for_telemetry(value, key=None): + """Recursively scrub a JSON-compatible configuration snapshot by key and value.""" + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return "" + if is_path_like_config_key_for_telemetry(key): + return "" + if isinstance(value, Mapping): + items = {} + collisions = set() + for child_key, child_value in value.items(): + safe_key = scrub_string_for_telemetry(str(child_key)) + if not safe_key: + continue + if safe_key in items: + collisions.add(safe_key) + else: + items[safe_key] = scrub_config_snapshot_for_telemetry(child_value, child_key) + return {safe_key: items[safe_key] for safe_key in sorted(items) if safe_key not in collisions} + if isinstance(value, (list, tuple)): + return [scrub_config_snapshot_for_telemetry(item, key) for item in value] + if isinstance(value, str): + return scrub_string_for_telemetry(value) + return value + + +def _token_start(value: str, index: int) -> int: + while index > 0 and not value[index - 1].isspace() and value[index - 1] not in "\"'": + index -= 1 + return index + + +def _is_drive_path_anchor(value: str, index: int) -> bool: + char = value[index] + if not char.isascii() or not char.isalpha(): + return False + if index > 0 and value[index - 1] not in "\"' \t=([{,;": + return False + if index + 2 >= len(value): + return False + return value[index + 1] == ":" and value[index + 2] in "/\\" + + +def _is_sensitive_slash_option(value: str, index: int) -> bool: + key_start = index + 1 + if key_start >= len(value) or not value[key_start].isascii() or not value[key_start].isalpha(): + return False + key_end = key_start + while key_end < len(value) and value[key_end].isascii() and (value[key_end].isalnum() or value[key_end] in "_.-"): + key_end += 1 + if key_end == key_start or not is_sensitive_config_key_for_telemetry(value[key_start:key_end]): + return False + separator = key_end + while separator < len(value) and value[separator].isspace(): + separator += 1 + return separator < len(value) and (value[separator] in "=:" or separator > key_end) + + +def _find_path_anchor(value: str): + index = 0 + slash_token_end = 0 + slash_token_start = 0 + slash_token_analyzed = False + relative_slash_anchor = None + while index < len(value): + char = value[index] + if value.startswith(("./", "../", ".\\", "..\\"), index) and ( + index == 0 or value[index - 1].isspace() or value[index - 1] in "\"'=([{,;" + ): + return index + if char == ":" and index + 2 < len(value) and value[index + 1 : index + 3] == "//": + scheme_start = index + while scheme_start > 0 and ( + value[scheme_start - 1].isascii() + and (value[scheme_start - 1].isalnum() or value[scheme_start - 1] in "+-.") + ): + scheme_start -= 1 + return scheme_start + if char == "\\" and index + 1 < len(value) and value[index + 1] == "\\": + return index + if char == "~" and index + 1 < len(value) and value[index + 1] in "/\\": + return index + if _is_drive_path_anchor(value, index): + return index + if char == "\\": + if ( + index + 1 < len(value) + and value[index + 1] not in "\\\r\n \t" + and (index == 0 or value[index - 1] in "\"' \t=([{,;") + ): + return index + separators = 0 + for candidate in value[index:]: + if candidate in "\r\n": + break + if candidate == "\\": + separators += 1 + if separators >= 2: + return _token_start(value, index) + if char == "/": + if index >= slash_token_end: + slash_token_start = _token_start(value, index) + slash_token_end = index + while ( + slash_token_end < len(value) + and not value[slash_token_end].isspace() + and value[slash_token_end] not in "\"'" + ): + slash_token_end += 1 + slash_token_analyzed = False + if ( + index + 1 < len(value) + and value[index + 1] not in "/\r\n \t" + and (index == 0 or value[index - 1] in "\"' \t=([{,;") + and not _is_sensitive_slash_option(value, index) + ): + return index + if not slash_token_analyzed: + slash_token_analyzed = True + segments = 0 + cursor = index + while cursor < slash_token_end and value[cursor] == "/": + separator_end = cursor + 1 + while separator_end < slash_token_end and value[separator_end] == "/": + separator_end += 1 + cursor = separator_end + segment_start = cursor + while cursor < slash_token_end and value[cursor] not in "/\r\n \t": + cursor += 1 + if cursor == segment_start: + break + segments += 1 + if segments >= 2: + return relative_slash_anchor if relative_slash_anchor is not None else slash_token_start + if segments == 1 and slash_token_start < index: + if relative_slash_anchor is not None: + return relative_slash_anchor + token = value[slash_token_start:slash_token_end] + if token.lower() not in {"and/or", "n/a", "read/write"} and any(char.isalpha() for char in token): + relative_slash_anchor = slash_token_start + index += 1 + return None + + +def _is_secret_key_char(char: str) -> bool: + return char.isascii() and (char.isalnum() or char in "_.-") + + +def _is_secret_key_boundary(char: str) -> bool: + return char.isspace() or char in "?&#;,\"'([{/-" + + +def _find_sensitive_value_anchor(value: str): + index = 0 + while index < len(value): + char = value[index] + if not char.isascii() or not char.isalpha(): + index += 1 + continue + if index > 0 and not _is_secret_key_boundary(value[index - 1]): + index += 1 + continue + + key_end = index + 1 + while key_end < len(value) and _is_secret_key_char(value[key_end]): + key_end += 1 + if not is_sensitive_config_key_for_telemetry(value[index:key_end]): + index = key_end + continue + + separator = key_end + if separator < len(value) and value[separator] in "\"'": + separator += 1 + before_whitespace = separator + while separator < len(value) and value[separator].isspace(): + separator += 1 + assignment = separator < len(value) and value[separator] in "=:" + cli_option = index > 0 and value[index - 1] in "-/" + delimited_cli_value = False + if cli_option and not assignment: + separator = key_end + while separator < len(value) and (value[separator].isspace() or value[separator] in "\"',[](){}"): + if value[separator] in "\"',[](){}": + delimited_cli_value = True + separator += 1 + separated_cli_value = cli_option and before_whitespace < separator < len(value) + separated_cli_value = separated_cli_value and (value[separator] != "-" or delimited_cli_value) + if not assignment and not separated_cli_value: + index = key_end + continue + + value_start = separator + 1 if assignment else separator + while value_start < len(value) and value[value_start].isspace(): + value_start += 1 + if value_start < len(value) and value[value_start] not in "&;\r\n": + return value_start + index = key_end + return None + + +def _is_user_info_terminator(char: str) -> bool: + return char.isspace() or char in '"\\/?#[]{}' + + +def _is_authority_terminator(char: str) -> bool: + return char.isspace() or char in "\"')},;/?#" + + +def _find_credential_url_anchor(value: str): + token_start = 0 + colon = None + index = 0 + while index < len(value): + char = value[index] + if colon is None: + if _is_user_info_terminator(char): + token_start = index + 1 + elif char == ":" and index > token_start: + colon = index + index += 1 + continue + + if _is_user_info_terminator(char): + token_start = index + 1 + colon = None + index += 1 + continue + if char != "@" or colon + 1 == index: + index += 1 + continue + + host_start = index + 1 + if host_start == len(value): + index += 1 + continue + if value[host_start] == "[": + host_end = value.find("]", host_start + 1) + if host_end > host_start + 1: + return token_start + index += 1 + continue + + host_end = host_start + while host_end < len(value) and not _is_authority_terminator(value[host_end]): + host_end += 1 + if host_end > host_start: + return token_start + index = host_end + 1 + token_start = index + colon = None + return None + + +def _truncate_utf8(value: str, max_bytes: int) -> str: + encoded = value.encode("utf-8") + if len(encoded) <= max_bytes: + return value + return encoded[:max_bytes].decode("utf-8", errors="ignore") + + +def _scrub_string_for_telemetry(value: str, max_bytes: int) -> str: + path_anchor = _find_path_anchor(value) + secret_value = _find_sensitive_value_anchor(value) + credential_anchor = _find_credential_url_anchor(value) + anchors = ( + (path_anchor, "[path]"), + (secret_value, "[secret]"), + (credential_anchor, "[secret]"), + ) + anchor, marker = min( + ((anchor, marker) for anchor, marker in anchors if anchor is not None), + default=(None, None), + ) + if anchor is not None: + return _truncate_utf8(value[:anchor], max_bytes - len(marker)) + marker + return _truncate_utf8(value, max_bytes) + + +def scrub_string_for_telemetry(value: str) -> str: + """Redact and cap a general telemetry string.""" + return _scrub_string_for_telemetry(value, MAX_TELEMETRY_STRING_LENGTH) + + +def scrub_error_message_for_telemetry(value: str) -> str: + """Redact and cap an error message at 40,960 UTF-8 bytes.""" + return _scrub_string_for_telemetry(value, MAX_ERROR_MESSAGE_LENGTH) + + +def scrub_value_for_telemetry(value, key=None): + """Recursively scrub strings and path-like values before serialization.""" + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return "" + if is_path_like_config_key_for_telemetry(key): + return "" + if isinstance(value, os.PathLike): + return "[path]" + if isinstance(value, str): + return scrub_string_for_telemetry(value) + if isinstance(value, (bytes, bytearray)): + try: + return scrub_string_for_telemetry(bytes(value).decode("utf-8")) + except UnicodeDecodeError: + return "[binary]" + if value is None or isinstance( + value, + (bool, int, float, datetime, date, time, timedelta, UUID), + ): + return value + if isinstance(value, Mapping): + items = {} + collisions = set() + for child_key, child in value.items(): + if isinstance(child_key, os.PathLike): + safe_key = "[path]" + elif isinstance(child_key, str): + safe_key = scrub_string_for_telemetry(child_key) + else: + try: + safe_key = scrub_string_for_telemetry(str(child_key)) + except Exception: + safe_key = f"[unsupported:{type(child_key).__name__}]" + if safe_key: + if safe_key in items: + collisions.add(safe_key) + else: + items[safe_key] = scrub_value_for_telemetry(child, child_key) + return {safe_key: items[safe_key] for safe_key in sorted(items) if safe_key not in collisions} + if isinstance(value, list): + return [scrub_value_for_telemetry(child, key) for child in value] + if isinstance(value, tuple): + return tuple(scrub_value_for_telemetry(child, key) for child in value) + if isinstance(value, AbstractSet): + children = [scrub_value_for_telemetry(child, key) for child in value] + return sorted(children, key=lambda child: (type(child).__name__, repr(child))) + try: + return scrub_string_for_telemetry(str(value)) + except Exception: + return f"[unsupported:{type(value).__name__}]" diff --git a/olive/telemetry/uploader.py b/olive/telemetry/uploader.py new file mode 100644 index 0000000000..2bedf8e507 --- /dev/null +++ b/olive/telemetry/uploader.py @@ -0,0 +1,343 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Background uploader that drains the SQLite offline store to OneCollector. + +Reads the oldest batch of events, POSTs them, and: +- deletes them on success (HTTP 2xx), +- deletes them on a permanent, non-retryable failure (e.g. HTTP 4xx) so a + poison event cannot block the queue forever, +- leaves them on a transient failure (network error / HTTP 5xx / timeout) to be + retried on the next cycle or the next process run. + +Durability is provided by the on-disk store, so the process can exit at any time +without losing events and without an exit-time flush. +""" + +import threading +import time +from enum import Enum +from typing import NamedTuple, Optional + +from olive.telemetry.library.options import CompressionType, OneCollectorTransportOptions +from olive.telemetry.library.payload_builder import PayloadBuilder +from olive.telemetry.library.transport import HttpJsonPostTransport +from olive.telemetry.offline_store import OfflineEventStore +from olive.telemetry.process_lock import ProcessDrainLock + + +class DrainOutcome(Enum): + EMPTY = "empty" + PROGRESS = "progress" + SPLIT = "split" + ACKNOWLEDGE_RETRY = "acknowledge_retry" + DELETE_RETRY = "delete_retry" + STORAGE_RETRY = "storage_retry" + TRANSPORT_RETRY = "transport_retry" + + +class DrainResult(NamedTuple): + """Rows handled includes both successful uploads and intentional permanent drops.""" + + handled: int + left: int + outcome: DrainOutcome + + +class EventUploader: + """Drains the offline store and ships events over HTTP on a daemon thread. + + When several processes share one telemetry database, a single-holder + advisory lock ensures only one of them drains at a time, so the same event + is never uploaded by two concurrent drainers. Processes that do not hold the + lock keep writing events durably; the holder drains them. + """ + + def __init__( + self, + store: OfflineEventStore, + instrumentation_key: str, + endpoint: str = OneCollectorTransportOptions.DEFAULT_ENDPOINT, + compression: CompressionType = CompressionType.DEFLATE, + drain_interval_seconds: float = 2.0, + max_items_per_drain: int = 256, + send_timeout_seconds: float = 10.0, + idle_backoff_seconds: float = 30.0, + ): + self._store = store + self._drain_interval = drain_interval_seconds + self._max_items = max_items_per_drain + self._send_timeout = send_timeout_seconds + self._idle_backoff = idle_backoff_seconds + self._drain_lock = ProcessDrainLock(store.db_path + ".lock") + + self._transport = HttpJsonPostTransport( + endpoint=endpoint, + ikey=instrumentation_key, + compression=compression, + ) + + self._wake = threading.Event() + self._stop = threading.Event() + self._retain_rows = threading.Event() + self._mutation_lock = threading.Lock() + self._thread: Optional[threading.Thread] = None + self._pending_ack_ids: list[int] = [] + self._pending_delete_ids: list[int] = [] + self._split_batch_size: Optional[int] = None + + # ----- control ------------------------------------------------------- + + def start(self) -> None: + if self._thread is not None: + return + self._thread = threading.Thread(target=self._run, name="olive-telemetry-uploader", daemon=True) + self._thread.start() + + def request_drain(self) -> None: + """Nudge this process when it is the designated drainer.""" + if self._drain_lock.held: + self._wake.set() + + def stop_loop(self, join_timeout_seconds: float = 5.0) -> bool: + """Stop the background loop and wait briefly for it to exit. + + Returns True if the thread actually stopped (so a caller may safely drain + as the sole drainer), False if it is still alive (e.g. stuck in an + in-flight send) — in which case the caller must NOT drain, to avoid + double-sending the rows the thread is still processing. + """ + self._stop.set() + self._wake.set() + thread = self._thread + if thread is None: + return True + thread.join(join_timeout_seconds) + stopped = not thread.is_alive() + if stopped: + self._thread = None + return stopped + + def signal_stop(self) -> None: + """Ask the loop to stop without blocking the caller. + + The daemon thread winds down on its next wake; the drain lock is released + when it exits (or by the OS at process exit). Use this for the opt-out + path so disabling telemetry never blocks the host. + """ + self._stop.set() + self._wake.set() + + def retain_queued_rows(self) -> None: + """Stop new draining and prevent any later queue deletion.""" + with self._mutation_lock: + self._retain_rows.set() + self.signal_stop() + + def close(self) -> None: + """Release the single-drainer lock (urllib holds no persistent connection).""" + self._drain_lock.release() + + def stop(self, timeout_seconds: float = 12.0) -> None: + """Stop the loop and release the drain lock (convenience).""" + if self.stop_loop(timeout_seconds): + self.close() + + # ----- draining ------------------------------------------------------ + + def _delete_unless_retained(self, ids: list[int], deadline: Optional[float] = None) -> Optional[bool]: + with self._mutation_lock: + if self._retain_rows.is_set(): + return None + return self._store.delete(ids, deadline) + + def _finish_handled_rows(self, ids: list[int], deadline: Optional[float] = None) -> DrainResult: + """Acknowledge and remove rows after upload or intentional permanent discard.""" + self._pending_ack_ids = ids + with self._mutation_lock: + if self._retain_rows.is_set(): + return DrainResult(0, len(ids), DrainOutcome.TRANSPORT_RETRY) + if not self._store.acknowledge(ids, deadline): + return DrainResult(0, len(ids), DrainOutcome.ACKNOWLEDGE_RETRY) + self._pending_ack_ids = [] + self._pending_delete_ids = ids + deleted = self._store.delete(ids, deadline) + if not deleted: + return DrainResult(0, len(ids), DrainOutcome.DELETE_RETRY) + self._pending_delete_ids = [] + return DrainResult(len(ids), 0, DrainOutcome.PROGRESS) + + def drain_once(self, deadline: Optional[float] = None) -> DrainResult: + """Attempt to upload one batch. + + The outcome distinguishes local work that should continue immediately + (poison isolation or acknowledged deletion) from a retryable transport + failure that should back off. + """ + if self._retain_rows.is_set(): + return DrainResult( + 0, + len(self._pending_ack_ids) + len(self._pending_delete_ids), + DrainOutcome.TRANSPORT_RETRY, + ) + if self._pending_ack_ids: + return self._finish_handled_rows(self._pending_ack_ids, deadline) + if self._pending_delete_ids: + pending_ids = self._pending_delete_ids + deleted = self._delete_unless_retained(pending_ids, deadline) + if deleted is None: + return DrainResult(0, len(pending_ids), DrainOutcome.TRANSPORT_RETRY) + if not deleted: + return DrainResult(0, len(pending_ids), DrainOutcome.DELETE_RETRY) + self._pending_delete_ids = [] + return DrainResult(len(pending_ids), 0, DrainOutcome.PROGRESS) + + batch_limit = min(self._max_items, self._split_batch_size or self._max_items) + acknowledged_ids = self._store.get_acknowledged_ids(batch_limit, deadline) + if acknowledged_ids is None: + return DrainResult(0, 0, DrainOutcome.STORAGE_RETRY) + if acknowledged_ids: + self._pending_delete_ids = acknowledged_ids + deleted = self._delete_unless_retained(acknowledged_ids, deadline) + if deleted is None: + return DrainResult(0, len(acknowledged_ids), DrainOutcome.TRANSPORT_RETRY) + if not deleted: + return DrainResult(0, len(acknowledged_ids), DrainOutcome.DELETE_RETRY) + self._pending_delete_ids = [] + return DrainResult(len(acknowledged_ids), 0, DrainOutcome.PROGRESS) + + batch = self._store.get_batch_for_upload(batch_limit, deadline) + if batch is None: + return DrainResult(0, 0, DrainOutcome.STORAGE_RETRY) + if not batch: + return DrainResult(0, 0, DrainOutcome.EMPTY) + + builder = PayloadBuilder( + max_size_bytes=OneCollectorTransportOptions.DEFAULT_MAX_PAYLOAD_SIZE_BYTES, + max_items=OneCollectorTransportOptions.DEFAULT_MAX_ITEMS_PER_PAYLOAD, + ) + included: list[int] = [] + for row_id, payload in batch: + if not builder.can_add(payload): + if builder.is_empty: + return self._finish_handled_rows([row_id], deadline) + break + builder.add(payload) + included.append(row_id) + payload_bytes = builder.build() + + timeout = self._send_timeout + if deadline is not None: + timeout = min(timeout, max(0.0, deadline - time.monotonic())) + if timeout <= 0.0: + return DrainResult(0, len(included), DrainOutcome.TRANSPORT_RETRY) + admission_released = False + + def release_admission() -> None: + nonlocal admission_released + if not admission_released: + admission_released = True + self._mutation_lock.release() + + # A context manager cannot be used here: on_send_admitted (release_admission) + # releases the lock as soon as the transport has read the payload, well before + # send() returns. A `with` block would then try to release again at scope exit, + # double-releasing the lock. release_admission() guards against that with the + # admission_released flag, but only a manual acquire (not `with`) is compatible + # with a release that must happen earlier than block exit. + self._mutation_lock.acquire() # pylint: disable=consider-using-with + try: + if self._retain_rows.is_set(): + return DrainResult(0, len(included), DrainOutcome.TRANSPORT_RETRY) + success, status = self._transport.send( + payload_bytes, + timeout, + item_count=len(included), + on_send_admitted=release_admission, + ) + except Exception: + success, status = (False, None) + finally: + release_admission() + + if success: + self._split_batch_size = None + return self._finish_handled_rows(included, deadline) + if not HttpJsonPostTransport.is_retryable(status): + if status in {400, 413, 422} and len(included) > 1: + self._split_batch_size = max(1, len(included) // 2) + return DrainResult(0, len(included), DrainOutcome.SPLIT) + # Permanent rejection (e.g. 4xx): drop so it can't block the queue. + self._split_batch_size = None + return self._finish_handled_rows(included, deadline) + # Transient failure: leave the rows for the next attempt. + return DrainResult(0, len(included), DrainOutcome.TRANSPORT_RETRY) + + def flush(self, max_seconds: float = 5.0) -> None: + """Best-effort drain of all pending events, bounded by max_seconds. + + Only drains after the background loop stops and if this process can hold + the single-drainer lock; otherwise events remain durable for a later run. + """ + if self._thread is not None and self._thread.is_alive(): + return + if not self._drain_lock.acquire(): + return + try: + deadline = time.monotonic() + max(0.0, max_seconds) + delete_retries = 0 + while time.monotonic() < deadline: + result = self.drain_once(deadline) + if result.outcome is DrainOutcome.EMPTY: + return # queue empty + if result.outcome in {DrainOutcome.STORAGE_RETRY, DrainOutcome.TRANSPORT_RETRY}: + return # transient failure; leave the rest for next run + if result.outcome in {DrainOutcome.ACKNOWLEDGE_RETRY, DrainOutcome.DELETE_RETRY}: + delete_retries += 1 + if delete_retries >= 2: + return + else: + delete_retries = 0 + finally: + self._drain_lock.release() + + def _run(self) -> None: + try: + while not self._stop.is_set(): + self._wake.clear() + transient_failure = False + # Only one process drains at a time. If another holds the lock, skip + # draining this cycle; our events remain durable for the holder. + if self._drain_lock.acquire(): + try: + delete_retries = 0 + while not self._stop.is_set(): + result = self.drain_once() + if result.outcome is DrainOutcome.EMPTY: + break + if result.outcome is DrainOutcome.TRANSPORT_RETRY: + transient_failure = True + break + if result.outcome is DrainOutcome.STORAGE_RETRY: + transient_failure = True + break + if result.outcome in {DrainOutcome.ACKNOWLEDGE_RETRY, DrainOutcome.DELETE_RETRY}: + delete_retries += 1 + if delete_retries >= 2: + transient_failure = True + break + else: + delete_retries = 0 + except Exception: + transient_failure = True + else: + transient_failure = True + + wait = self._idle_backoff if transient_failure else self._drain_interval + self._wake.wait(wait) + finally: + # Release the single-drainer lock when the loop exits so another + # process can take over (also released by close()/OS on exit). + self._drain_lock.release() diff --git a/olive/telemetry/utils.py b/olive/telemetry/utils.py index 52a39acded..08c68ef63f 100644 --- a/olive/telemetry/utils.py +++ b/olive/telemetry/utils.py @@ -2,15 +2,10 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -import base64 import functools import os import platform -import tempfile -import traceback from pathlib import Path -from types import TracebackType -from typing import Optional ORT_SUPPORT_DIR = r"Microsoft/DeveloperTools/.onnxruntime" @@ -18,15 +13,26 @@ def _resolve_home_dir() -> Path: """Resolve the user home directory with fallbacks for container environments.""" home = os.getenv("HOME") - if home: - return Path(home).expanduser() + if home and Path(home).is_absolute(): + return Path(home) + if platform.system() != "Windows": + try: + import pwd + + passwd_home = Path(pwd.getpwuid(os.getuid()).pw_dir) + if passwd_home.is_absolute(): + return passwd_home + except (AttributeError, ImportError, KeyError, OSError): + # Fall through to pathlib's platform-independent home resolution. + pass try: - return Path.home() + fallback_home = Path.home() + if fallback_home.is_absolute(): + return fallback_home except (RuntimeError, KeyError): - # /var/tmp persists across reboots unlike /tmp (FHS spec) - if platform.system() != "Windows": - return Path("/var/tmp") - return Path(tempfile.gettempdir()) + # Neither OS account data nor pathlib could provide a usable home. + pass + raise RuntimeError("No absolute per-user telemetry storage directory is available") @functools.lru_cache(maxsize=1) @@ -34,9 +40,9 @@ def get_telemetry_base_dir() -> Path: os_name = platform.system() if os_name == "Windows": base_dir = os.environ.get("LOCALAPPDATA") or os.environ.get("APPDATA") - if not base_dir: - base_dir = str(Path.home() / "AppData" / "Local") - return Path(base_dir) / "Microsoft" / ".onnxruntime" + if base_dir and Path(base_dir).is_absolute(): + return Path(base_dir) / ORT_SUPPORT_DIR + return _resolve_home_dir() / "AppData" / "Local" / ORT_SUPPORT_DIR if os_name == "Darwin": home = _resolve_home_dir() @@ -44,102 +50,6 @@ def get_telemetry_base_dir() -> Path: # Use XDG_CACHE_HOME if set, otherwise fall back to $HOME/.cache cache_dir = os.getenv("XDG_CACHE_HOME") - if not cache_dir: - cache_dir = str(_resolve_home_dir() / ".cache") - - return Path(cache_dir).expanduser() / ORT_SUPPORT_DIR - - -def _format_exception_message(ex: BaseException, tb: Optional[TracebackType] = None) -> str: - """Format an exception and trim local paths for readability.""" - folder = "Olive" - file_line = 'File "' - formatted = traceback.format_exception(type(ex), ex, tb, limit=5) - lines = [] - for line in formatted: - line_trunc = line.strip() - if line_trunc.startswith(file_line) and folder in line_trunc: - idx = line_trunc.find(folder) - if idx != -1: - line_trunc = line_trunc[idx + len(folder) :] - elif line_trunc.startswith(file_line): - idx = line_trunc[len(file_line) :].find('"') - line_trunc = line_trunc[idx + len(file_line) :] - lines.append(line_trunc) - return "\n".join(lines) - - -class _ExclusiveFileLock: - """Cross-platform exclusive file lock context manager. - - Uses fcntl on Unix/Linux/macOS, msvcrt on Windows. - Prevents cache corruption when multiple processes access the same file. - - Design decisions: - - Lock is held for the entire duration of file access (prevents partial reads/writes) - - Lock is released automatically on close (even on exceptions) - - Platform-specific implementation (fcntl for POSIX, msvcrt for Windows) - - Assumptions: - - File locking is supported on the platform - - Lock is advisory on some systems (cooperative locking) - """ - - def __init__(self, file_path: Path, mode: str): - self.file_path = file_path - self.mode = mode - self.file = None - - def __enter__(self): - self.file = open(self.file_path, self.mode, encoding="utf-8") - - try: - # Platform-specific locking - if os.name == "posix": - import fcntl - - fcntl.flock(self.file.fileno(), fcntl.LOCK_EX) - elif os.name == "nt": - import msvcrt - - # Lock 1 byte at position 0 - msvcrt.locking(self.file.fileno(), msvcrt.LK_LOCK, 1) - except Exception: - self.file.close() - self.file = None - raise - - return self.file - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.file: - # Unlock happens automatically on close - self.file.close() - - -def _exclusive_file_lock(file_path: Path, mode: str): - """Create an exclusive file lock context manager. - - :param file_path: Path to the file to lock. - :param mode: File open mode ('r', 'a', 'w', etc.). - :return: Context manager that returns an open file handle. - """ - return _ExclusiveFileLock(file_path, mode) - - -def _encode_cache_line(plaintext: str) -> str: - """Encode a single cache line using base64. - - :param plaintext: The plaintext string to encode. - :return: Base64-encoded string (safe for a single text line). - """ - return base64.b64encode(plaintext.encode("utf-8")).decode("ascii") - - -def _decode_cache_line(encoded: str) -> str: - """Decode a single base64-encoded cache line. - - :param encoded: The base64-encoded string. - :return: The decoded plaintext string. - """ - return base64.b64decode(encoded.encode("ascii")).decode("utf-8") + if cache_dir and Path(cache_dir).is_absolute(): + return Path(cache_dir) / ORT_SUPPORT_DIR + return _resolve_home_dir() / ".cache" / ORT_SUPPORT_DIR diff --git a/olive/workflows/run/run.py b/olive/workflows/run/run.py index 6b6928daf9..ad57b38ec5 100644 --- a/olive/workflows/run/run.py +++ b/olive/workflows/run/run.py @@ -7,7 +7,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from copy import deepcopy from pathlib import Path -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union from olive.cache import isolated_cache_env from olive.common.utils import set_tempdir @@ -16,6 +16,15 @@ from olive.package_config import OlivePackageConfig from olive.systems.accelerator_creator import create_accelerator from olive.systems.common import SystemType +from olive.telemetry.recipe_telemetry import _build_recipe_result_metadata, _load_config_input_for_telemetry +from olive.telemetry.telemetry import Telemetry, is_ci_environment +from olive.telemetry.telemetry_extensions import ( + _format_exception_message, + _is_exception_logged, + _mark_exception_logged, + log_error, + log_recipe_result, +) from olive.workflows.run.builds import MultiBuildRunConfig, get_build_cache_dir, parse_run_config from olive.workflows.run.config import RunConfig @@ -150,25 +159,76 @@ def run( list_required_packages: bool = False, package_config: Optional[Union[str, Path, dict]] = None, tempdir: Optional[Union[str, Path]] = None, + recipe_telemetry_metadata: Optional[dict[str, Any]] = None, + emit_recipe_telemetry: bool = True, + emit_error_telemetry: bool = True, ): # set tempdir set_tempdir(tempdir) + package_config_input = package_config + try: + package_config_telemetry_input = ( + _load_config_input_for_telemetry(package_config_input) if package_config_input is not None else None + ) + except Exception: + package_config_telemetry_input = None + + package_config_provided = package_config is not None if package_config is None: package_config = OlivePackageConfig.get_default_config_path() - package_config = OlivePackageConfig.parse_file_or_obj(package_config) - parsed_config = parse_run_config(run_config) - if isinstance(parsed_config, MultiBuildRunConfig): - if list_required_packages: - _list_required_packages(package_config, parsed_config.values()) - return None - return _run_builds_in_parallel(package_config, parsed_config) - - if list_required_packages: - _list_required_packages(package_config, [parsed_config]) - return None - return _run_single(package_config, parsed_config) + parsed_run_config = None + success = False + exception = None + try: + package_config = OlivePackageConfig.parse_file_or_obj(package_config) + parsed_config = parse_run_config(run_config) + if isinstance(parsed_config, MultiBuildRunConfig): + if list_required_packages: + _list_required_packages(package_config, parsed_config.values()) + workflow_output = None + else: + workflow_output = _run_builds_in_parallel(package_config, parsed_config) + else: + parsed_run_config = parsed_config + if list_required_packages: + _list_required_packages(package_config, [parsed_run_config]) + workflow_output = None + else: + workflow_output = _run_single(package_config, parsed_run_config) + success = True + return workflow_output + except Exception as exc: + exception = exc + raise + finally: + if exception is not None and emit_error_telemetry and not _is_exception_logged(exception): + log_error( + exception_type=type(exception).__name__, + exception_message=_format_exception_message(exception, exception.__traceback__), + ) + _mark_exception_logged(exception) + if emit_recipe_telemetry: + try: + metadata = _build_recipe_result_metadata( + run_config, + None, + parsed_run_config, + recipe_telemetry_metadata, + list_required_packages=list_required_packages, + package_config_input=package_config_telemetry_input, + package_config_provided=package_config_provided, + ) + recipe_name = metadata.pop("recipe_name", None) + if recipe_name: + log_recipe_result(recipe_name, success=success, metadata=metadata) + except Exception: + logger.debug("Failed to emit recipe result telemetry.", exc_info=True) + if is_ci_environment(): + telemetry = Telemetry.get_existing_instance() + if telemetry is not None: + telemetry.shutdown() def _run_builds_in_parallel(package_config: OlivePackageConfig, parsed_config: MultiBuildRunConfig) -> OrderedDict: @@ -234,8 +294,7 @@ def _run_named_build(package_config: OlivePackageConfig, build_name: str, run_co def _run_single(package_config: OlivePackageConfig, run_config: RunConfig): if run_config.engine.host and run_config.engine.host.type == SystemType.Docker: docker_system = run_config.engine.host.create_system() - return docker_system.run_workflow(run_config) - + return docker_system.run_workflow(deepcopy(run_config)) set_default_logger_severity(run_config.engine.log_severity_level) return run_engine(package_config, run_config) diff --git a/requirements.txt b/requirements.txt index 661e166586..9b8ebce78e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,6 @@ onnx onnx-shape-inference>=0.3.1 onnx_ir>=0.1.2 onnxscript>=0.5.3 -opentelemetry-sdk>=1.39.1 optuna pandas pydantic>=2.0 diff --git a/test/cli/init/test_init_command.py b/test/cli/init/test_init_command.py index 257f8c6dad..7ae505aeec 100644 --- a/test/cli/init/test_init_command.py +++ b/test/cli/init/test_init_command.py @@ -2,7 +2,9 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # ------------------------------------------------------------------------- -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +import pytest class TestInitCommand: @@ -31,7 +33,46 @@ def test_run(self, mock_wizard_cls): args = parser.parse_args(["init", "-o", "./my-output"]) cmd = InitCommand(parser, args, []) - cmd.run() + telemetry = MagicMock(accepts_detailed_events=True) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.log_action") as log_action, + patch("olive.telemetry.telemetry_extensions.log_error") as log_error, + ): + cmd.run() mock_wizard_cls.assert_called_once_with(default_output_path="./my-output") mock_wizard_cls.return_value.start.assert_called_once() + log_action.assert_called_once() + assert log_action.call_args.kwargs["action_name"] == "Init" + assert log_action.call_args.kwargs["success"] is True + log_error.assert_not_called() + + @patch("olive.cli.init.wizard.InitWizard") + def test_run_failure_emits_one_action_and_error(self, mock_wizard_cls): + from argparse import ArgumentParser + + from olive.cli.init import InitCommand + + parser = ArgumentParser() + sub_parsers = parser.add_subparsers() + InitCommand.register_subcommand(sub_parsers) + args = parser.parse_args(["init"]) + cmd = InitCommand(parser, args, []) + mock_wizard_cls.return_value.start.side_effect = RuntimeError("boom") + + telemetry = MagicMock(accepts_detailed_events=True) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.log_action") as log_action, + patch("olive.telemetry.telemetry_extensions.log_error") as log_error, + pytest.raises(RuntimeError, match="boom"), + ): + cmd.run() + + log_action.assert_called_once() + assert log_action.call_args.kwargs["action_name"] == "Init" + assert log_action.call_args.kwargs["success"] is False + log_error.assert_called_once() diff --git a/test/cli/test_api.py b/test/cli/test_api.py index f6bfd151f8..f65c79f4f4 100644 --- a/test/cli/test_api.py +++ b/test/cli/test_api.py @@ -91,6 +91,18 @@ def test_optimize_function_basic(self, mock_cmd_cls): mock_cmd.run.assert_called_once() assert result is mock_output + def test_optimize_function_latches_disable_telemetry(self): + from olive import optimize + from olive.cli.optimize import OptimizeCommand + + with ( + patch("olive.cli.api.disable_telemetry") as disable, + patch.object(OptimizeCommand, "run", return_value=MagicMock()), + ): + optimize("test_model", disable_telemetry=True) + + disable.assert_called_once_with() + @patch("olive.cli.api.CaptureOnnxGraphCommand") def test_capture_cmd_basic(self, mock_cmd_cls): from olive import capture_onnx_graph diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index f12ab83592..b52c83c4ec 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -3,17 +3,119 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import json +import os import subprocess import sys +from argparse import Namespace from pathlib import Path from unittest.mock import MagicMock, patch import pytest from olive.cli.base import TEST_OUTPUT_MARKER_FILE +from olive.cli.launcher import get_cli_parser from olive.cli.launcher import main as cli_main +def test_launcher_handles_commands_without_disable_telemetry(): + parser = MagicMock() + service = MagicMock() + telemetry = MagicMock() + parser.parse_known_args.return_value = (Namespace(func=lambda *_: service), []) + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.Telemetry") as mock_telemetry, + ): + mock_telemetry.get_or_create_if_enabled.return_value = telemetry + cli_main([]) + + service.run.assert_called_once() + telemetry.shutdown.assert_called_once() + + +def test_launcher_without_subcommand_does_not_initialize_telemetry(): + parser = MagicMock() + parser.parse_known_args.return_value = (Namespace(), []) + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.Telemetry") as mock_telemetry, + pytest.raises(SystemExit), + ): + cli_main([]) + + parser.print_help.assert_called_once() + mock_telemetry.assert_not_called() + + +def test_launcher_shuts_down_telemetry_on_command_failure(): + parser = MagicMock() + service = MagicMock() + telemetry = MagicMock() + service.run.side_effect = RuntimeError("boom") + parser.parse_known_args.return_value = ( + Namespace(func=lambda *_: service, disable_telemetry=False), + [], + ) + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.Telemetry") as mock_telemetry, + ): + mock_telemetry.get_or_create_if_enabled.return_value = telemetry + with pytest.raises(RuntimeError, match="boom"): + cli_main([]) + + telemetry.shutdown.assert_called_once() + + +def test_launcher_latches_opt_out_before_constructing_telemetry(monkeypatch): + monkeypatch.delenv("OLIVE_DISABLE_TELEMETRY", raising=False) + parser = MagicMock() + service = MagicMock() + parser.parse_known_args.side_effect = [ + (Namespace(func=lambda *_: service, disable_telemetry=True), []), + (Namespace(func=lambda *_: service, disable_telemetry=False), []), + ] + call_order = [] + observed_during_run = [] + disabled = False + service.run.side_effect = lambda: observed_during_run.append(os.environ.get("OLIVE_DISABLE_TELEMETRY")) + + def latch_disable(): + nonlocal disabled + disabled = True + call_order.append("disable") + + def get_or_create(): + if disabled: + return None + call_order.append("construct") + return MagicMock() + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.disable_telemetry", side_effect=latch_disable), + patch("olive.cli.launcher.Telemetry.get_or_create_if_enabled", side_effect=get_or_create), + ): + cli_main([]) + cli_main([]) + + assert call_order == ["disable"] + assert observed_during_run == [None, None] + assert "OLIVE_DISABLE_TELEMETRY" not in os.environ + + +def test_init_command_parses_full_telemetry_opt_out(): + parser = get_cli_parser() + + args, unknown_args = parser.parse_known_args(["init", "--disable_telemetry"]) + + assert args.disable_telemetry is True + assert unknown_args == [] + + @pytest.mark.parametrize("console_script", [True, False]) @pytest.mark.parametrize( "command", @@ -108,7 +210,18 @@ def test_workflow_run_command(mock_run, tempdir, list_required_packages, tmp_pat # assert mock_run.assert_called_once_with( - {"key": "value"}, package_config=None, tempdir=tempdir, list_required_packages=list_required_packages + {"key": "value"}, + package_config=None, + tempdir=tempdir, + list_required_packages=list_required_packages, + recipe_telemetry_metadata={ + "recipe_command": "WorkflowRun", + "recipe_source": "config_file", + "recipe_format": "json", + "execution_mode": "list_required_packages" if list_required_packages else "run", + "package_config_provided": False, + }, + emit_error_telemetry=False, ) @@ -234,12 +347,30 @@ def test_workflow_run_command_with_overrides(mock_repo_exists, mock_run, tmp_pat list_required_packages=False, package_config=None, tempdir=None, + recipe_telemetry_metadata={ + "recipe_command": "WorkflowRun", + "recipe_source": "config_file", + "recipe_format": "json", + "execution_mode": "run", + "package_config_provided": False, + "config_overrides": { + "input_model": { + "type": "HfModel", + "model_path": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "load_kwargs": {"attn_implementation": "sdpa", "trust_remote_code": False}, + }, + "output_dir": str(Path("new_output_path").resolve()), + "log_severity_level": 2, + }, + }, + emit_error_telemetry=False, ) @patch("olive.workflows.run") def test_workflow_run_command_with_test_override(mock_run, tmp_path): mock_run.return_value = None + llama_env_path = str(tmp_path / "llama_env") config_path = tmp_path / "config.json" config_path.write_text( json.dumps( @@ -253,7 +384,14 @@ def test_workflow_run_command_with_test_override(mock_run, tmp_path): } ) ) - command_args = ["run", "--run-config", str(config_path), "--test"] + command_args = [ + "run", + "--run-config", + str(config_path), + "--test", + "--test_llama_path", + llama_env_path, + ] cli_main(command_args) @@ -271,6 +409,11 @@ def test_workflow_run_command_with_test_override(mock_run, tmp_path): "output_dir": output_dir, "passes": { "save_test_model_config": {"type": "SaveTestModelConfig"}, + "convert_hf_to_gguf": { + "type": "ConvertHfToGGUF", + "llama_cpp_env_path": llama_env_path, + "reference_model_path": test_model_path, + }, "discrepancy_check": { "type": "OnnxDiscrepancyCheck", "reference_model_path": test_model_path, @@ -278,15 +421,47 @@ def test_workflow_run_command_with_test_override(mock_run, tmp_path): "test_metrics": ["mae"], "max_mae": 0.1, "timing_iterations": 0, + "llama_cpp": True, + "llama_cpp_env_path": llama_env_path, }, }, }, list_required_packages=False, package_config=None, tempdir=None, + recipe_telemetry_metadata={ + "recipe_command": "WorkflowRun", + "recipe_source": "config_file", + "recipe_format": "json", + "execution_mode": "run", + "package_config_provided": False, + }, + emit_error_telemetry=False, ) +@patch("olive.cli.run.warn_unused_test_metrics") +@patch("olive.workflows.run") +def test_workflow_run_command_warns_for_test_options_without_test(mock_run, mock_warn, tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"key": "value"})) + llama_env_path = str(tmp_path / "llama_env") + + cli_main( + [ + "run", + "--run-config", + str(config_path), + "--test_metrics", + "mae", + "--test_llama_path", + llama_env_path, + ] + ) + + mock_warn.assert_called_once_with(False, ["mae"], llama_env_path) + + def test_workflow_run_command_with_test_rejects_non_test_output_dir(tmp_path): config_path = tmp_path / "config.json" output_dir = tmp_path / "output" diff --git a/olive/telemetry/constants.py b/test/cli/test_run_pass_action.py similarity index 52% rename from olive/telemetry/constants.py rename to test/cli/test_run_pass_action.py index ca9e150b1b..2fc3a0e1a1 100644 --- a/olive/telemetry/constants.py +++ b/test/cli/test_run_pass_action.py @@ -2,7 +2,11 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +import inspect -"""OneCollector connection string.""" +from olive.cli.run_pass import RunPassCommand -CONNECTION_STRING = "SW5zdHJ1bWVudGF0aW9uS2V5PTlkNWRkYWVjNjFlMjQ1NjdiNzg4YTIwYWVhMzI0NjMxLTcyMzdkN2M2LWVlNjEtNGNmZC1iYjdiLTU5MDNhOTcyYzJlNC03MDQ3" + +def test_run_pass_action_decorates_run_method_not_command_class(): + assert inspect.isclass(RunPassCommand) + assert hasattr(RunPassCommand.run, "__wrapped__") diff --git a/test/conftest.py b/test/conftest.py index db97c685af..fb70f9b4c5 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -2,13 +2,14 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +import os import shutil from unittest.mock import patch import pytest from packaging import version -from olive.telemetry.telemetry import Telemetry +import olive.telemetry.telemetry as telemetry_module from test.utils import create_onnx_model_file, delete_onnx_model_files @@ -44,4 +45,11 @@ def maybe_patch_inc(): @pytest.fixture(scope="session", autouse=True) def disable_telemetry(): - Telemetry().disable_telemetry() + # Apply the full opt-out before singleton construction so tests create no + # telemetry identity, store, uploader, heartbeat thread, or network traffic. + with patch.dict(os.environ, {"ORT_DISABLE_TELEMETRY": "1"}): + telemetry = telemetry_module.Telemetry() + try: + yield + finally: + telemetry.shutdown() diff --git a/test/systems/docker/test_docker_system.py b/test/systems/docker/test_docker_system.py index 5430b68587..2a5f805ca5 100644 --- a/test/systems/docker/test_docker_system.py +++ b/test/systems/docker/test_docker_system.py @@ -2,6 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +import json from unittest.mock import MagicMock, patch import pytest @@ -10,7 +11,7 @@ from olive.systems.system_config import DockerTargetUserConfig from test.utils import ONNX_MODEL_PATH -# pylint: disable=attribute-defined-outside-init,protected-access +# pylint: disable=attribute-defined-outside-init,duplicate-code,protected-access class TestDockerSystem: @@ -140,6 +141,64 @@ def test_run_workflow(self, mock_find_resources, mock_tempdir, mock_from_env, tm # Verify cleanup mock_container.remove.assert_called_once() + @patch("olive.systems.docker.docker_system.docker.from_env") + def test_prepare_environment_forwards_ci_to_workflow_container(self, mock_from_env, monkeypatch): + mock_docker_client = MagicMock() + mock_from_env.return_value = mock_docker_client + mock_docker_client.images.get.return_value = MagicMock() + monkeypatch.setenv("TF_BUILD", "True") + docker_config = self.get_default_docker_config() + docker_system = DockerSystem( + image_name=docker_config.image_name, + build_context_path=docker_config.build_context_path, + dockerfile=docker_config.dockerfile, + work_dir=docker_config.work_dir, + ) + + environment = docker_system._prepare_environment({}) + + assert environment["CI"] == "1" + + @patch("olive.systems.docker.docker_system.docker.from_env") + def test_prepare_environment_forwards_full_telemetry_opt_out(self, mock_from_env): + from olive.telemetry.telemetry import Telemetry + + mock_from_env.return_value.images.get.return_value = MagicMock() + docker_config = self.get_default_docker_config() + docker_system = DockerSystem( + image_name=docker_config.image_name, + build_context_path=docker_config.build_context_path, + dockerfile=docker_config.dockerfile, + work_dir=docker_config.work_dir, + ) + + with patch.object(Telemetry, "_process_disabled", True): + environment = docker_system._prepare_environment({}) + + assert environment["OLIVE_DISABLE_TELEMETRY"] == "1" + + def test_workflow_runner_disables_inner_recipe_telemetry(self, tmp_path, monkeypatch): + from olive.systems.docker import workflow_runner + + monkeypatch.delenv("HF_TOKEN", raising=False) + config = {"input_model": {"type": "ONNXModel", "model_path": "model.onnx"}} + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(config)) + + telemetry = MagicMock() + with ( + patch.object(workflow_runner, "olive_run") as mock_olive_run, + patch.object(workflow_runner.Telemetry, "get_existing_instance", return_value=telemetry), + ): + workflow_runner.runner_entry(config_path) + + mock_olive_run.assert_called_once_with( + config, + emit_error_telemetry=False, + emit_recipe_telemetry=False, + ) + telemetry.shutdown.assert_called_once_with(flush=True) + @patch("olive.systems.docker.docker_system.docker.from_env") @patch("olive.systems.docker.docker_system.tempfile.TemporaryDirectory") @patch("olive.systems.docker.docker_system.find_all_resources") diff --git a/test/test_telemetry.py b/test/test_telemetry.py new file mode 100644 index 0000000000..741d2feb71 --- /dev/null +++ b/test/test_telemetry.py @@ -0,0 +1,2107 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=duplicate-code,protected-access,redefined-outer-name +"""Tests for the SQLite-backed telemetry pipeline. + +Covers enabled, full opt-out, and CI recipe-only semantics; the +ALLOWED_KEYS whitelist filtering, the durable SQLite store, the single-drainer +process lock, the background uploader's success/poison/transient handling, and +the Common Schema serialization helpers. No test touches the network or the real +user profile: the HTTP transport is stubbed and the store directory is +redirected to a temp dir. +""" + +import hashlib +import json +import os +import sqlite3 +import stat +import subprocess +import sys +import tempfile +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import olive.telemetry.deviceid._store as deviceid_store_mod +import olive.telemetry.library.transport as transport_mod +import olive.telemetry.telemetry as tmod +import olive.telemetry.utils as telemetry_utils +from olive.telemetry.library.connection_string_parser import ConnectionStringParser +from olive.telemetry.library.options import OneCollectorTransportOptions +from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper as Serializer +from olive.telemetry.offline_store import OfflineEventStore +from olive.telemetry.process_lock import ProcessDrainLock +from olive.telemetry.telemetry_redaction import scrub_value_for_telemetry +from olive.telemetry.uploader import DrainOutcome, DrainResult, EventUploader + +ACTION_EVENT_NAME = tmod.ACTION_EVENT_NAME +ERROR_EVENT_NAME = tmod.ERROR_EVENT_NAME +HEARTBEAT_EVENT_NAME = tmod.HEARTBEAT_EVENT_NAME +RECIPE_EVENT_NAME = tmod.RECIPE_EVENT_NAME +Telemetry = tmod.Telemetry +is_ci_environment = tmod.is_ci_environment + +_FULL_DISABLE_VAR = "ORT_DISABLE_TELEMETRY" +_OPT_OUT_VAR = "OLIVE_DISABLE_TELEMETRY" +_CI_VARS = ( + "CI", + "TF_BUILD", + "GITHUB_ACTIONS", + "GITLAB_CI", + "CIRCLECI", + "TRAVIS", + "JENKINS_URL", + "CODEBUILD_BUILD_ID", + "BUILDKITE", + "TEAMCITY_VERSION", + "APPVEYOR", + "BITBUCKET_BUILD_NUMBER", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", +) + + +@pytest.fixture +def tenv(tmp_path, monkeypatch): + """Hermetic telemetry environment. + + Clears CI/opt-out signals so each test sets its own mode, stubs the HTTP + transport (recording every send in ``.sends``), and redirects the durable + store off the real profile. + """ + Telemetry._instance = None + Telemetry._process_disabled = False + Telemetry._heartbeat_enqueued = False + for var in (_FULL_DISABLE_VAR, _OPT_OUT_VAR, *_CI_VARS): + monkeypatch.delenv(var, raising=False) + + sends = [] + + def _record_send(self, payload, timeout_sec, item_count=1, on_send_admitted=None): + if on_send_admitted is not None: + on_send_admitted() + sends.append({"item_count": item_count, "size": len(payload), "payload": payload}) + return True, 204 + + monkeypatch.setattr(transport_mod.HttpJsonPostTransport, "send", _record_send) + monkeypatch.setattr(tmod, "get_telemetry_base_dir", lambda: tmp_path) + monkeypatch.setattr(telemetry_utils, "get_telemetry_base_dir", lambda: tmp_path) + monkeypatch.setattr(deviceid_store_mod, "get_telemetry_base_dir", lambda: tmp_path) + + yield SimpleNamespace(sends=sends, tmp_path=tmp_path) + + inst = Telemetry._instance + if inst is not None: + uploader = getattr(inst, "_uploader", None) + if uploader is not None: + uploader.stop_loop(5) + Telemetry._instance = None + Telemetry._process_disabled = False + Telemetry._heartbeat_enqueued = False + + +def _quiesce(t): + """Stop the background loop and drain its durable queue deterministically.""" + if t._uploader is not None: + t._uploader.stop_loop(5) + for _ in range(20): + if t._store is None or t._store.count() == 0: + break + t._uploader.drain_once() + + +def _sent_event_names(sends): + names = [] + for s in sends: + payload = bytes(s["payload"]) + names.extend( + token.decode() + for token in (b"OliveHeartbeat", b"OliveRecipe", b"OliveAction", b"OliveError") + if token in payload + ) + return names + + +# -------------------------------------------------------------------------- +# Full opt-out semantics +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", " yes ", "on", "y"]) +@pytest.mark.parametrize("variable", [_FULL_DISABLE_VAR, _OPT_OUT_VAR]) +def test_environment_opt_out_uses_canonical_allowlist(tenv, monkeypatch, variable, value): + monkeypatch.setenv(variable, value) + + assert tmod.is_telemetry_disabled_by_environment() is True + + +@pytest.mark.parametrize("value", ["", "0", "false", "no", "off", "unexpected"]) +@pytest.mark.parametrize("variable", [_FULL_DISABLE_VAR, _OPT_OUT_VAR]) +def test_environment_opt_out_rejects_non_allowlisted_values(tenv, monkeypatch, variable, value): + monkeypatch.setenv(variable, value) + + assert tmod.is_telemetry_disabled_by_environment() is False + + +def test_ci_is_recipe_only_with_no_heartbeat(tenv, monkeypatch): + monkeypatch.setenv("CI", "1") + t = Telemetry() + + # CI suppresses the device-id heartbeat but still persists recipe events. + assert t._store is not None + assert t.accepts_detailed_events is False + assert t._uploader._thread is None + + t.log(RECIPE_EVENT_NAME, {"recipe_name": "r", "success": True}) + t.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + + _quiesce(t) + names = _sent_event_names(tenv.sends) + assert "OliveHeartbeat" not in names + assert "OliveRecipe" in names + + +def test_ci_recipe_queue_does_not_drain_local_details(tenv, monkeypatch): + local_store = OfflineEventStore(str(tenv.tmp_path / tmod.DB_FILE_NAME)) + local_store.store(b'{"name":"OliveAction"}') + monkeypatch.setenv("CI", "1") + + telemetry = Telemetry() + telemetry.log(RECIPE_EVENT_NAME, {"recipe_name": "r", "success": True}) + _quiesce(telemetry) + + assert local_store.count() == 1 + assert "OliveAction" not in _sent_event_names(tenv.sends) + assert "OliveRecipe" in _sent_event_names(tenv.sends) + local_store.close() + + +def test_ci_shutdown_flushes_recipe(tenv, monkeypatch): + monkeypatch.setenv("CI", "1") + with patch.object(EventUploader, "start") as start: + telemetry = Telemetry() + telemetry.log(RECIPE_EVENT_NAME, {"recipe_name": "r", "success": True}) + + telemetry.shutdown() + + start.assert_not_called() + assert _sent_event_names(tenv.sends) == ["OliveRecipe"] + + +def test_ci_user_opt_out_fully_disables_telemetry(tenv, monkeypatch): + monkeypatch.setenv("CI", "1") + monkeypatch.setenv(_OPT_OUT_VAR, "1") + with ( + patch.object(tmod, "OfflineEventStore") as mock_store, + patch.object(tmod, "get_hashed_device_id_and_status") as mock_device_id, + ): + telemetry = Telemetry() + + assert Telemetry._instance is None + assert telemetry._enabled is False + assert telemetry._store is None + mock_store.assert_not_called() + mock_device_id.assert_not_called() + assert tenv.sends == [] + + +def test_user_opt_out_sends_and_persists_nothing(tenv, monkeypatch): + monkeypatch.setenv(_OPT_OUT_VAR, "1") + with ( + patch.object(tmod, "OfflineEventStore") as mock_store, + patch.object(tmod, "EventUploader") as mock_uploader, + patch.object(tmod, "get_hashed_device_id_and_status") as mock_device_id, + ): + t = Telemetry() + + assert t._enabled is False + assert t._store is None + assert t._uploader is None + assert t.accepts_detailed_events is False + mock_store.assert_not_called() + mock_uploader.assert_not_called() + mock_device_id.assert_not_called() + assert tenv.sends == [] + + +def test_user_opt_out_is_latched_across_reinitialization(tenv, monkeypatch): + monkeypatch.setenv(_OPT_OUT_VAR, "true") + first = Telemetry() + first.shutdown() + monkeypatch.delenv(_OPT_OUT_VAR) + + second = Telemetry() + + assert Telemetry._instance is None + assert second is not first + assert second._enabled is False + assert second._store is None + assert second._uploader is None + assert tenv.sends == [] + + +def test_full_disable_and_ci_send_nothing(tenv, monkeypatch): + monkeypatch.setenv(_FULL_DISABLE_VAR, "1") + monkeypatch.setenv("CI", "1") + t = Telemetry() + _quiesce(t) + + # Explicit full suppression + CI: record and send nothing at all. + assert Telemetry._instance is None + assert t._enabled is False + assert t._store is None + assert tenv.sends == [] + + +def test_full_disable_sends_nothing(tenv, monkeypatch): + monkeypatch.setenv(_FULL_DISABLE_VAR, "1") + t = Telemetry() + _quiesce(t) + + assert Telemetry._instance is None + assert t._enabled is False + assert t._store is None + assert t._uploader is None + assert tenv.sends == [] + + +def test_public_disable_before_initialization_only_latches_suppression(tenv): + from olive.telemetry.telemetry_extensions import log_action, log_error, log_recipe_result + + with ( + patch.object(tmod, "OfflineEventStore") as mock_store, + patch.object(tmod, "EventUploader") as mock_uploader, + patch.object(tmod, "get_hashed_device_id_and_status") as mock_device_id, + ): + tmod.disable_telemetry() + assert Telemetry._instance is None + log_action("cli", "work", 1.0, True) + log_error("RuntimeError", "boom") + log_recipe_result("recipe", True) + assert Telemetry._instance is None + telemetry = Telemetry() + tmod.disable_telemetry() + + assert Telemetry._instance is None + mock_store.assert_not_called() + mock_uploader.assert_not_called() + mock_device_id.assert_not_called() + assert telemetry._enabled is False + assert telemetry._store is None + assert tenv.sends == [] + + +def test_environment_opt_out_helpers_do_not_publish_singleton(tenv, monkeypatch): + from olive.telemetry.telemetry_extensions import log_action, log_error, log_recipe_result + + monkeypatch.setenv(_OPT_OUT_VAR, "1") + log_action("cli", "work", 1.0, True) + log_error("RuntimeError", "boom") + log_recipe_result("recipe", True) + + assert Telemetry._instance is None + + +def test_environment_opt_out_after_initialization_retains_queue(tenv, monkeypatch): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + telemetry.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1, "success": True}) + row_count = telemetry._store.count() + + monkeypatch.setenv(_OPT_OUT_VAR, "1") + assert Telemetry.get_or_create_if_enabled() is None + + assert telemetry._disabled is True + remaining_store = OfflineEventStore(str(tenv.tmp_path / tmod.DB_FILE_NAME)) + assert remaining_store.count() == row_count + remaining_store.close() + + +def test_enabled_records_heartbeat_and_events(tenv): + import uuid + + t = Telemetry() + session_guid = uuid.UUID(t._app_session_guid) + assert session_guid.version == 4 + assert session_guid.variant == uuid.RFC_4122 + + assert t._enabled is True + assert t._store is not None + assert t.accepts_detailed_events is True + + t.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + + _quiesce(t) + names = _sent_event_names(tenv.sends) + assert "OliveHeartbeat" in names + assert "OliveAction" in names + + +def test_heartbeat_is_stored_durably_before_uploader_starts(tenv): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + + batch = telemetry._store.get_batch(10) + assert len(batch) == 1 + assert json.loads(batch[0][1])["name"] == HEARTBEAT_EVENT_NAME + assert tenv.sends == [] + + +def test_heartbeat_releases_minimal_fallback_when_enrichment_fails(tenv): + with ( + patch.object(EventUploader, "start"), + patch.object( + tmod, + "get_hashed_device_id_and_status", + return_value=("c:device", SimpleNamespace(value="Existing")), + ), + patch.object(tmod.platform, "system", side_effect=RuntimeError("unavailable")), + ): + telemetry = Telemetry() + + batch = telemetry._store.get_batch(10) + assert len(batch) == 1 + data = json.loads(batch[0][1])["data"] + assert data["deviceId"] == "c:device" + assert data["deviceIdStatus"] == "Existing" + assert "os" not in data + + +def test_disable_telemetry_stops_detailed_events(tenv): + t = Telemetry() + _quiesce(t) + t.disable_telemetry() + + assert t._enabled is False + before = t._store.count() if t._store is not None else 0 + t.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + after = t._store.count() if t._store is not None else 0 + assert after == before + t.shutdown() + assert Telemetry()._enabled is False + assert _sent_event_names(tenv.sends).count("OliveHeartbeat") == 1 + + +def test_runtime_disable_does_not_emit_an_additional_heartbeat(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + tenv.sends.clear() + telemetry.disable_telemetry() + telemetry.disable_telemetry() + + assert tenv.sends == [] + + +def test_runtime_disable_retains_all_unsent_rows(tenv): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + other_store = OfflineEventStore(telemetry._store.db_path) + other_store.store(b'{"other":1}') + telemetry.log( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}, + ) + + telemetry.disable_telemetry() + + payloads = [payload for _, payload in other_store.get_batch(10)] + assert b'{"other":1}' in payloads + assert {json.loads(payload)["name"] for payload in payloads if payload != b'{"other":1}'} == { + HEARTBEAT_EVENT_NAME, + ACTION_EVENT_NAME, + } + other_store.close() + + +def test_runtime_disable_serializes_with_inflight_log(tenv): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + entered = threading.Event() + release = threading.Event() + original_store = telemetry._store.store + + def blocked_store(payload): + entered.set() + assert release.wait(5) + return original_store(payload) + + with patch.object(telemetry._store, "store", side_effect=blocked_store): + logging_thread = threading.Thread( + target=telemetry.log, + args=( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}, + ), + ) + logging_thread.start() + assert entered.wait(5) + disable_thread = threading.Thread(target=telemetry.disable_telemetry) + disable_thread.start() + release.set() + logging_thread.join(5) + disable_thread.join(5) + + assert not logging_thread.is_alive() + assert not disable_thread.is_alive() + assert OfflineEventStore(str(tenv.tmp_path / tmod.DB_FILE_NAME)).count() == 2 + + +def test_shutdown_closes_store_without_uploader(): + t = object.__new__(Telemetry) + t._disabled = False + t._uploader = None + t._store = MagicMock() + t._initialized = True + + store = t._store + t.shutdown() + + store.close.assert_called_once() + assert t._store is None + assert t._initialized is False + + +def test_shutdown_uses_one_overall_budget(): + t = object.__new__(Telemetry) + t._disabled = False + t._uploader = MagicMock() + t._uploader.stop_loop.return_value = True + t._store = MagicMock() + uploader = t._uploader + + monotonic = MagicMock(side_effect=[100.0, 100.5, 101.0]) + with patch("olive.telemetry.telemetry.time", SimpleNamespace(monotonic=monotonic)): + t.shutdown(flush=True) + + uploader.stop_loop.assert_called_once_with(join_timeout_seconds=1.5) + uploader.flush.assert_called_once_with(1.0) + assert t._uploader is None + assert t._store is None + assert t._initialized is False + + +def test_shutdown_does_not_flush_durable_queue_by_default(): + telemetry = object.__new__(Telemetry) + telemetry._disabled = False + telemetry._recipe_only_ci_telemetry = False + telemetry._uploader = MagicMock() + telemetry._uploader.stop_loop.return_value = True + telemetry._store = MagicMock() + uploader = telemetry._uploader + + telemetry.shutdown() + + uploader.flush.assert_not_called() + + +def test_shutdown_does_not_wait_or_flush_after_full_disable(): + telemetry = object.__new__(Telemetry) + telemetry._initialized = True + telemetry._disabled = True + telemetry._uploader = MagicMock() + telemetry._uploader.stop_loop.return_value = True + telemetry._store = MagicMock() + uploader = telemetry._uploader + + telemetry.shutdown(flush=True) + + uploader.stop_loop.assert_called_once_with(join_timeout_seconds=0) + uploader.flush.assert_not_called() + + +def test_reinitialization_does_not_enqueue_second_heartbeat(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + telemetry.shutdown() + + restarted = Telemetry() + restarted.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + _quiesce(restarted) + + names = _sent_event_names(tenv.sends) + assert names.count("OliveHeartbeat") == 1 + assert names.count("OliveAction") == 1 + + +def test_live_uploader_keeps_store_open(): + telemetry = object.__new__(Telemetry) + telemetry._disabled = False + telemetry._uploader = MagicMock() + telemetry._uploader.stop_loop.return_value = False + telemetry._store = MagicMock() + + telemetry.shutdown() + + telemetry._store.close.assert_not_called() + + +def test_closed_store_disables_telemetry(tenv): + closed_store = MagicMock(is_open=False) + with ( + patch.object(tmod, "OfflineEventStore", return_value=closed_store), + patch.object(tmod, "EventUploader") as mock_uploader, + ): + t = Telemetry() + + assert t._enabled is False + assert t._store is None + assert Telemetry._heartbeat_enqueued is False + mock_uploader.assert_not_called() + + +def test_closed_store_allows_initialization_retry(tenv): + closed_store = MagicMock(is_open=False) + open_store = MagicMock(is_open=True) + with ( + patch.object(tmod, "OfflineEventStore", side_effect=[closed_store, open_store]), + patch.object(tmod, "EventUploader") as mock_uploader, + ): + first = Telemetry() + assert first._initialized is False + + second = Telemetry() + + assert second is first + assert second._initialized is True + assert second._enabled is True + mock_uploader.assert_called_once_with(open_store, instrumentation_key=second._instrumentation_key) + + +# -------------------------------------------------------------------------- +# Whitelist filtering / payload building +# -------------------------------------------------------------------------- + + +def test_build_payload_drops_non_whitelisted_keys(tenv): + t = Telemetry() + _quiesce(t) + + payload = t._build_payload( + ACTION_EVENT_NAME, + { + "invoked_from": "cli", + "action_name": "WorkflowRun", + "duration_ms": 1.0, + "success": True, + "secret": "SHOULD_NOT_BE_SENT", + }, + ) + data = json.loads(payload)["data"] + assert "secret" not in data + assert data["actionName"] == "WorkflowRun" + # Defaults are stamped on every event. + assert data["appName"] == "Olive" + assert data["LibraryVersion"] + assert data["AppSessionGuid"] + assert "appVersion" not in data + assert "appSessionGuid" not in data + + +def test_build_payload_returns_none_for_unknown_event(tenv): + t = Telemetry() + _quiesce(t) + assert t._build_payload("TotallyUnknownEvent", {"k": "v"}) is None + + +def test_build_payload_heartbeat_uses_flat_os_fields(tenv): + t = Telemetry() + _quiesce(t) + + payload = t._build_payload( + HEARTBEAT_EVENT_NAME, + { + "device_id": "DEVICE", + "device_id_status": "ok", + "os": "Windows", + "os_version": "10.0.22631", + "os_release": "11", + "os_arch": "AMD64", + "leak": "DROP", + }, + ) + data = json.loads(payload)["data"] + assert data["deviceId"] == "DEVICE" + assert data["deviceIdStatus"] == "ok" + assert data["os"] == "Windows" + assert data["osVersion"] == "10.0.22631" + assert "leak" not in data + + +def test_build_payload_drops_non_scalar_whitelisted_values(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + { + "recipe_name": "optimize", + "success": True, + "config_overrides": { + "url": "https://user:secret@example.test/model", + Path("/home/alice/private"): [Path("/home/alice/model.onnx")], + }, + }, + ) + data = json.loads(payload)["data"] + + assert "configOverrides" not in data + + +def test_build_payload_drops_object_before_stringification(tenv): + class SensitiveObject: + def __str__(self): + return r"C:\Users\Alice\private.txt" + + telemetry = Telemetry() + _quiesce(telemetry) + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + {"recipe_name": "optimize", "success": True, "model_task": SensitiveObject()}, + ) + + serialized = payload.decode("utf-8") + assert "modelTask" not in json.loads(payload)["data"] + assert "Alice" not in serialized + + +def test_final_scrubber_redacts_nested_credential_aliases(): + scrubbed = scrub_value_for_telemetry( + { + "modelTask": { + "access_token": "snake-secret", + "accesstoken": "compact-secret", + "apiKey": "camel-secret", + "auth_header": "header-secret", + "credential_value": "credential-secret", + "docker_env": {"TOKEN": "env-secret"}, + "modelpath": "private/model.onnx", + "environmentVariables": {"HOME": "/home/private"}, + } + } + ) + + assert scrubbed["modelTask"] == { + "access_token": "", + "accesstoken": "", + "apiKey": "", + "auth_header": "", + "credential_value": "", + "docker_env": "", + "environmentVariables": "", + "modelpath": "", + } + + +def test_recipe_snapshot_json_remains_parseable_after_url_redaction(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + snapshot = json.dumps({"auth": "token=[secret]", "endpoint": "[path]", "z_value": "kept"}) + + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + {"recipe_name": "optimize", "success": True, "config_overrides": snapshot}, + ) + serialized_snapshot = json.loads(payload)["data"]["configOverrides"] + + assert json.loads(serialized_snapshot) == { + "auth": "", + "endpoint": "[path]", + "z_value": "kept", + } + + raw_snapshot = json.dumps( + { + "access_token": "snake-secret", + "accesstoken": "compact-secret", + "modelpath": "private/model.onnx", + } + ) + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + {"recipe_name": "optimize", "success": True, "config_overrides": raw_snapshot}, + ) + serialized_snapshot = json.loads(payload)["data"]["configOverrides"] + assert json.loads(serialized_snapshot) == { + "access_token": "", + "accesstoken": "", + "modelpath": "", + } + + +def test_recipe_snapshot_with_non_finite_value_preserves_event(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + snapshot = json.dumps({"invalid": float("nan")}) + + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + {"recipe_name": "optimize", "success": True, "config_overrides": snapshot}, + ) + data = json.loads(payload)["data"] + + assert data["recipeName"] == "optimize" + assert data["success"] is True + assert json.loads(data["configOverrides"]) == {"truncated": "[truncated]"} + + +def test_final_scrubber_scrubs_text_bytes_and_drops_binary(): + scrubbed = scrub_value_for_telemetry( + { + "text": rb"C:\Users\Alice Smith\model.onnx", + "binary": b"\xff\x00", + } + ) + + assert scrubbed["text"] == "[path]" + assert scrubbed["binary"] == "[binary]" + + +def test_non_finite_event_is_rejected_without_affecting_next_event(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + telemetry._uploader = None + + telemetry.log( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "bad", "duration_ms": float("nan"), "success": True}, + ) + telemetry.log( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "good", "duration_ms": 1.0, "success": True}, + ) + + assert telemetry._store.count() == 1 + + +def test_device_id_is_canonical_shared_hash(): + import olive.telemetry.deviceid.deviceid as deviceid + + raw_id = "123e4567-e89b-42d3-a456-426614174000" + with patch.dict( + deviceid._device_id_state, + {"device_id": raw_id, "status": deviceid.DeviceIdStatus.EXISTING}, + clear=True, + ): + hashed, status = deviceid.get_hashed_device_id_and_status() + + expected = hashlib.sha256(raw_id.encode("utf-8")).hexdigest() + assert hashed == f"c:{expected}" + assert len(hashed) == 66 + assert status == deviceid.DeviceIdStatus.EXISTING + + +def test_device_id_hash_matches_native_algorithm_vector(): + import olive.telemetry.deviceid.deviceid as deviceid + + raw_id = "01234567-89ab-4def-8123-456789abcdef" + with patch.dict( + deviceid._device_id_state, + {"device_id": raw_id, "status": deviceid.DeviceIdStatus.EXISTING}, + clear=True, + ): + hashed, _ = deviceid.get_hashed_device_id_and_status() + + assert hashed == f"c:{hashlib.sha256(raw_id.encode('utf-8')).hexdigest()}" + assert hashed == "c:6225bd190d6ccf87766a49c9986d174def3391fe175a61525e49a1d2334d6a43" + + +def test_global_metadata_is_merged_then_filtered(tenv): + t = Telemetry() + _quiesce(t) + + t.add_global_metadata({"app_version": "9.9.9", "app_instance_id": "mallory@example.test", "not_allowed": "DROP"}) + payload = t._build_payload( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}, + ) + data = json.loads(payload)["data"] + assert data["LibraryVersion"] == tmod.VERSION + assert data["AppSessionGuid"] == t._app_session_guid + assert "mallory" not in payload.decode("utf-8") + assert "not_allowed" not in data + + +def test_event_attributes_override_metadata(tenv): + t = Telemetry() + _quiesce(t) + payload = t._build_payload( + ERROR_EVENT_NAME, + {"exception_type": "ValueError", "exception_message": "safe"}, + {"exception_message": r"C:\Users\Mallory\secret.txt"}, + ) + assert json.loads(payload)["data"]["exceptionMessage"] == "safe" + + +def test_error_event_whitelist(tenv): + t = Telemetry() + _quiesce(t) + payload = t._build_payload( + ERROR_EVENT_NAME, + {"exception_type": "RuntimeError", "exception_message": "boom", "stack": "SENSITIVE"}, + ) + data = json.loads(payload)["data"] + assert data["exceptionType"] == "RuntimeError" + assert data["exceptionMessage"] == "boom" + assert "stack" not in data + + +@pytest.mark.parametrize("event_name", sorted(tmod.ALLOWED_KEYS)) +def test_all_whitelisted_fields_use_canonical_names(tenv, event_name): + t = Telemetry() + _quiesce(t) + source = dict.fromkeys(tmod.ALLOWED_KEYS[event_name], "value") + for snapshot_key in ("config_overrides", "package_config_overrides"): + if snapshot_key in source: + source[snapshot_key] = '{"key":"value"}' + + payload = t._build_payload(event_name, source) + data = json.loads(payload)["data"] + expected = {tmod.FIELD_NAMES.get(key, key) for key in source} + expected.update({"appName", "LibraryVersion", "AppSessionGuid"}) + + assert set(data) == expected + for source_name, canonical_name in tmod.FIELD_NAMES.items(): + if source_name != canonical_name and source_name in source: + assert source_name not in data + + +# -------------------------------------------------------------------------- +# CI detection +# -------------------------------------------------------------------------- + + +def test_is_ci_environment(monkeypatch): + for var in (_FULL_DISABLE_VAR, _OPT_OUT_VAR, *_CI_VARS): + monkeypatch.delenv(var, raising=False) + assert is_ci_environment() is False + monkeypatch.setenv("GITHUB_ACTIONS", "true") + assert is_ci_environment() is True + + +@pytest.mark.parametrize("value", ["", "0", "false", " no ", "OFF"]) +def test_false_ci_values_are_not_ci(monkeypatch, value): + with patch.dict(os.environ, {"CI": value}, clear=True): + assert is_ci_environment() is False + + +# -------------------------------------------------------------------------- +# Durable SQLite store +# -------------------------------------------------------------------------- + + +def _new_store(**kwargs): + db = os.path.join(tempfile.mkdtemp(), "olive_telemetry.db") + return OfflineEventStore(db, **kwargs) + + +def test_store_is_fifo(): + store = _new_store() + for i in range(5): + store.store(f'{{"e":{i}}}'.encode()) + assert store.count() == 5 + batch = store.get_batch(3) + assert [payload for _, payload in batch] == [b'{"e":0}', b'{"e":1}', b'{"e":2}'] + + +def test_store_closes_connection_when_initialization_fails(tmp_path): + connection = MagicMock() + connection.execute.side_effect = RuntimeError("pragma failed") + + with patch("olive.telemetry.offline_store.sqlite3.connect", return_value=connection): + store = OfflineEventStore(str(tmp_path / "failed.db")) + + assert store.is_open is False + connection.close.assert_called_once() + + +def test_store_delete(): + store = _new_store() + store.store(b'{"a":1}') + store.store(b'{"b":2}') + ids = [row_id for row_id, _ in store.get_batch(10)] + store.delete(ids[:1]) + assert store.count() == 1 + + +def test_store_trims_over_capacity(): + store = _new_store(max_records=8) + for i in range(40): + store.store(f'{{"i":{i}}}'.encode()) + assert store.count() <= 8 + + +def test_store_rejects_empty_payload(): + store = _new_store() + assert store.store(b"") is False + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX permissions") +def test_store_uses_owner_only_permissions(): + store = _new_store() + db_path = Path(store.db_path) + assert stat.S_IMODE(db_path.parent.stat().st_mode) == 0o700 + assert stat.S_IMODE(db_path.stat().st_mode) == 0o600 + + +def test_empty_permission_path_does_not_chmod_cwd(): + from olive.telemetry.offline_store import _chmod_best_effort + + with patch("olive.telemetry.offline_store.Path.chmod") as mock_chmod: + _chmod_best_effort("", 0o700) + + mock_chmod.assert_not_called() + + +def test_store_sets_schema_version(): + store = _new_store() + + assert store._conn.execute("PRAGMA user_version").fetchone()[0] == 3 + + +def test_store_migrates_availability_column(tmp_path): + db_path = tmp_path / "v1.db" + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE events (id INTEGER PRIMARY KEY AUTOINCREMENT, payload BLOB NOT NULL)") + conn.execute("PRAGMA user_version=1") + conn.commit() + conn.close() + + store = OfflineEventStore(str(db_path)) + + columns = {row[1] for row in store._conn.execute("PRAGMA table_info(events)")} + assert columns == {"id", "payload", "available_at", "acknowledged"} + assert store._conn.execute("PRAGMA user_version").fetchone()[0] == 3 + assert store.store(b'{"migrated":1}') is True + + +def test_reserved_event_is_hidden_until_release(): + store = _new_store() + row_id = store.reserve(b'{"minimal":1}', 60) + + assert row_id is not None + assert store.count() == 1 + assert store.get_batch(10) == [] + + released = store.release(row_id, b'{"enriched":1}') + assert released is True + assert store.get_batch(10) == [(row_id, b'{"enriched":1}')] + + +def test_failed_delete_is_reported_and_rolled_back(): + store = _new_store() + store.store(b'{"a":1}') + row_id = store.get_batch(1)[0][0] + store._conn.execute("CREATE TRIGGER fail_delete BEFORE DELETE ON events BEGIN SELECT RAISE(FAIL, 'blocked'); END") + store._conn.commit() + + deleted = store.delete([row_id]) + assert deleted is False + assert store.count() == 1 + + +def test_store_operations_bound_busy_timeout_to_deadline(): + store = _new_store() + store.store(b'{"a":1}') + row_id = store.get_batch(1)[0][0] + statements = [] + store._conn.set_trace_callback(statements.append) + + with patch("olive.telemetry.offline_store.time.monotonic", return_value=100.0): + deleted = store.delete([row_id], deadline=100.025) + + assert deleted + bounded = [statement for statement in statements if statement.startswith("PRAGMA busy_timeout=")] + assert bounded[0] in {"PRAGMA busy_timeout=24", "PRAGMA busy_timeout=25"} + assert bounded[-1] == "PRAGMA busy_timeout=3000" + + +# -------------------------------------------------------------------------- +# Single-drainer process lock +# -------------------------------------------------------------------------- + + +def _lock_path(): + return os.path.join(tempfile.mkdtemp(), "olive_telemetry.db.lock") + + +def test_lock_is_mutually_exclusive(): + path = _lock_path() + a = ProcessDrainLock(path) + b = ProcessDrainLock(path) + assert a.acquire() is True + assert b.acquire() is False # held by a + a.release() + assert b.acquire() is True # released + b.release() + + +def test_lock_reacquire_is_idempotent(): + a = ProcessDrainLock(_lock_path()) + assert a.acquire() is True + assert a.acquire() is True # already held + assert a.held is True + a.release() + assert a.held is False + + +# -------------------------------------------------------------------------- +# Uploader drain classification (no real network) +# -------------------------------------------------------------------------- + + +def _store_and_uploader(): + db = os.path.join(tempfile.mkdtemp(), "olive_telemetry.db") + store = OfflineEventStore(db) + uploader = EventUploader(store, instrumentation_key="abc-def") + return store, uploader + + +def test_uploader_deletes_on_success(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = lambda *a, **k: (True, 204) + result = uploader.drain_once() + assert (result.handled, result.left, result.outcome) == (1, 0, DrainOutcome.PROGRESS) + assert store.count() == 0 + + +def test_uploader_retention_latch_preserves_inflight_success(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + entered = threading.Event() + release = threading.Event() + result = [] + + def blocked_send(*_args, **_kwargs): + _kwargs["on_send_admitted"]() + entered.set() + assert release.wait(5) + return True, 204 + + uploader._transport.send = blocked_send + drain_thread = threading.Thread(target=lambda: result.append(uploader.drain_once())) + drain_thread.start() + assert entered.wait(5) + + uploader.retain_queued_rows() + release.set() + drain_thread.join(5) + + assert not drain_thread.is_alive() + assert result[0].outcome is DrainOutcome.TRANSPORT_RETRY + assert store.count() == 1 + + +def test_uploader_retention_latch_prevents_new_send(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock() + + uploader.retain_queued_rows() + + assert uploader.drain_once().outcome is DrainOutcome.TRANSPORT_RETRY + uploader._transport.send.assert_not_called() + assert store.count() == 1 + + +def test_uploader_retries_failed_delete_without_reposting(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + original_delete = store.delete + delete_attempts = 0 + + def flaky_delete(ids, deadline=None): + nonlocal delete_attempts + delete_attempts += 1 + return False if delete_attempts == 1 else original_delete(ids, deadline) + + store.delete = flaky_delete + + assert uploader.drain_once().outcome is DrainOutcome.DELETE_RETRY + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert store.count() == 0 + uploader._transport.send.assert_called_once() + + +def test_uploader_retries_failed_acknowledgement_without_reposting(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + original_acknowledge = store.acknowledge + acknowledge_attempts = 0 + + def flaky_acknowledge(ids, deadline=None): + nonlocal acknowledge_attempts + acknowledge_attempts += 1 + return False if acknowledge_attempts == 1 else original_acknowledge(ids, deadline) + + store.acknowledge = flaky_acknowledge + + assert uploader.drain_once().outcome is DrainOutcome.ACKNOWLEDGE_RETRY + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert store.count() == 0 + uploader._transport.send.assert_called_once() + + +def test_acknowledged_rows_are_deleted_without_reposting_after_restart(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + store.delete = MagicMock(return_value=False) + + assert uploader.drain_once().outcome is DrainOutcome.DELETE_RETRY + assert store._conn.execute("SELECT acknowledged FROM events").fetchone()[0] == 1 + db_path = store.db_path + store.close() + + reopened = OfflineEventStore(db_path) + restarted = EventUploader(reopened, instrumentation_key="abc-def") + restarted._transport.send = MagicMock() + assert restarted.drain_once().outcome is DrainOutcome.PROGRESS + assert reopened.count() == 0 + restarted._transport.send.assert_not_called() + + +def test_uploader_reports_storage_read_failure(): + store, uploader = _store_and_uploader() + store.get_batch_for_upload = MagicMock(return_value=None) + + assert uploader.drain_once().outcome is DrainOutcome.STORAGE_RETRY + + +def test_uploader_uses_only_remaining_deadline(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(False, None)) + + with patch("olive.telemetry.uploader.time.monotonic", return_value=100.75): + result = uploader.drain_once(deadline=101.0) + + assert (result.handled, result.left, result.outcome) == (0, 1, DrainOutcome.TRANSPORT_RETRY) + assert uploader._transport.send.call_args.args[1] == pytest.approx(0.25) + + +def test_request_drain_only_wakes_lock_holder(): + _, uploader = _store_and_uploader() + uploader._wake = MagicMock() + uploader._drain_lock = MagicMock(held=False) + + uploader.request_drain() + uploader._wake.set.assert_not_called() + + uploader._drain_lock.held = True + uploader.request_drain() + uploader._wake.set.assert_called_once() + + +def test_uploader_drops_poison_4xx(): + store, uploader = _store_and_uploader() + store.store(b'{"bad":1}') + uploader._transport.send = lambda *a, **k: (False, 400) + uploader.drain_once() + assert store.count() == 0 # dropped, not retried forever + + +def test_uploader_handles_oversized_single_row_without_transport(): + store, uploader = _store_and_uploader() + store.store(b'{"oversized":true}') + uploader._transport.send = MagicMock() + + with patch.object(OneCollectorTransportOptions, "DEFAULT_MAX_PAYLOAD_SIZE_BYTES", 1): + result = uploader.drain_once() + + assert (result.handled, result.left, result.outcome) == (1, 0, DrainOutcome.PROGRESS) + assert not hasattr(result, "delivered") + assert store.count() == 0 + uploader._transport.send.assert_not_called() + + +def test_uploader_retains_transient_5xx(): + store, uploader = _store_and_uploader() + store.store(b'{"later":1}') + uploader._transport.send = lambda *a, **k: (False, 503) + result = uploader.drain_once() + assert (result.handled, result.left, result.outcome) == (0, 1, DrainOutcome.TRANSPORT_RETRY) + assert store.count() == 1 # kept for retry + + +@pytest.mark.parametrize("status", [507, 520, 599]) +def test_uploader_retains_all_server_errors(status): + store, uploader = _store_and_uploader() + store.store(b'{"later":1}') + uploader._transport.send = lambda *a, **k: (False, status) + + assert uploader.drain_once().outcome is DrainOutcome.TRANSPORT_RETRY + assert store.count() == 1 + + +@pytest.mark.parametrize("status", [400, 413, 422]) +def test_uploader_isolates_rejected_event(status): + store, uploader = _store_and_uploader() + store.store(b'{"bad":1}') + store.store(b'{"valid":1}') + + def send(payload, timeout, item_count=1, on_send_admitted=None): + if on_send_admitted is not None: + on_send_admitted() + if item_count > 1 or b'"bad"' in payload: + return (False, status) + return (True, 204) + + uploader._transport.send = MagicMock(side_effect=send) + + assert uploader.drain_once().outcome is DrainOutcome.SPLIT + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert store.count() == 0 + assert uploader._transport.send.call_count == 3 + + +@pytest.mark.parametrize("status", [400, 413, 422]) +def test_flush_completes_poison_isolation_within_one_process(status): + store, uploader = _store_and_uploader() + store.store(b'{"bad":1}') + store.store(b'{"valid":1}') + + def send(payload, timeout, item_count=1, on_send_admitted=None): + if on_send_admitted is not None: + on_send_admitted() + if item_count > 1 or b'"bad"' in payload: + return (False, status) + return (True, 204) + + uploader._transport.send = MagicMock(side_effect=send) + uploader.flush(1) + + assert store.count() == 0 + assert uploader._transport.send.call_count == 3 + + +def test_flush_retries_acknowledged_delete_without_reposting(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + original_delete = store.delete + delete_attempts = 0 + + def flaky_delete(ids, deadline=None): + nonlocal delete_attempts + delete_attempts += 1 + return False if delete_attempts == 1 else original_delete(ids, deadline) + + store.delete = flaky_delete + uploader.flush(1) + + assert store.count() == 0 + uploader._transport.send.assert_called_once() + + +def test_flush_does_not_touch_lock_while_thread_is_alive(): + _, uploader = _store_and_uploader() + uploader._thread = MagicMock() + uploader._thread.is_alive.return_value = True + uploader._drain_lock.acquire = MagicMock() + uploader._drain_lock.release = MagicMock() + + uploader.flush(0.01) + + uploader._drain_lock.acquire.assert_not_called() + uploader._drain_lock.release.assert_not_called() + + +def test_uploader_backs_off_after_lock_contention(): + _, uploader = _store_and_uploader() + uploader._drain_lock = MagicMock() + uploader._drain_lock.acquire.return_value = False + uploader._stop = MagicMock() + uploader._stop.is_set.side_effect = [False, True] + uploader._wake = MagicMock() + + uploader._run() + + uploader._wake.wait.assert_called_once_with(uploader._idle_backoff) + + +def test_uploader_backs_off_after_drain_exception(): + _, uploader = _store_and_uploader() + uploader._drain_lock = MagicMock() + uploader._drain_lock.acquire.return_value = True + uploader._stop = MagicMock() + uploader._stop.is_set.side_effect = [False, False, True] + uploader._wake = MagicMock() + uploader.drain_once = MagicMock(side_effect=RuntimeError("transient failure")) + + uploader._run() + + uploader._wake.wait.assert_called_once_with(uploader._idle_backoff) + uploader._drain_lock.release.assert_called_once() + + +def test_uploader_backs_off_after_storage_failure(): + _, uploader = _store_and_uploader() + uploader._drain_lock = MagicMock() + uploader._drain_lock.acquire.return_value = True + uploader._stop = MagicMock() + uploader._stop.is_set.side_effect = [False, False, True] + uploader._wake = MagicMock() + uploader.drain_once = MagicMock(return_value=DrainResult(0, 0, DrainOutcome.STORAGE_RETRY)) + + uploader._run() + + uploader._wake.wait.assert_called_once_with(uploader._idle_backoff) + uploader._drain_lock.release.assert_called_once() + + +# -------------------------------------------------------------------------- +# Serialization + connection string parsing +# -------------------------------------------------------------------------- + + +def test_serialize_basic_types(): + assert Serializer.serialize_value(None) is None + assert Serializer.serialize_value(True) is True + assert Serializer.serialize_value(42) == 42 + assert Serializer.serialize_value("hello") == "hello" + assert Serializer.serialize_value([1, "two", 3.0]) == [1, "two", 3.0] + assert Serializer.serialize_value({"k": "v"}) == {"k": "v"} + assert Serializer.serialize_value({0: "zero", "": "skip"}) == {"0": "zero"} + assert Serializer.serialize_value({False: "false"}) == {"False": "false"} + + +def test_redacted_mapping_key_collisions_are_dropped_deterministically(): + first = {"/first": "a", "/second": "b", "safe": "kept"} + second = dict(reversed(list(first.items()))) + + assert scrub_value_for_telemetry(first) == {"safe": "kept"} + assert scrub_value_for_telemetry(second) == {"safe": "kept"} + assert Serializer.serialize_value(first) == {"safe": "kept"} + assert Serializer.serialize_value(second) == {"safe": "kept"} + + +def test_snapshot_key_collisions_are_dropped_deterministically(): + from olive.telemetry.telemetry_redaction import scrub_config_snapshot_for_telemetry + + first = {"/first": "a", "/second": "b", "safe": "kept"} + second = dict(reversed(list(first.items()))) + + assert scrub_config_snapshot_for_telemetry(first) == {"safe": "kept"} + assert scrub_config_snapshot_for_telemetry(second) == {"safe": "kept"} + + +def test_create_event_envelope(): + envelope = Serializer.create_event_envelope( + event_name="TestEvent", + timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc), + ikey="o:test-key", + data={"key": "value"}, + ) + assert envelope["name"] == "TestEvent" + assert envelope["iKey"] == "o:test-key" + assert envelope["data"] == {"key": "value"} + + +def test_connection_string_parser(): + assert ConnectionStringParser("InstrumentationKey=abc-def-ghi").instrumentation_key == "abc-def-ghi" + with pytest.raises(ValueError, match="Connection string cannot be empty"): + ConnectionStringParser("") + with pytest.raises(ValueError, match="InstrumentationKey"): + ConnectionStringParser("SomeOtherKey=value") + + +def test_http_retry_shares_one_timeout_budget(): + import urllib.error + import urllib.request + + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + with ( + patch( + "olive.telemetry.library.transport.urllib.request.urlopen", + side_effect=urllib.error.URLError("offline"), + ) as urlopen, + patch("olive.telemetry.library.transport.time.monotonic", side_effect=[10.0, 10.0, 10.75]), + ): + assert transport._do_request(request, 1.0) == (False, None) + + assert urlopen.call_count == 2 + assert urlopen.call_args_list[0].kwargs["timeout"] == pytest.approx(1.0) + assert urlopen.call_args_list[1].kwargs["timeout"] == pytest.approx(0.25) + + +def test_http_error_body_is_not_read_during_bounded_send(): + import urllib.error + import urllib.request + + response = MagicMock() + http_error = urllib.error.HTTPError("https://example.invalid", 503, "unavailable", {}, response) + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + + with patch("olive.telemetry.library.transport.urllib.request.urlopen", side_effect=http_error): + assert transport._do_request(request, 1.0) == (False, 503) + + response.read.assert_not_called() + response.close.assert_called_once() + + +def test_http_deadline_bounds_blocked_resolution(): + import urllib.request + + entered = threading.Event() + release = threading.Event() + response = MagicMock(status=204) + response.__enter__.return_value = response + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + + def blocked_urlopen(*_args, **_kwargs): + entered.set() + release.wait(5) + return response + + start = time.perf_counter() + with patch("olive.telemetry.library.transport.urllib.request.urlopen", side_effect=blocked_urlopen): + assert transport._do_request(request, 0.05) == (False, None) + assert entered.wait(1) + elapsed = time.perf_counter() - start + release.set() + + assert elapsed < 0.5 + + +def test_http_timeout_does_not_start_replacement_worker(): + import urllib.request + + entered = threading.Event() + release = threading.Event() + response = MagicMock(status=204) + response.__enter__.return_value = response + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + + def blocked_urlopen(*_args, **_kwargs): + entered.set() + release.wait(5) + return response + + with patch( + "olive.telemetry.library.transport.urllib.request.urlopen", + side_effect=blocked_urlopen, + ) as urlopen: + assert transport._do_request(request, 0.05) == (False, None) + assert entered.wait(1) + assert transport._do_request(request, 0.05) == (False, None) + assert urlopen.call_count == 1 + release.set() + transport._inflight_worker.join(1) + assert transport._do_request(request, 0.05) == (True, 204) + assert urlopen.call_count == 1 + + +@pytest.mark.parametrize("status", [500, 507, 520, 599]) +def test_all_server_errors_are_retryable(status): + assert transport_mod.HttpJsonPostTransport.is_retryable(status) + + +def test_serialization_orders_maps_and_sets_deterministically(): + helper = Serializer + timestamp = datetime(2025, 1, 1, tzinfo=timezone.utc) + first = helper.create_event_envelope("event", timestamp, "o:key", {"z": {"b", "a"}, "a": 1}) + second = helper.create_event_envelope("event", timestamp, "o:key", {"a": 1, "z": {"a", "b"}}) + + assert helper.serialize_to_json_bytes(first) == helper.serialize_to_json_bytes(second) + + +# -------------------------------------------------------------------------- +# Exception-message path redaction (privacy) +# -------------------------------------------------------------------------- + + +def test_redact_paths_and_general_length_contract(): + from olive.telemetry.telemetry_redaction import MAX_TELEMETRY_STRING_LENGTH, scrub_string_for_telemetry + + assert scrub_string_for_telemetry(r"C:\Users\alice\model.onnx") == "[path]" + assert scrub_string_for_telemetry(r"\secret.onnx") == "[path]" + assert scrub_string_for_telemetry(r"failed \secret.onnx") == "failed [path]" + assert scrub_string_for_telemetry("/var/data/run/output.log") == "[path]" + assert scrub_string_for_telemetry("/secret.onnx") == "[path]" + # Last segment is a directory/username (no extension) -> fully redacted. + assert scrub_string_for_telemetry("/home/bob") == "[path]" + # UNC paths are redacted too. + assert scrub_string_for_telemetry(r"\\server\share\secret") == "[path]" + assert scrub_string_for_telemetry(r"failed C:\Users\Alice Smith\models\phi.onnx") == "failed [path]" + assert scrub_string_for_telemetry("failed /home/Alice Smith/models/phi.onnx") == "failed [path]" + assert scrub_string_for_telemetry("a/b/c") == "[path]" + assert scrub_string_for_telemetry(r"Load Users\bob\model.onnx failed") == "Load [path]" + assert scrub_string_for_telemetry("Users/Alice Smith/models/phi.onnx") == "[path]" + assert scrub_string_for_telemetry("Users/Alice Smith/model.onnx") == "[path]" + assert scrub_string_for_telemetry("models/foo.onnx") == "models/foo.onnx" + assert scrub_string_for_telemetry("ratio 3/4 and and/or") == "ratio 3/4 and and/or" + assert scrub_string_for_telemetry("before /home/alice/model.onnx\nafter") == "before [path]" + assert scrub_string_for_telemetry("https://example.test/model?token=secret") == "[path]" + assert scrub_string_for_telemetry("download(s3://private-bucket/model)") == "download([path]" + assert scrub_string_for_telemetry("GET https://host/x?path=/home/alice/model.onnx failed") == "GET [path]" + assert scrub_string_for_telemetry("open ./private/model.onnx") == "open [path]" + assert scrub_string_for_telemetry(r"open ..\private\model.onnx") == "open [path]" + assert scrub_string_for_telemetry("request token=supersecret") == "request token=[secret]" + assert scrub_string_for_telemetry("token=first&api_key=second") == "token=[secret]" + assert scrub_string_for_telemetry("fetch example.test?access_token=top-secret") == ( + "fetch example.test?access_token=[secret]" + ) + assert scrub_string_for_telemetry("redirect#access_token=top-secret") == ("redirect#access_token=[secret]") + assert scrub_string_for_telemetry("auth.token=top-secret") == "auth.token=[secret]" + assert scrub_string_for_telemetry("auth=top-secret") == "auth=[secret]" + assert scrub_string_for_telemetry("refreshToken=top-secret") == "refreshToken=[secret]" + assert scrub_string_for_telemetry("AWS_SECRET_ACCESS_KEY=top-secret") == ("AWS_SECRET_ACCESS_KEY=[secret]") + assert scrub_string_for_telemetry("PWD=odbc-value") == "PWD=[secret]" + assert scrub_string_for_telemetry("Authorization: ******") == "Authorization: [secret]" + assert scrub_string_for_telemetry("failure --api-key top-secret") == "failure --api-key [secret]" + assert scrub_string_for_telemetry("Command ['tool', '--api-key', 'top-secret']") == ( + "Command ['tool', '--api-key', '[secret]" + ) + assert scrub_string_for_telemetry('--password "-abc"') == '--password "[secret]' + assert scrub_string_for_telemetry("Command ['tool', '--api-key', '-abc']") == ( + "Command ['tool', '--api-key', '[secret]" + ) + assert scrub_string_for_telemetry("--password -abc") == "--password -abc" + assert scrub_string_for_telemetry("model dir /_apikey:top-secret") == "model dir [path]" + assert scrub_string_for_telemetry("arg /1token=top-secret") == "arg [path]" + assert scrub_string_for_telemetry("connect /2fa_token=top-secret") == "connect [path]" + assert scrub_string_for_telemetry("failure /password:top-secret") == "failure /password:[secret]" + assert scrub_string_for_telemetry("connect user:password@localhost/model") == "connect [secret]" + assert scrub_string_for_telemetry("n/a read/write domain\\user") == "n/a read/write domain\\user" + assert scrub_string_for_telemetry("meta-llama/Llama-3.1-8B-Instruct") == "meta-llama/Llama-3.1-8B-Instruct" + assert scrub_string_for_telemetry("tokenizer=enabled oauth=enabled") == "tokenizer=enabled oauth=enabled" + assert ( + len(scrub_string_for_telemetry("x" * (MAX_TELEMETRY_STRING_LENGTH + 100)).encode("utf-8")) + == MAX_TELEMETRY_STRING_LENGTH + ) + assert scrub_string_for_telemetry("x" * (MAX_TELEMETRY_STRING_LENGTH - 1) + "€") == ( + "x" * (MAX_TELEMETRY_STRING_LENGTH - 1) + ) + redacted_at_limit = scrub_string_for_telemetry("x" * (MAX_TELEMETRY_STRING_LENGTH - 5) + " /a/b/c") + assert len(redacted_at_limit.encode("utf-8")) == MAX_TELEMETRY_STRING_LENGTH + assert redacted_at_limit.endswith("[path]") + + +def test_secret_scanner_advances_past_rejected_key_tokens(): + from olive.telemetry import telemetry_redaction + + value = "-" + "field-" * 6_800 + "value" + with patch( + "olive.telemetry.telemetry_redaction.is_sensitive_config_key_for_telemetry", + wraps=telemetry_redaction.is_sensitive_config_key_for_telemetry, + ) as is_sensitive_key: + scrubbed = telemetry_redaction.scrub_string_for_telemetry(value) + + assert scrubbed == value + assert is_sensitive_key.call_count == 1 + + +def test_error_messages_are_capped_at_40960_utf8_bytes(): + from olive.telemetry.telemetry_extensions import log_error + from olive.telemetry.telemetry_redaction import MAX_ERROR_MESSAGE_LENGTH + + telemetry = MagicMock() + with patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry): + log_error("RuntimeError", "x" * (MAX_ERROR_MESSAGE_LENGTH + 100)) + truncated = telemetry.log.call_args.args[1]["exception_message"] + assert len(truncated.encode("utf-8")) == MAX_ERROR_MESSAGE_LENGTH + + log_error("RuntimeError", "x" * (MAX_ERROR_MESSAGE_LENGTH - 1) + "€") + multibyte = telemetry.log.call_args.args[1]["exception_message"] + assert multibyte == "x" * (MAX_ERROR_MESSAGE_LENGTH - 1) + + log_error("RuntimeError", "x" * (MAX_ERROR_MESSAGE_LENGTH - 5) + " /a/b/c") + redacted_at_limit = telemetry.log.call_args.args[1]["exception_message"] + assert len(redacted_at_limit.encode("utf-8")) == MAX_ERROR_MESSAGE_LENGTH + assert redacted_at_limit.endswith("[path]") + + +def test_error_payload_preserves_error_specific_size_limit(tenv): + from olive.telemetry.telemetry_redaction import MAX_ERROR_MESSAGE_LENGTH + + telemetry = Telemetry() + _quiesce(telemetry) + payload = telemetry._build_payload( + ERROR_EVENT_NAME, + { + "exception_type": "RuntimeError", + "exception_message": "x" * (MAX_ERROR_MESSAGE_LENGTH + 100), + }, + ) + + assert len(json.loads(payload)["data"]["exceptionMessage"].encode("utf-8")) == MAX_ERROR_MESSAGE_LENGTH + + +def test_format_exception_message_redacts_paths_in_message(): + from olive.telemetry.telemetry_extensions import _format_exception_message + + exc = RuntimeError(r"failed to read C:\Users\alice\secret\weights.bin") + message = _format_exception_message(exc, exc.__traceback__) + assert "alice" not in message + assert "[path]" in message + + +def test_format_exception_message_handles_unprintable_exception(): + from olive.telemetry.telemetry_extensions import _format_exception_message + + class UnprintableError(Exception): + def __str__(self): + raise RuntimeError("cannot render") + + assert _format_exception_message(UnprintableError()).endswith("UnprintableError: ") + + +def test_public_helpers_never_propagate_failures(): + from olive.telemetry.telemetry_extensions import log_action, log_error, log_recipe_result + + with patch("olive.telemetry.telemetry_extensions._get_logger", side_effect=RuntimeError("telemetry failed")): + log_action("test", "work", 1.0, True, metadata=["not", "a", "dict"]) + log_error("RuntimeError", "boom", metadata=["not", "a", "dict"]) + log_recipe_result("recipe", True, metadata=["not", "a", "dict"]) + + +def test_log_recipe_result_never_propagates_when_telemetry_log_raises(): + from olive.telemetry.telemetry_extensions import log_recipe_result + + telemetry = MagicMock() + telemetry.log.side_effect = RuntimeError("log failed") + with patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry): + log_recipe_result("recipe", True) + + telemetry.log.assert_called_once() + + +def _raise_error_called_with_source_secret(_secret): + raise RuntimeError("boom") + + +def test_format_exception_message_omits_source_code(): + from olive.telemetry.telemetry_extensions import _format_exception_message + + try: + _raise_error_called_with_source_secret("source-secret") + except RuntimeError as ex: + message = _format_exception_message(ex, ex.__traceback__) + + assert "source-secret" not in message + assert __file__ not in message + assert 'File "[path]"' in message + assert "in _raise_error_called_with_source_secret" in message + assert message.endswith("RuntimeError: boom") + + +def test_device_id_store_uses_owner_only_creation_mode(tmp_path): + with ( + patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path), + patch.object(Path, "mkdir") as mock_mkdir, + ): + deviceid_store_mod.Store().store_id("test-device-id") + + mock_mkdir.assert_called_once_with(mode=0o700, parents=True, exist_ok=True) + + +def test_relative_cache_environment_uses_absolute_home(tmp_path, monkeypatch): + telemetry_utils.get_telemetry_base_dir.cache_clear() + monkeypatch.setenv("XDG_CACHE_HOME", "relative-cache") + monkeypatch.setenv("HOME", str(tmp_path)) + with patch.object(telemetry_utils.platform, "system", return_value="Linux"): + base_dir = telemetry_utils.get_telemetry_base_dir() + telemetry_utils.get_telemetry_base_dir.cache_clear() + + assert base_dir == tmp_path / ".cache" / telemetry_utils.ORT_SUPPORT_DIR + assert base_dir.is_absolute() + + +def test_relative_windows_appdata_uses_absolute_home(tmp_path, monkeypatch): + telemetry_utils.get_telemetry_base_dir.cache_clear() + monkeypatch.setenv("LOCALAPPDATA", "relative-local") + monkeypatch.setenv("APPDATA", "relative-roaming") + monkeypatch.setenv("HOME", "relative-home") + with ( + patch.object(telemetry_utils.platform, "system", return_value="Windows"), + patch.object(telemetry_utils.Path, "home", return_value=tmp_path), + ): + base_dir = telemetry_utils.get_telemetry_base_dir() + telemetry_utils.get_telemetry_base_dir.cache_clear() + + assert base_dir == tmp_path / "AppData" / "Local" / telemetry_utils.ORT_SUPPORT_DIR + assert base_dir.is_absolute() + + +def test_windows_telemetry_base_dir_uses_canonical_developer_tools_path(tmp_path, monkeypatch): + telemetry_utils.get_telemetry_base_dir.cache_clear() + monkeypatch.setenv("LOCALAPPDATA", str(tmp_path)) + with patch.object(telemetry_utils.platform, "system", return_value="Windows"): + base_dir = telemetry_utils.get_telemetry_base_dir() + telemetry_utils.get_telemetry_base_dir.cache_clear() + + assert base_dir == tmp_path / telemetry_utils.ORT_SUPPORT_DIR + + +def test_home_resolution_fails_without_absolute_per_user_directory(monkeypatch): + fake_pwd = SimpleNamespace(getpwuid=MagicMock(side_effect=KeyError("missing"))) + monkeypatch.setenv("HOME", "relative-home") + with ( + patch.object(telemetry_utils.platform, "system", return_value="Linux"), + patch.object(telemetry_utils.Path, "home", return_value=Path("relative-home")), + patch.dict(sys.modules, {"pwd": fake_pwd}), + pytest.raises(RuntimeError, match="No absolute per-user"), + ): + telemetry_utils._resolve_home_dir() + + +def test_device_id_is_ephemeral_when_per_user_storage_is_unavailable(): + import olive.telemetry.deviceid.deviceid as deviceid + + deviceid._device_id_state.update({"device_id": None, "status": deviceid.DeviceIdStatus.NEW}) + with ( + patch.object(deviceid.platform, "system", return_value="Linux"), + patch.object(deviceid, "Store", side_effect=RuntimeError("no home")), + ): + generated = deviceid.get_device_id() + + assert deviceid._is_valid_device_id(generated) + assert deviceid._device_id_state["status"] == deviceid.DeviceIdStatus.FAILED + + +@pytest.mark.parametrize("system", ["FreeBSD", "OpenBSD", "AIX", "SunOS"]) +def test_device_id_uses_file_store_on_other_posix_platforms(system): + import olive.telemetry.deviceid.deviceid as deviceid + + stored_id = "123e4567-e89b-42d3-a456-426614174000" + store = MagicMock() + store.retrieve_id = stored_id + deviceid._device_id_state.update({"device_id": None, "status": deviceid.DeviceIdStatus.NEW}) + + with ( + patch.object(deviceid.platform, "system", return_value=system), + patch.object(deviceid, "os", SimpleNamespace(name="posix")), + patch.object(deviceid, "Store", return_value=store) as store_type, + ): + assert deviceid.get_device_id() == stored_id + + store_type.assert_called_once_with() + assert deviceid._device_id_state["status"] == deviceid.DeviceIdStatus.EXISTING + + +def test_missing_device_id_raises_file_not_found(tmp_path): + with ( + patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path), + pytest.raises(FileNotFoundError), + ): + _ = deviceid_store_mod.Store().retrieve_id + + +def test_corrupted_device_id_is_atomically_repaired(tmp_path): + import olive.telemetry.deviceid.deviceid as deviceid + + device_id_path = tmp_path / "deviceid" + device_id_path.write_bytes(b"\xff\xfe") + deviceid._device_id_state.update({"device_id": None, "status": deviceid.DeviceIdStatus.NEW}) + + with ( + patch.object(deviceid.platform, "system", return_value="Linux"), + patch.object(deviceid, "get_telemetry_base_dir", return_value=tmp_path), + patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path), + ): + repaired = deviceid.get_device_id() + + assert deviceid._is_valid_device_id(repaired) + assert device_id_path.read_text(encoding="utf-8") == repaired + assert deviceid._device_id_state["status"] == deviceid.DeviceIdStatus.CORRUPTED + + +def test_file_store_does_not_overwrite_concurrent_winner(tmp_path): + with patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path): + store = deviceid_store_mod.Store() + assert store.store_id("first") is True + assert store.store_id("second") is False + + assert (tmp_path / "deviceid").read_text(encoding="utf-8") == "first" + + +def test_concurrent_processes_publish_one_device_id(tmp_path): + script = ( + "import platform; " + "import olive.telemetry.deviceid.deviceid as d; " + "import olive.telemetry.utils as u; " + "platform.system=lambda:'Linux'; " + "u.get_telemetry_base_dir.cache_clear(); " + "print(d.get_device_id())" + ) + env = os.environ.copy() + env["XDG_CACHE_HOME"] = str(tmp_path) + + def run_process(_): + return subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env=env, + timeout=90, + check=False, + ) + + with ThreadPoolExecutor(max_workers=6) as executor: + completed = list(executor.map(run_process, range(6))) + results = [] + for process in completed: + assert process.returncode == 0, process.stderr + results.append(process.stdout.strip()) + + assert len(set(results)) == 1 + + +@pytest.mark.skipif(os.name != "nt", reason="Windows named mutex") +def test_windows_device_id_protocol_serializes_cross_process_publication(tmp_path): + script = """ +import os +from pathlib import Path +import olive.telemetry.deviceid.deviceid as d + +class FileBackedWindowsStore: + @property + def retrieve_id(self): + path = Path(os.environ["OLIVE_TEST_DEVICE_ID"]) + if not path.exists(): + raise FileNotFoundError + return path.read_text(encoding="utf-8") + + def store_id(self, device_id, replace_existing=False): + Path(os.environ["OLIVE_TEST_DEVICE_ID"]).write_text(device_id, encoding="utf-8") + return True + +d.platform.system = lambda: "Windows" +d.WindowsStore = FileBackedWindowsStore +d.get_telemetry_base_dir = lambda: Path(os.environ["OLIVE_TEST_DEVICE_DIR"]) +print(d.get_device_id()) +""" + env = os.environ.copy() + env["OLIVE_TEST_DEVICE_DIR"] = str(tmp_path) + env["OLIVE_TEST_DEVICE_ID"] = str(tmp_path / "registry-deviceid") + + def run_process(_): + return subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env=env, + timeout=90, + check=False, + ) + + with ThreadPoolExecutor(max_workers=6) as executor: + completed = list(executor.map(run_process, range(6))) + results = [] + for process in completed: + assert process.returncode == 0, process.stderr + results.append(process.stdout.strip()) + + assert len(set(results)) == 1 + + +def test_windows_device_id_store_uses_least_privilege_access(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_QUERY_VALUE=0x0001, + KEY_SET_VALUE=0x0002, + KEY_CREATE_SUB_KEY=0x0004, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + ) + winreg.QueryValueEx.side_effect = FileNotFoundError + key_handle = object() + winreg.CreateKeyEx.return_value.__enter__.return_value = key_handle + + with patch.dict("sys.modules", {"winreg": winreg}): + deviceid_store_mod.WindowsStore().store_id("test-device-id") + + winreg.CreateKeyEx.assert_called_once_with( + winreg.HKEY_CURRENT_USER, + deviceid_store_mod.REGISTRY_PATH, + reserved=0, + access=winreg.KEY_QUERY_VALUE | winreg.KEY_SET_VALUE | winreg.KEY_CREATE_SUB_KEY | winreg.KEY_WOW64_64KEY, + ) + winreg.SetValueEx.assert_called_once_with( + key_handle, + deviceid_store_mod.REGISTRY_KEY, + 0, + winreg.REG_SZ, + "test-device-id", + ) + + +def test_windows_device_id_store_preserves_existing_value(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_QUERY_VALUE=0x0001, + KEY_SET_VALUE=0x0002, + KEY_CREATE_SUB_KEY=0x0004, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + ) + winreg.QueryValueEx.return_value = ("existing-device-id", winreg.REG_SZ) + + with patch.dict("sys.modules", {"winreg": winreg}): + deviceid_store_mod.WindowsStore().store_id("new-device-id") + + winreg.SetValueEx.assert_not_called() + + +def test_windows_device_id_store_replaces_existing_value_when_requested(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_QUERY_VALUE=0x0001, + KEY_SET_VALUE=0x0002, + KEY_CREATE_SUB_KEY=0x0004, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + ) + key_handle = object() + winreg.CreateKeyEx.return_value.__enter__.return_value = key_handle + + with patch.dict("sys.modules", {"winreg": winreg}): + deviceid_store_mod.WindowsStore().store_id("new-device-id", replace_existing=True) + + winreg.QueryValueEx.assert_not_called() + winreg.SetValueEx.assert_called_once_with( + key_handle, + deviceid_store_mod.REGISTRY_KEY, + 0, + winreg.REG_SZ, + "new-device-id", + ) + + +def test_windows_device_id_store_rejects_wrong_registry_type(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_READ=0x0001, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + REG_BINARY=3, + ) + winreg.QueryValueEx.return_value = (b"not-a-string", winreg.REG_BINARY) + + with ( + patch.dict("sys.modules", {"winreg": winreg}), + pytest.raises(ValueError, match="not a string"), + ): + _ = deviceid_store_mod.WindowsStore().retrieve_id + + +def test_nested_actions_log_error_once(): + from olive.telemetry.telemetry_extensions import action + + telemetry = MagicMock(accepts_detailed_events=True) + + @action + @action + def fail(): + raise ValueError("boom") + + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions.log_error") as mock_log_error, + pytest.raises(ValueError, match="boom"), + ): + fail() + + mock_log_error.assert_called_once() + + +def test_positional_function_uses_function_action_name(): + from olive.telemetry.telemetry_extensions import action + + telemetry = MagicMock(accepts_detailed_events=True) + + @action + def work(value): + return value + + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.log_action") as mock_log_action, + ): + assert work("value") == "value" + + assert mock_log_action.call_args.kwargs["action_name"] == "work" + + +def test_action_context_without_start_time_reports_zero_duration(): + from olive.telemetry.telemetry_extensions import ActionContext + + telemetry = MagicMock(accepts_detailed_events=True) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.time.perf_counter", return_value=100.0), + patch("olive.telemetry.telemetry_extensions.log_action") as mock_log_action, + ): + context = ActionContext("work") + context.__exit__(None, None, None) + + assert mock_log_action.call_args.kwargs["duration_ms"] == 0 + + +def test_disabled_action_skips_stack_inspection(): + from olive.telemetry.telemetry_extensions import action + + telemetry = MagicMock(accepts_detailed_events=False) + + @action + def work(): + return 42 + + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from") as mock_resolve, + ): + assert work() == 42 + + mock_resolve.assert_not_called() + + +def test_disabled_action_context_skips_stack_inspection(): + from olive.telemetry.telemetry_extensions import ActionContext + + telemetry = MagicMock(accepts_detailed_events=False) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from") as mock_resolve, + patch("olive.telemetry.telemetry_extensions.log_action") as mock_log_action, + ActionContext("work"), + ): + result = 42 + + assert result == 42 + mock_resolve.assert_not_called() + mock_log_action.assert_not_called() diff --git a/test/workflows/test_workflow_run.py b/test/workflows/test_workflow_run.py index 82cc4980bf..0679db7c36 100644 --- a/test/workflows/test_workflow_run.py +++ b/test/workflows/test_workflow_run.py @@ -1,10 +1,20 @@ +# pylint: disable=protected-access +import json import sys from copy import deepcopy from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest +from olive.telemetry.recipe_telemetry import ( + _NO_OVERRIDE, + _build_recipe_hash, + _classify_input_model_source, + _classify_run_config_source, + _extract_config_overrides, + _sanitize_config_snapshot, +) from olive.workflows import run as olive_run from test.utils import ( get_pytorch_model, @@ -125,3 +135,432 @@ def test_run_packages(): # cleanup requirements_file_path.unlink() + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +@patch("olive.telemetry.recipe_telemetry.is_ci_environment", return_value=False) +def test_run_logs_recipe_result_success(_, mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [{"device": "gpu", "execution_providers": ["CUDAExecutionProvider"]}], + } + }, + "engine": {"target": "local_system"}, + "passes": {"dynamic_quant": {"type": "OnnxDynamicQuantization"}}, + } + expected_output = object() + mock_run_engine.return_value = expected_output + + output = olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + assert output is expected_output + mock_log_recipe_result.assert_called_once() + assert mock_log_recipe_result.call_args.args[0] == "Quantize" + assert mock_log_recipe_result.call_args.kwargs["success"] is True + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert metadata["recipe_command"] == "Quantize" + assert metadata["recipe_source"] == "generated_cli" + assert metadata["recipe_format"] == "generated" + assert metadata["workflow_id"] == "default_workflow" + assert metadata["input_model_type"] == "hfmodel" + assert metadata["input_model_source"] == "string_name" + assert metadata["model_task"] == "text-generation" + assert metadata["target_system_type"] == "LocalSystem" + assert metadata["target_device"] == "gpu" + assert metadata["target_execution_provider"] == "CUDAExecutionProvider" + assert metadata["target_execution_providers"] == "CUDAExecutionProvider" + assert metadata["host_system_type"] == "LocalSystem" + assert "host_device" not in metadata + assert "host_execution_provider" not in metadata + assert "host_execution_providers" not in metadata + assert metadata["pass_types"] == "onnxdynamicquantization" + assert metadata["pass_count"] == 1 + assert metadata["data_config_count"] == 0 + assert metadata["search_enabled"] is False + assert metadata["package_config_provided"] is False + assert metadata["is_ci"] is False + assert metadata["recipe_hash"] + assert "input_model_name_hash" not in metadata + assert "config_overrides" not in metadata + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_config_overrides_when_recipe_metadata_provides_overrides(mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + } + } + mock_run_engine.return_value = object() + + olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "WorkflowRun", + "config_overrides": { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + }, + "engine": {"target": "local_system"}, + "data_path": Path("data"), + }, + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + config_overrides = json.loads(metadata["config_overrides"]) + assert config_overrides["input_model"]["model_path"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert config_overrides["engine"]["target"] == "" + assert config_overrides["data_path"] == "" + + +def test_missing_baseline_keys_are_not_reported_as_overrides(): + value = {"passes": {"optimize": {"type": "OrtTransformersOptimization"}}} + baseline = { + "passes": {"optimize": {"type": "OrtTransformersOptimization"}}, + "extra_dependencies": {"cpu": ["onnxruntime"]}, + } + + assert _extract_config_overrides(value, baseline) is _NO_OVERRIDE + + +def test_config_snapshot_redacts_unknown_object_type_names(): + class ProjectSpecificObject: + pass + + assert _sanitize_config_snapshot({"value": ProjectSpecificObject()}) == {"value": ""} + + +def test_recipe_hash_does_not_fingerprint_credentials_or_environment_values(): + first = { + "input_model": {"type": "HfModel", "model_path": "Qwen/Qwen2.5-0.5B-Instruct"}, + "environmentVariables": {"API_KEY": "first-secret"}, + "service": {"clientSecret": "first-client-secret"}, + "auth": "first-auth-secret", + "auth_header": "first-header-secret", + "credential_value": "first-credential-secret", + "docker_env": {"TOKEN": "first-env-secret"}, + } + second = deepcopy(first) + second["environmentVariables"]["API_KEY"] = "second-secret" + second["service"]["clientSecret"] = "second-client-secret" + second["auth"] = "second-auth-secret" + second["auth_header"] = "second-header-secret" + second["credential_value"] = "second-credential-secret" + second["docker_env"]["TOKEN"] = "second-env-secret" + + assert _build_recipe_hash(first) == _build_recipe_hash(second) + + +def test_recipe_hash_does_not_fingerprint_paths_under_generic_keys(): + first = {"script_args": [r"C:\Users\Alice\private.json"], "label": "same"} + second = {"script_args": ["/home/bob/private.json"], "label": "same"} + + assert _build_recipe_hash(first) == _build_recipe_hash(second) + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_recipe_metadata_redacts_credentials_from_preformatted_overrides(mock_run_engine, mock_log_recipe_result): + mock_run_engine.return_value = object() + olive_run( + {"input_model": {"type": "HfModel", "model_path": "Qwen/Qwen2.5-0.5B-Instruct"}}, + recipe_telemetry_metadata={ + "config_overrides": json.dumps( + { + "auth": "token=private", + "endpoint": "https://private.example/model", + "environmentVariables": {"API_KEY": "top-secret"}, + "service": { + "apiKey": "api-secret", + "aws_access_key_id": "aws-secret", + "batch_size": 1, + "clientSecret": "client-secret", + "serviceCredential": "credential-secret", + }, + "modelPath": "private/model.onnx", + } + ), + "package_config_overrides": json.dumps({"feed": {"access_token": "package-secret"}}), + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + serialized = metadata["config_overrides"] + overrides = json.loads(serialized) + assert overrides["environmentVariables"] == "" + assert overrides["auth"] == "" + assert overrides["endpoint"] == "[path]" + assert overrides["modelPath"] == "" + assert overrides["service"] == { + "apiKey": "", + "aws_access_key_id": "", + "batch_size": 1, + "clientSecret": "", + "serviceCredential": "", + } + assert "top-secret" not in serialized + for secret in ("api-secret", "aws-secret", "client-secret", "credential-secret", "private/model.onnx"): + assert secret not in serialized + assert json.loads(metadata["package_config_overrides"]) == {"feed": {"access_token": ""}} + assert "package-secret" not in metadata["package_config_overrides"] + + +@patch("olive.workflows.run.run.log_error") +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_recipe_result_failure(mock_run_engine, mock_log_recipe_result, mock_log_error): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + }, + "passes": {"dynamic_quant": {"type": "OnnxDynamicQuantization"}}, + } + mock_run_engine.side_effect = ValueError("recipe failed") + + with pytest.raises(ValueError, match="recipe failed"): + olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + mock_log_recipe_result.assert_called_once() + assert mock_log_recipe_result.call_args.args[0] == "Quantize" + assert mock_log_recipe_result.call_args.kwargs["success"] is False + assert "exception_type" not in mock_log_recipe_result.call_args.kwargs + mock_log_error.assert_called_once() + assert mock_log_error.call_args.kwargs["exception_type"] == "ValueError" + assert "recipe failed" in mock_log_error.call_args.kwargs["exception_message"] + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_skips_recipe_result_when_recipe_telemetry_is_not_emitted(mock_run_engine, mock_log_recipe_result): + expected_output = object() + mock_run_engine.return_value = expected_output + + output = olive_run( + { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + } + }, + emit_recipe_telemetry=False, + ) + + assert output is expected_output + mock_log_recipe_result.assert_not_called() + + +@patch("olive.workflows.run.run.is_ci_environment", return_value=True) +@patch("olive.workflows.run.run.Telemetry.get_existing_instance") +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_programmatic_ci_run_flushes_recipe_with_bounded_shutdown( + mock_run_engine, mock_log_recipe_result, mock_get_telemetry, _mock_is_ci +): + expected_output = object() + telemetry = Mock() + mock_run_engine.return_value = expected_output + mock_get_telemetry.return_value = telemetry + + output = olive_run( + { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + } + } + ) + + assert output is expected_output + mock_log_recipe_result.assert_called_once() + telemetry.shutdown.assert_called_once_with() + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.systems.system_config.SystemConfig.create_system") +def test_run_logs_single_parent_recipe_result_for_docker_host(mock_create_system, mock_log_recipe_result): + expected_output = object() + docker_system = Mock() + + def run_workflow(container_run_config): + container_run_config.engine.host = container_run_config.engine.target + return expected_output + + docker_system.run_workflow.side_effect = run_workflow + mock_create_system.return_value = docker_system + config = { + "input_model": {"type": "ONNXModel", "model_path": "model.onnx"}, + "systems": { + "docker_system": { + "type": "Docker", + "config": { + "dockerfile": "Dockerfile", + "build_context_path": "build_context", + "image_name": "test-image:latest", + "work_dir": "/olive-ws", + }, + }, + "local_system": {"type": "LocalSystem"}, + }, + "engine": {"host": "docker_system", "target": "local_system"}, + } + + output = olive_run(config) + + assert output is expected_output + mock_log_recipe_result.assert_called_once() + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert metadata["host_system_type"] == "Docker" + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_recipe_host_metadata_without_explicit_target(mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + }, + "systems": { + "host_system": { + "type": "LocalSystem", + "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], + } + }, + "engine": {"host": "host_system"}, + } + mock_run_engine.return_value = object() + + olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert "target_system_type" not in metadata + assert "target_device" not in metadata + assert "target_execution_provider" not in metadata + assert "target_execution_providers" not in metadata + assert metadata["host_system_type"] == "LocalSystem" + assert metadata["host_device"] == "cpu" + assert metadata["host_execution_provider"] == "CPUExecutionProvider" + assert metadata["host_execution_providers"] == "CPUExecutionProvider" + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_package_config_overrides_when_package_config_provided(mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + } + } + mock_run_engine.return_value = object() + + olive_run( + config, + package_config={ + "passes": { + "AddOliveMetadata": { + "module_path": "olive.passes.onnx.add_metadata.AddOliveMetadata", + "supported_providers": ["CPUExecutionProvider"], + } + }, + "extra_dependencies": {"custom_accelerator": ["custom-package"]}, + }, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert metadata["package_config_provided"] is True + package_config_overrides = json.loads(metadata["package_config_overrides"]) + assert package_config_overrides["passes"][0]["supported_providers"] == ["CPUExecutionProvider"] + assert "module_path" not in package_config_overrides["passes"][0] + assert package_config_overrides["extra_dependencies"]["custom_accelerator"] == ["custom-package"] + + +def test_classify_run_config_source_handles_non_pathlike_object(): + assert _classify_run_config_source(object()) == ("config_object", "object") + + +def test_classify_input_model_source_does_not_depend_on_local_filesystem(tmp_path, monkeypatch): + assert _classify_input_model_source("Qwen/Qwen2.5-0.5B-Instruct") == "string_name" + + monkeypatch.chdir(tmp_path) + (tmp_path / "bert-base-uncased").mkdir() + + assert _classify_input_model_source("bert-base-uncased") == "string_name" + assert _classify_input_model_source("./model.onnx") == "local_file" + assert _classify_input_model_source("model.onnx") == "local_file" + + +def test_recipe_hash_does_not_depend_on_local_model_path_presence(tmp_path, monkeypatch): + config = { + "input_model": {"type": "HfModel", "config": {"model_path": "bert-base-uncased"}}, + "engine": {"output_dir": "output"}, + } + recipe_hash = _build_recipe_hash(config) + + monkeypatch.chdir(tmp_path) + (tmp_path / "bert-base-uncased").mkdir() + + assert _build_recipe_hash(config) == recipe_hash + + +def test_recipe_hash_handles_path_values(): + config = { + "input_model": {"type": "HfModel", "config": {"model_path": Path("model")}}, + "custom_value": Path("custom"), + } + + assert _build_recipe_hash(config)