From c833e33193c67b197a852ac8687052b04670965e Mon Sep 17 00:00:00 2001 From: Jordan Gonzalez <30836115+duncanista@users.noreply.github.com> Date: Mon, 2 Feb 2026 14:09:32 -0500 Subject: [PATCH 1/2] omit creating cold start on LMI --- datadog_lambda/cold_start.py | 20 ++++++++++++++++ datadog_lambda/wrapper.py | 9 ++++++- tests/test_cold_start.py | 46 ++++++++++++++++++++++++++++++++++++ 3 files changed, 74 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py index a40e2fcb3..6f0031ba3 100644 --- a/datadog_lambda/cold_start.py +++ b/datadog_lambda/cold_start.py @@ -1,3 +1,4 @@ +import os import time from typing import List, Hashable import logging @@ -63,6 +64,21 @@ def get_proactive_init_tag(): ) +def is_managed_instances_mode(): + """ + Checks if the Lambda function is running in managed instances mode. + In managed instances mode, we should not create cold start tracing spans + as the gap between the sandbox initialization and the first + invocation might not be be a great experience. + + Returns: + bool: True if running in managed instances mode, False otherwise + """ + return ( + os.environ.get("AWS_LAMBDA_INITIALIZATION_TYPE") == "lambda-managed-instances" + ) + + class ImportNode(object): def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None): self.module_name = module_name @@ -145,6 +161,10 @@ def wrapped_find_spec(*args, **kwargs): def initialize_cold_start_tracing(): + # Skip cold start tracing initialization in managed instances mode + if is_managed_instances_mode(): + return + if is_new_sandbox() and config.cold_start_tracing: from sys import meta_path diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 5a32e2ff9..b2f83f138 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -16,6 +16,7 @@ is_cold_start, is_proactive_init, is_new_sandbox, + is_managed_instances_mode, ColdStartTracer, ) from datadog_lambda.config import config @@ -338,7 +339,13 @@ def _after(self, event, context): create_dd_dummy_metadata_subsegment( self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY ) - should_trace_cold_start = config.cold_start_tracing and is_new_sandbox() + # Skip creating cold start spans in managed instances mode + # In managed instances, the tracer library handles cold start independently + should_trace_cold_start = ( + config.cold_start_tracing + and is_new_sandbox() + and not is_managed_instances_mode() + ) if should_trace_cold_start: trace_ctx = tracer.current_trace_context() diff --git a/tests/test_cold_start.py b/tests/test_cold_start.py index d75b5f43a..858db9a51 100644 --- a/tests/test_cold_start.py +++ b/tests/test_cold_start.py @@ -12,6 +12,52 @@ class TestColdStartTracingSetup(unittest.TestCase): + def test_is_managed_instances_mode_when_set(self): + os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] = "lambda-managed-instances" + self.assertTrue(cold_start.is_managed_instances_mode()) + # Clean up + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + + def test_is_not_managed_instances_mode_when_not_set(self): + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + self.assertFalse(cold_start.is_managed_instances_mode()) + + def test_is_not_managed_instances_mode_with_different_value(self): + os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] = "on-demand" + self.assertFalse(cold_start.is_managed_instances_mode()) + # Clean up + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + + def test_initialize_cold_start_tracing_skips_in_managed_instances(self): + # Set managed instances mode + os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] = "lambda-managed-instances" + os.environ["DD_COLD_START_TRACING"] = "true" + cold_start._cold_start = True + cold_start._lambda_container_initialized = False + + # Reset node stacks and wrapped loaders to get clean state + cold_start.reset_node_stacks() + cold_start.already_wrapped_loaders.clear() + + # Count wrapped loaders before + wrapped_loaders_before = len(cold_start.already_wrapped_loaders) + + # Initialize cold start tracing - should skip in managed instances mode + cold_start.initialize_cold_start_tracing() + + # Verify no loaders were wrapped + wrapped_loaders_after = len(cold_start.already_wrapped_loaders) + self.assertEqual(wrapped_loaders_before, wrapped_loaders_after) + + # Clean up + if "AWS_LAMBDA_INITIALIZATION_TYPE" in os.environ: + del os.environ["AWS_LAMBDA_INITIALIZATION_TYPE"] + if "DD_COLD_START_TRACING" in os.environ: + del os.environ["DD_COLD_START_TRACING"] + def test_proactive_init(self): cold_start._cold_start = True cold_start._proactive_initialization = False From 21b61b8ec17c791f61a8f23944ebcb6a97f6b2cf Mon Sep 17 00:00:00 2001 From: Jordan Gonzalez <30836115+duncanista@users.noreply.github.com> Date: Mon, 2 Feb 2026 14:31:58 -0500 Subject: [PATCH 2/2] set env in config instead of using os to do it --- datadog_lambda/cold_start.py | 5 +---- datadog_lambda/config.py | 1 + 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py index 6f0031ba3..260d491b8 100644 --- a/datadog_lambda/cold_start.py +++ b/datadog_lambda/cold_start.py @@ -1,4 +1,3 @@ -import os import time from typing import List, Hashable import logging @@ -74,9 +73,7 @@ def is_managed_instances_mode(): Returns: bool: True if running in managed instances mode, False otherwise """ - return ( - os.environ.get("AWS_LAMBDA_INITIALIZATION_TYPE") == "lambda-managed-instances" - ) + return config.aws_lambda_initialization_type == "lambda-managed-instances" class ImportNode(object): diff --git a/datadog_lambda/config.py b/datadog_lambda/config.py index eda6b5828..9960953ec 100644 --- a/datadog_lambda/config.py +++ b/datadog_lambda/config.py @@ -98,6 +98,7 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False): integration_test = _get_env("DD_INTEGRATION_TEST", "false", as_bool) aws_lambda_function_name = _get_env("AWS_LAMBDA_FUNCTION_NAME") + aws_lambda_initialization_type = _get_env("AWS_LAMBDA_INITIALIZATION_TYPE") @property def function_name(self):