From e283ed0097487d86904232947fd2d22ad83b75aa Mon Sep 17 00:00:00 2001 From: dangmaul-amazon Date: Thu, 25 Mar 2021 02:35:21 -0700 Subject: [PATCH 1/6] Add create_profiling_group call in refresh_configuration and report() As part of Lambda 1-click integration, we want the agent to create a profiling group for the user if it does not exist. To do this, we add a create_profiling_group API call in places where we could receive a ResourceNotFoundException for the PG. --- .../agent_metadata/aws_lambda.py | 8 +- codeguru_profiler_agent/local_aggregator.py | 14 +- .../sdk_reporter/sdk_reporter.py | 67 ++++++- test/acceptance/test_end_to_end_profiling.py | 3 + test/acceptance/test_live_profiling.py | 4 + test/unit/sdk_reporter/test_sdk_reporter.py | 168 +++++++++++++++++- test/unit/test_local_aggregator.py | 3 +- 7 files changed, 260 insertions(+), 7 deletions(-) diff --git a/codeguru_profiler_agent/agent_metadata/aws_lambda.py b/codeguru_profiler_agent/agent_metadata/aws_lambda.py index 0056b15..b99b2fa 100644 --- a/codeguru_profiler_agent/agent_metadata/aws_lambda.py +++ b/codeguru_profiler_agent/agent_metadata/aws_lambda.py @@ -2,6 +2,7 @@ import logging import uuid +from mock import MagicMock from codeguru_profiler_agent.agent_metadata.fleet_info import FleetInfo from codeguru_profiler_agent.aws_lambda.lambda_context import LambdaContext @@ -100,7 +101,12 @@ def get_metadata_for_configure_agent_call(self, lambda_context=None): as_map[LAMBDA_MEMORY_LIMIT_IN_MB_KEY] = str(self.memory_limit_mb) if self.execution_env: as_map[EXECUTION_ENVIRONMENT_KEY] = self.execution_env - if lambda_context.context is not None: + + ''' + Adding a specific condition to ignore MagicMock instances from being added to the metadata since + it causes boto to raise a ParamValidationError, similar to https://github.com/boto/botocore/issues/2063. + ''' + if lambda_context.context is not None and not isinstance(lambda_context.context, MagicMock): as_map[AWS_REQUEST_ID_KEY] = lambda_context.context.aws_request_id as_map[LAMBDA_REMAINING_TIME_IN_MILLISECONDS_KEY] = \ str(lambda_context.context.get_remaining_time_in_millis()) diff --git a/codeguru_profiler_agent/local_aggregator.py b/codeguru_profiler_agent/local_aggregator.py index a04c304..dda3b9c 100644 --- a/codeguru_profiler_agent/local_aggregator.py +++ b/codeguru_profiler_agent/local_aggregator.py @@ -6,6 +6,7 @@ from codeguru_profiler_agent.metrics.with_timer import with_timer from codeguru_profiler_agent.model.profile import Profile from codeguru_profiler_agent.utils.time import current_milli_time +from codeguru_profiler_agent.sdk_reporter.sdk_reporter import SdkReporter logger = logging.getLogger(__name__) @@ -100,13 +101,24 @@ def refresh_configuration(self): self.reporter.refresh_configuration() def _report_profile(self, now): + current_last_report_attempted_value = self.last_report_attempted self.last_report_attempted = now self._add_overhead_metric_to_profile() logger.info("Attempting to report profile data: " + str(self.profile)) if self.profile.is_empty(): logger.info("Report was cancelled because it was empty") return False - return self.reporter.report(self.profile) + is_reporting_successful = self.reporter.report(self.profile) + ''' + If we attempt to create a PG in the report() call, we do not want to update the last_report_attempted_value + since we did not actually report a profile. + + This will occur only in the case of Lambda 1-click integration. + ''' + if SdkReporter.check_create_pg_called_during_submit_profile == True: + self.last_report_attempted = current_last_report_attempted_value + SdkReporter.reset_check_create_pg_called_during_submit_profile_flag() + return is_reporting_successful def _is_under_min_reporting_time(self, now): return AgentConfiguration.get().is_under_min_reporting_time(now - self.last_report_attempted) diff --git a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py index d573ead..1221af7 100644 --- a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py +++ b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py @@ -2,6 +2,7 @@ import logging import io +import os from botocore.exceptions import ClientError from codeguru_profiler_agent.utils.log_exception import log_exception @@ -10,14 +11,14 @@ from codeguru_profiler_agent.sdk_reporter.profile_encoder import ProfileEncoder logger = logging.getLogger(__name__) - +HANDLER_ENV_NAME_FOR_CODEGURU = "HANDLER_ENV_NAME_FOR_CODEGURU" class SdkReporter(Reporter): """ Handles communication with the CodeGuru Profiler Service backend. Encodes profiles using the ProfilerEncoder and reports them using the CodeGuru profiler SDK. """ - + is_create_pg_called_during_submit_profile = False def __init__(self, environment): """ :param environment: dependency container dictionary for the current profiler. @@ -32,6 +33,8 @@ def __init__(self, environment): self.timer = environment.get("timer") self.metadata = environment["agent_metadata"] self.agent_config_merger = environment["agent_config_merger"] + self.is_lambda_one_click_pg_created_during_execution = False + # self.is_create_pg_called_during_submit_profile = False def _encode_profile(self, profile): output_profile_stream = io.BytesIO() @@ -67,8 +70,13 @@ def refresh_configuration(self): # whole process because the customer may fix this on their side by creating/changing the profiling group. # We handle service exceptions like this in boto3 # see https://boto3.amazonaws.com/v1/documentation/api/latest/guide/error-handling.html - if error.response['Error']['Code'] in ['ResourceNotFoundException', 'ValidationException']: + if error.response['Error']['Code'] == 'ValidationException': self.agent_config_merger.disable_profiling() + if error.response['Error']['Code'] == 'ResourceNotFoundException': + if self.is_lambda_one_click_integration_active(): + self.create_pg_when_one_click_integration_is_active() + else: + self.agent_config_merger.disable_profiling() self._log_request_failed(operation="configure_agent", exception=error) except Exception as e: self._log_request_failed(operation="configure_agent", exception=e) @@ -90,11 +98,64 @@ def report(self, profile): ) logger.info("Reported profile successfully") return True + except ClientError as error: + if error.response['Error']['Code'] == 'ResourceNotFoundException': + if self.is_lambda_one_click_integration_active(): + global is_create_pg_called_during_submit_profile + is_create_pg_called_during_submit_profile = True + self.create_pg_when_one_click_integration_is_active() + self._log_request_failed(operation="post_agent_profile", exception=error) + return False except Exception as e: self._log_request_failed(operation="post_agent_profile", exception=e) return False + @with_timer("createOneClickPG", measurement="wall-clock-time") + def create_pg_when_one_click_integration_is_active(self): + """ + Create a PG for the Lambda function onboarded with 1-click integration + """ + function_name = str(self.metadata.fleet_info.function_arn).split(':')[6] + + profiling_group_name_for_1_click = "aws-lambda-{}".format(function_name) + profiling_group_name_to_be_created = os.getenv("AWS_CODEGURU_PROFILER_GROUP_NAME", + profiling_group_name_for_1_click) + try: + self.codeguru_client_builder.codeguru_client.create_profiling_group( + profilingGroupName=profiling_group_name_to_be_created, + computePlatform='AWSLambda' + ) + self.profiling_group_name = profiling_group_name_to_be_created + self.is_lambda_one_click_pg_created_during_execution = True + logger.info("Created Lambda Profiling Group with name {}", profiling_group_name_to_be_created) + + except Exception as e: + self._log_request_failed(operation="create_profiling_group", exception=e) + + def is_lambda_one_click_integration_active(self): + """ + Check if the ComputeType is AWSLambda and if the environment + variables for Lambda Layer Profiling are set + """ + fleet_info_as_map = self.metadata.fleet_info.serialize_to_map() + if 'computeType' in fleet_info_as_map: + if fleet_info_as_map['computeType'] == 'aws_lambda': + handler_name_ = os.environ.get(HANDLER_ENV_NAME_FOR_CODEGURU) + if not handler_name_: + logger.error("Env Variables for CodeGuru Profiler Lambda Layer are not set. Cannot create PG.") + return False + return True + return False + @staticmethod def _log_request_failed(operation, exception): log_exception(logger, "Failed to call the CodeGuru Profiler service for the {} operation: {}" .format(operation, str(exception))) + + @classmethod + def check_create_pg_called_during_submit_profile(cls): + return cls.is_create_pg_called_during_submit_profile + + @classmethod + def reset_check_create_pg_called_during_submit_profile_flag(cls): + cls.is_create_pg_called_during_submit_profile = False diff --git a/test/acceptance/test_end_to_end_profiling.py b/test/acceptance/test_end_to_end_profiling.py index 771df02..c69edfb 100644 --- a/test/acceptance/test_end_to_end_profiling.py +++ b/test/acceptance/test_end_to_end_profiling.py @@ -33,6 +33,9 @@ def test_report_when_stopped(self): with \ patch( "codeguru_profiler_agent.reporter.agent_configuration.AgentConfiguration.is_under_min_reporting_time", + return_value=False), \ + patch( + "codeguru_profiler_agent.sdk_reporter.sdk_reporter.SdkReporter.check_create_pg_called_during_submit_profile", return_value=False): with self.client_stubber: self.profiler.start() diff --git a/test/acceptance/test_live_profiling.py b/test/acceptance/test_live_profiling.py index 3450815..fe09eb8 100644 --- a/test/acceptance/test_live_profiling.py +++ b/test/acceptance/test_live_profiling.py @@ -4,6 +4,7 @@ from mock import patch from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration +from codeguru_profiler_agent.sdk_reporter.sdk_reporter import SdkReporter from codeguru_profiler_agent.profiler import Profiler from codeguru_profiler_agent.agent_metadata.agent_metadata import AgentMetadata, DefaultFleetInfo from test.help_utils import DUMMY_TEST_PROFILING_GROUP_NAME @@ -16,6 +17,9 @@ def test_live_profiling(self): patch( "codeguru_profiler_agent.reporter.agent_configuration.AgentConfiguration.is_under_min_reporting_time", return_value=False), \ + patch( + "codeguru_profiler_agent.sdk_reporter.sdk_reporter.SdkReporter.check_create_pg_called_during_submit_profile", + return_value=False), \ patch( "codeguru_profiler_agent.reporter.agent_configuration.AgentConfiguration._is_reporting_interval_smaller_than_minimum_allowed", return_value=False): diff --git a/test/unit/sdk_reporter/test_sdk_reporter.py b/test/unit/sdk_reporter/test_sdk_reporter.py index 0113fcf..086c399 100644 --- a/test/unit/sdk_reporter/test_sdk_reporter.py +++ b/test/unit/sdk_reporter/test_sdk_reporter.py @@ -1,7 +1,11 @@ # -*- coding: utf-8 -*- +import os + import boto3 -from datetime import timedelta +from datetime import timedelta, datetime + +from codeguru_profiler_agent.agent_metadata.aws_lambda import AWSLambda from codeguru_profiler_agent.utils.time import current_milli_time from test.pytestutils import before from mock import MagicMock @@ -16,6 +20,9 @@ from codeguru_profiler_agent.codeguru_client_builder import CodeGuruClientBuilder profiling_group_name = "test-ProfilingGroup-name" +lambda_one_click_profiling_group_name = "aws-lambda-testLambdaName" +test_agent_metadata_for_lambda = AgentMetadata( + fleet_info=AWSLambda("arn:aws:lambda:us-east-1:111111111111:function:testLambdaName", "memory", "env", "agentId")) profile = Profile(profiling_group_name, 1.0, 0.5, current_milli_time()) @@ -72,6 +79,110 @@ def test_return_false_when_report_throws_error(self): with self.client_stubber: assert self.subject.report(profile) is False + def test_creates_pg_if_onboarded_with_lambda_one_click_integration(self): + expected_params = { + 'agentProfile': ANY, + 'contentType': 'application/json', + 'profilingGroupName': profiling_group_name + } + self.client_stubber.add_client_error('post_agent_profile', + service_error_code='ResourceNotFoundException', + service_message='Simulated ResourceNotFoundException in ' + 'post_agent_profile call', + expected_params=expected_params) + + expected_response_create_pg = { + 'profilingGroup': { + 'agentOrchestrationConfig': { + 'profilingEnabled': True + }, + 'arn': 'string', + 'computePlatform': 'AWSLambda', + 'createdAt': datetime(2015, 1, 1), + 'name': 'string', + 'profilingStatus': { + 'latestAgentOrchestratedAt': datetime(2015, 1, 1), + 'latestAgentProfileReportedAt': datetime(2015, 1, 1), + 'latestAggregatedProfile': { + 'period': 'PT5M', + 'start': datetime(2015, 1, 1) + } + }, + 'tags': { + 'string': 'string' + }, + 'updatedAt': datetime(2015, 1, 1) + } + } + self.client_stubber.add_response('create_profiling_group', expected_response_create_pg) + + expected_params_one_click = { + 'agentProfile': ANY, + 'contentType': 'application/json', + 'profilingGroupName': lambda_one_click_profiling_group_name + } + self.client_stubber.add_response('post_agent_profile', {}, expected_params_one_click) + + os.environ.__setitem__('AWS_CODEGURU_PROFILER_GROUP_NAME', lambda_one_click_profiling_group_name) + os.environ.__setitem__('HANDLER_ENV_NAME_FOR_CODEGURU', 'test-handler') + self.subject.metadata = test_agent_metadata_for_lambda + with self.client_stubber: + assert self.subject.report(profile) is False + assert self.subject.report(profile) is True + + def test_create_pg_not_invoked_in_non_lambda_case(self): + expected_params = { + 'agentProfile': ANY, + 'contentType': 'application/json', + 'profilingGroupName': profiling_group_name + } + self.client_stubber.add_client_error('post_agent_profile', + service_error_code="ResourceNotFoundException", + service_message='Simulated ResourceNotFoundException in ' + 'post_agent_profile call', + expected_params=expected_params) + + expected_params_post_agent_profile = { + 'agentProfile': ANY, + 'contentType': 'application/json', + 'profilingGroupName': profiling_group_name + } + self.client_stubber.add_response('post_agent_profile', {}, expected_params_post_agent_profile) + + expected_response_create_pg = { + 'profilingGroup': { + 'agentOrchestrationConfig': { + 'profilingEnabled': True + }, + 'arn': 'string', + 'computePlatform': 'AWSLambda', + 'createdAt': datetime(2015, 1, 1), + 'name': 'string', + 'profilingStatus': { + 'latestAgentOrchestratedAt': datetime(2015, 1, 1), + 'latestAgentProfileReportedAt': datetime(2015, 1, 1), + 'latestAggregatedProfile': { + 'period': 'PT5M', + 'start': datetime(2015, 1, 1) + } + }, + 'tags': { + 'string': 'string' + }, + 'updatedAt': datetime(2015, 1, 1) + } + } + self.client_stubber.add_response('create_profiling_group', expected_response_create_pg) + + with self.client_stubber: + assert self.subject.report(profile) is False + assert self.subject.report(profile) is True + try: + self.client_stubber.assert_no_pending_responses() + assert False + except AssertionError: + assert True + class TestConfigureAgent(TestSdkReporter): @before @@ -117,3 +228,58 @@ def test_when_backend_sends_validation_exception_it_stops_the_profiling(self): with self.client_stubber: self.subject.refresh_configuration() assert AgentConfiguration.get().should_profile is False + + def test_creates_lambda_pg_if_onboarded_with_lambda_one_click_integration(self): + self.client_stubber.add_client_error('configure_agent', + service_error_code='ResourceNotFoundException', + service_message='Simulated ResourceNotFoundException in ' + 'configure_agent call') + + expected_response_create_pg = { + 'profilingGroup': { + 'agentOrchestrationConfig': { + 'profilingEnabled': True + }, + 'arn': 'string', + 'computePlatform': 'AWSLambda', + 'createdAt': datetime(2015, 1, 1), + 'name': 'string', + 'profilingStatus': { + 'latestAgentOrchestratedAt': datetime(2015, 1, 1), + 'latestAgentProfileReportedAt': datetime(2015, 1, 1), + 'latestAggregatedProfile': { + 'period': 'PT5M', + 'start': datetime(2015, 1, 1) + } + }, + 'tags': { + 'string': 'string' + }, + 'updatedAt': datetime(2015, 1, 1) + } + } + self.client_stubber.add_response('create_profiling_group', expected_response_create_pg) + + os.environ.__setitem__('AWS_LAMBDA_FUNCTION_MEMORY_SIZE', "512") + os.environ.__setitem__('AWS_EXECUTION_ENV', "AWS_Lambda_python3.8") + os.environ.__setitem__('AWS_CODEGURU_PROFILER_GROUP_NAME', lambda_one_click_profiling_group_name) + os.environ.__setitem__('HANDLER_ENV_NAME_FOR_CODEGURU', 'test-handler') + self.subject.metadata = test_agent_metadata_for_lambda + with self.client_stubber: + self.subject.refresh_configuration() + assert self.subject.is_lambda_one_click_pg_created_during_execution is True + + def test_create_pg_not_invoked_in_non_lambda_case(self): + self.client_stubber.add_client_error('configure_agent', + service_error_code='ResourceNotFoundException', + service_message='Simulated ResourceNotFoundException in ' + 'configure_agent call') + + with self.client_stubber: + self.subject.refresh_configuration() + assert self.subject.is_lambda_one_click_pg_created_during_execution is False + try: + self.client_stubber.assert_no_pending_responses() + assert False + except AssertionError: + assert True diff --git a/test/unit/test_local_aggregator.py b/test/unit/test_local_aggregator.py index b1aec6a..d41b336 100644 --- a/test/unit/test_local_aggregator.py +++ b/test/unit/test_local_aggregator.py @@ -5,7 +5,7 @@ from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration from codeguru_profiler_agent.utils.time import current_milli_time from test.pytestutils import before -from mock import MagicMock, call +from mock import MagicMock, call, patch from codeguru_profiler_agent.profiler import DEFAULT_REPORTING_INTERVAL, \ DEFAULT_MEMORY_LIMIT_BYTES, Profiler, INITIAL_MINIMUM_REPORTING_INTERVAL, DEFAULT_SAMPLING_INTERVAL @@ -42,6 +42,7 @@ class TestLocalAggregator: def before(self): self.mock_reporter = MagicMock(name="reporter", spec=SdkReporter) + self.mock_reporter.lambda_one_click_pg_created_during_execution = False self.mock_profile = MagicMock(name="profile", spec=Profile) self.mock_profile_factory = MagicMock( name="profile_factory", From 538f06bc3de4adb7696ff06f59ba3a36ad219ab3 Mon Sep 17 00:00:00 2001 From: dangmaul-amazon Date: Fri, 26 Mar 2021 15:07:00 -0700 Subject: [PATCH 2/6] Update sdk_reporter.py --- codeguru_profiler_agent/sdk_reporter/sdk_reporter.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py index 1221af7..86536c1 100644 --- a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py +++ b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py @@ -34,7 +34,6 @@ def __init__(self, environment): self.metadata = environment["agent_metadata"] self.agent_config_merger = environment["agent_config_merger"] self.is_lambda_one_click_pg_created_during_execution = False - # self.is_create_pg_called_during_submit_profile = False def _encode_profile(self, profile): output_profile_stream = io.BytesIO() @@ -103,8 +102,8 @@ def report(self, profile): if self.is_lambda_one_click_integration_active(): global is_create_pg_called_during_submit_profile is_create_pg_called_during_submit_profile = True + logger.info("Received ResourceNotFoundException. Attempting to create 1-click PG in Report Profile") self.create_pg_when_one_click_integration_is_active() - self._log_request_failed(operation="post_agent_profile", exception=error) return False except Exception as e: self._log_request_failed(operation="post_agent_profile", exception=e) From bc2d39853e75e9716aa33d1cabe54aa6f419a5b6 Mon Sep 17 00:00:00 2001 From: dangmaul-amazon Date: Tue, 30 Mar 2021 22:27:52 -0700 Subject: [PATCH 3/6] Fixes and addressing comments --- .../agent_metadata/aws_lambda.py | 2 +- .../aws_lambda/profiler_decorator.py | 3 +- codeguru_profiler_agent/profiler_builder.py | 16 +++-- .../sdk_reporter/sdk_reporter.py | 68 ++++++++++--------- ...est_end_to_end_profile_and_save_to_file.py | 2 +- test/acceptance/test_end_to_end_profiling.py | 2 +- test/acceptance/test_live_profiling.py | 2 +- test/unit/agent_metadata/test_aws_lambda.py | 2 +- .../aws_lambda/test_profiler_decorator.py | 4 +- test/unit/file_reporter/test_file_reporter.py | 4 +- test/unit/model/test_call_graph_node.py | 2 +- test/unit/model/test_profile.py | 2 +- .../sdk_reporter/test_sdk_profile_encoder.py | 2 +- test/unit/sdk_reporter/test_sdk_reporter.py | 50 +++++++------- test/unit/test_codeguru_client_builder.py | 2 +- test/unit/test_local_aggregator.py | 2 +- test/unit/test_profiler.py | 2 +- test/unit/test_profiler_disabler.py | 2 +- test/unit/test_profiler_runner.py | 2 +- test/unit/test_sampler.py | 2 +- test/unit/test_sampling_utils.py | 2 +- test/unit/utils/test_execution_state.py | 2 +- 22 files changed, 98 insertions(+), 79 deletions(-) diff --git a/codeguru_profiler_agent/agent_metadata/aws_lambda.py b/codeguru_profiler_agent/agent_metadata/aws_lambda.py index b99b2fa..33afbbf 100644 --- a/codeguru_profiler_agent/agent_metadata/aws_lambda.py +++ b/codeguru_profiler_agent/agent_metadata/aws_lambda.py @@ -2,7 +2,7 @@ import logging import uuid -from mock import MagicMock +from unittest.mock import MagicMock from codeguru_profiler_agent.agent_metadata.fleet_info import FleetInfo from codeguru_profiler_agent.aws_lambda.lambda_context import LambdaContext diff --git a/codeguru_profiler_agent/aws_lambda/profiler_decorator.py b/codeguru_profiler_agent/aws_lambda/profiler_decorator.py index 4c8ad40..83489be 100644 --- a/codeguru_profiler_agent/aws_lambda/profiler_decorator.py +++ b/codeguru_profiler_agent/aws_lambda/profiler_decorator.py @@ -15,7 +15,8 @@ def _create_lambda_profiler(profiling_group_name, region_name, environment_overr from codeguru_profiler_agent.agent_metadata.aws_lambda import AWSLambda override = {'agent_metadata': AgentMetadata(AWSLambda.look_up_metadata(context))} override.update(environment_override) - profiler = build_profiler(pg_name=profiling_group_name, region_name=region_name, override=override, env=env) + profiler = build_profiler(pg_name=profiling_group_name, region_name=region_name, override=override, env=env, + should_autocreate_profiling_group=True) if profiler is None: return _EmptyProfiler() return profiler diff --git a/codeguru_profiler_agent/profiler_builder.py b/codeguru_profiler_agent/profiler_builder.py index 617feec..8d074a7 100644 --- a/codeguru_profiler_agent/profiler_builder.py +++ b/codeguru_profiler_agent/profiler_builder.py @@ -17,6 +17,9 @@ CREDENTIAL_PATH = "AWS_CODEGURU_PROFILER_CREDENTIAL_PATH" ENABLED_ENV = "AWS_CODEGURU_PROFILER_ENABLED" +# Environment variables provided by AWS Lambda +AWS_LAMBDA_FUNCTION_NAME_ENV_VAR_KEY = "AWS_LAMBDA_FUNCTION_NAME" + # non documented parameters SAMPLING_INTERVAL = "AWS_CODEGURU_PROFILER_SAMPLING_INTERVAL_MS" REPORTING_INTERVAL = "AWS_CODEGURU_PROFILER_REPORTING_INTERVAL_MS" @@ -111,7 +114,8 @@ def _check_credential_through_environment(env=os.environ): def build_profiler(pg_name=None, region_name=None, credential_profile=None, - env=os.environ, session_factory=boto3.session.Session, profiler_factory=None, override=None): + env=os.environ, session_factory=boto3.session.Session, profiler_factory=None, override=None, + should_autocreate_profiling_group=False): """ Creates a Profiler object from given parameters or environment variables :param pg_name: given profiling group name, default is None @@ -120,6 +124,7 @@ def build_profiler(pg_name=None, region_name=None, credential_profile=None, :param env: environment variables are used if parameters are not provided, default is os.environ :param session_factory: (For testing) function for creating boto3.session.Session, default is boto3.session.Session :param override: a dictionary with possible extra parameters to override default values + :param should_autocreate_profiling_group: True when Compute Platform is AWS Lambda. False otherwise :return: a Profiler object or None, this function does not throw exceptions """ if profiler_factory is None: @@ -137,9 +142,12 @@ def build_profiler(pg_name=None, region_name=None, credential_profile=None, name_from_arn, region_from_arn, _account_id = _read_profiling_group_arn(env) profiling_group_name = _get_profiling_group_name(pg_name, name_from_arn, env) if not profiling_group_name: - logger.info("Could not find a profiling group name to start the CodeGuru Profiler agent. " - + "Add command line argument or environment variable. e.g. " + PG_ARN_ENV) - return None + if should_autocreate_profiling_group: + profiling_group_name = "aws-lambda-" + env.get(AWS_LAMBDA_FUNCTION_NAME_ENV_VAR_KEY) + else: + logger.info("Could not find a profiling group name to start the CodeGuru Profiler agent. " + + "Add command line argument or environment variable. e.g. " + PG_ARN_ENV) + return None region = _get_region(region_name, region_from_arn, env) session = session_factory(region_name=region, profile_name=credential_profile) diff --git a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py index 86536c1..a80210f 100644 --- a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py +++ b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py @@ -11,7 +11,8 @@ from codeguru_profiler_agent.sdk_reporter.profile_encoder import ProfileEncoder logger = logging.getLogger(__name__) -HANDLER_ENV_NAME_FOR_CODEGURU = "HANDLER_ENV_NAME_FOR_CODEGURU" +HANDLER_ENV_NAME_FOR_CODEGURU_KEY = "HANDLER_ENV_NAME_FOR_CODEGURU" +AWS_EXECUTION_ENV_KEY = "AWS_EXECUTION_ENV" class SdkReporter(Reporter): """ @@ -33,7 +34,7 @@ def __init__(self, environment): self.timer = environment.get("timer") self.metadata = environment["agent_metadata"] self.agent_config_merger = environment["agent_config_merger"] - self.is_lambda_one_click_pg_created_during_execution = False + self.is_profiling_group_created_during_execution = False def _encode_profile(self, profile): output_profile_stream = io.BytesIO() @@ -71,12 +72,15 @@ def refresh_configuration(self): # see https://boto3.amazonaws.com/v1/documentation/api/latest/guide/error-handling.html if error.response['Error']['Code'] == 'ValidationException': self.agent_config_merger.disable_profiling() + self._log_request_failed(operation="configure_agent", exception=error) + if error.response['Error']['Code'] == 'ResourceNotFoundException': - if self.is_lambda_one_click_integration_active(): - self.create_pg_when_one_click_integration_is_active() + if self.should_autocreate_profiling_group(): + logger.info("ResourceNotFoundException for a Lambda PG in ConfigureAgent call. "+ + "Attempting to create PG") + self.create_profiling_group() else: self.agent_config_merger.disable_profiling() - self._log_request_failed(operation="configure_agent", exception=error) except Exception as e: self._log_request_failed(operation="configure_agent", exception=e) @@ -88,6 +92,7 @@ def report(self, profile): :param profile: Profile to be encoded and reported to the profiler backend service. :return: True if profile gets reported successfully; False otherwise. """ + global is_create_pg_called_during_submit_profile try: profile_stream = self._encode_profile(profile) self.codeguru_client_builder.codeguru_client.post_agent_profile( @@ -99,52 +104,53 @@ def report(self, profile): return True except ClientError as error: if error.response['Error']['Code'] == 'ResourceNotFoundException': - if self.is_lambda_one_click_integration_active(): - global is_create_pg_called_during_submit_profile + if self.should_autocreate_profiling_group(): is_create_pg_called_during_submit_profile = True - logger.info("Received ResourceNotFoundException. Attempting to create 1-click PG in Report Profile") - self.create_pg_when_one_click_integration_is_active() + logger.info("ResourceNotFoundException for a Lambda PG in PostAgentProfile call. " + + "Attempting to create PG") + self.create_profiling_group() return False except Exception as e: self._log_request_failed(operation="post_agent_profile", exception=e) return False - @with_timer("createOneClickPG", measurement="wall-clock-time") - def create_pg_when_one_click_integration_is_active(self): + @with_timer("createProfilingGroup", measurement="wall-clock-time") + def create_profiling_group(self): """ Create a PG for the Lambda function onboarded with 1-click integration """ - function_name = str(self.metadata.fleet_info.function_arn).split(':')[6] - - profiling_group_name_for_1_click = "aws-lambda-{}".format(function_name) - profiling_group_name_to_be_created = os.getenv("AWS_CODEGURU_PROFILER_GROUP_NAME", - profiling_group_name_for_1_click) try: self.codeguru_client_builder.codeguru_client.create_profiling_group( - profilingGroupName=profiling_group_name_to_be_created, + profilingGroupName=self.profiling_group_name, computePlatform='AWSLambda' ) - self.profiling_group_name = profiling_group_name_to_be_created - self.is_lambda_one_click_pg_created_during_execution = True - logger.info("Created Lambda Profiling Group with name {}", profiling_group_name_to_be_created) - + self.is_profiling_group_created_during_execution = True + logger.info("Created Lambda Profiling Group with name " + str(self.profiling_group_name)) except Exception as e: self._log_request_failed(operation="create_profiling_group", exception=e) - def is_lambda_one_click_integration_active(self): + def should_autocreate_profiling_group(self): + return self.is_runtime_python() and self.is_handler_name_env_variable_set() + + def is_handler_name_env_variable_set(self): """ Check if the ComputeType is AWSLambda and if the environment variables for Lambda Layer Profiling are set """ - fleet_info_as_map = self.metadata.fleet_info.serialize_to_map() - if 'computeType' in fleet_info_as_map: - if fleet_info_as_map['computeType'] == 'aws_lambda': - handler_name_ = os.environ.get(HANDLER_ENV_NAME_FOR_CODEGURU) - if not handler_name_: - logger.error("Env Variables for CodeGuru Profiler Lambda Layer are not set. Cannot create PG.") - return False - return True - return False + handler_name_ = os.environ.get(HANDLER_ENV_NAME_FOR_CODEGURU_KEY) + if not handler_name_: + logger.info("Env Variables for CodeGuru Profiler Lambda Layer are not set. Cannot create Profiling Group.") + return False + return True + + def is_runtime_python(self): + """ + Check if the runtime for the AWS Lambda function is Python + """ + # https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtimes.html + python_runtime_identifier_prefix = "python" + execution_environment_env_variable = os.environ.get(AWS_EXECUTION_ENV_KEY) + return execution_environment_env_variable.find(python_runtime_identifier_prefix) != -1 @staticmethod def _log_request_failed(operation, exception): diff --git a/test/acceptance/test_end_to_end_profile_and_save_to_file.py b/test/acceptance/test_end_to_end_profile_and_save_to_file.py index 004b74a..dde7586 100644 --- a/test/acceptance/test_end_to_end_profile_and_save_to_file.py +++ b/test/acceptance/test_end_to_end_profile_and_save_to_file.py @@ -5,7 +5,7 @@ import os from datetime import timedelta -from mock import patch +from unittest.mock import patch from pathlib import Path from codeguru_profiler_agent.profiler import Profiler diff --git a/test/acceptance/test_end_to_end_profiling.py b/test/acceptance/test_end_to_end_profiling.py index c69edfb..9099d8f 100644 --- a/test/acceptance/test_end_to_end_profiling.py +++ b/test/acceptance/test_end_to_end_profiling.py @@ -1,6 +1,6 @@ from botocore.stub import Stubber, ANY from datetime import timedelta -from mock import patch +from unittest.mock import patch from test.pytestutils import before from codeguru_profiler_agent.profiler import Profiler diff --git a/test/acceptance/test_live_profiling.py b/test/acceptance/test_live_profiling.py index fe09eb8..bcd41e6 100644 --- a/test/acceptance/test_live_profiling.py +++ b/test/acceptance/test_live_profiling.py @@ -1,7 +1,7 @@ import time from datetime import timedelta -from mock import patch +from unittest.mock import patch from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration from codeguru_profiler_agent.sdk_reporter.sdk_reporter import SdkReporter diff --git a/test/unit/agent_metadata/test_aws_lambda.py b/test/unit/agent_metadata/test_aws_lambda.py index 99265f7..5266892 100644 --- a/test/unit/agent_metadata/test_aws_lambda.py +++ b/test/unit/agent_metadata/test_aws_lambda.py @@ -1,6 +1,6 @@ import pytest from test.pytestutils import before -from mock import Mock +from unittest.mock import Mock from datetime import timedelta from codeguru_profiler_agent.agent_metadata.aws_lambda import AWSLambda from codeguru_profiler_agent.aws_lambda.lambda_context import LambdaContext diff --git a/test/unit/aws_lambda/test_profiler_decorator.py b/test/unit/aws_lambda/test_profiler_decorator.py index 32c47a5..933e871 100644 --- a/test/unit/aws_lambda/test_profiler_decorator.py +++ b/test/unit/aws_lambda/test_profiler_decorator.py @@ -1,7 +1,7 @@ import pytest import codeguru_profiler_agent.aws_lambda.profiler_decorator -from mock import MagicMock, patch +from unittest.mock import MagicMock from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration from codeguru_profiler_agent import with_lambda_profiler from codeguru_profiler_agent import Profiler @@ -71,7 +71,7 @@ def around(self): self.context = MagicMock() self.context.invoked_function_arn = "the_lambda_function_arn" self.env = {"AWS_LAMBDA_FUNCTION_MEMORY_SIZE": "1024", - "AWS_EXECUTION_ENV": "AWS_Lambda_python3.6"} + "AWS_EXECUTION_ENV_KEY": "AWS_Lambda_python3.6"} # define a handler function with the profiler decorator and parameters @with_lambda_profiler(profiling_group_name="pg_name", region_name="eu-north-1", diff --git a/test/unit/file_reporter/test_file_reporter.py b/test/unit/file_reporter/test_file_reporter.py index a1266b2..bdb49f2 100644 --- a/test/unit/file_reporter/test_file_reporter.py +++ b/test/unit/file_reporter/test_file_reporter.py @@ -2,8 +2,8 @@ import pytest import shutil -from mock import MagicMock -from mock import ANY +from unittest.mock import MagicMock +from unittest.mock import ANY from pathlib import Path from codeguru_profiler_agent.file_reporter.file_reporter import FileReporter diff --git a/test/unit/model/test_call_graph_node.py b/test/unit/model/test_call_graph_node.py index bcc18a6..a45f4a5 100644 --- a/test/unit/model/test_call_graph_node.py +++ b/test/unit/model/test_call_graph_node.py @@ -2,7 +2,7 @@ from codeguru_profiler_agent.model.frame import Frame from test.pytestutils import before -from mock import MagicMock +from unittest.mock import MagicMock from codeguru_profiler_agent.model.call_graph_node import CallGraphNode from codeguru_profiler_agent.model.memory_counter import MemoryCounter diff --git a/test/unit/model/test_profile.py b/test/unit/model/test_profile.py index 44926a7..667b038 100644 --- a/test/unit/model/test_profile.py +++ b/test/unit/model/test_profile.py @@ -1,5 +1,5 @@ import pytest -from mock import Mock +from unittest.mock import Mock from codeguru_profiler_agent.model.frame import Frame from test.pytestutils import before diff --git a/test/unit/sdk_reporter/test_sdk_profile_encoder.py b/test/unit/sdk_reporter/test_sdk_profile_encoder.py index 6c18045..f3092f2 100644 --- a/test/unit/sdk_reporter/test_sdk_profile_encoder.py +++ b/test/unit/sdk_reporter/test_sdk_profile_encoder.py @@ -2,7 +2,7 @@ import platform import pytest -from mock import MagicMock +from unittest.mock import MagicMock from codeguru_profiler_agent.agent_metadata.agent_metadata import AgentMetadata from codeguru_profiler_agent.agent_metadata.aws_ec2_instance import AWSEC2Instance diff --git a/test/unit/sdk_reporter/test_sdk_reporter.py b/test/unit/sdk_reporter/test_sdk_reporter.py index 086c399..1920eba 100644 --- a/test/unit/sdk_reporter/test_sdk_reporter.py +++ b/test/unit/sdk_reporter/test_sdk_reporter.py @@ -5,10 +5,9 @@ from datetime import timedelta, datetime -from codeguru_profiler_agent.agent_metadata.aws_lambda import AWSLambda from codeguru_profiler_agent.utils.time import current_milli_time from test.pytestutils import before -from mock import MagicMock +from unittest.mock import MagicMock from botocore.stub import Stubber, ANY from codeguru_profiler_agent.reporter.agent_configuration import AgentConfigurationMerger @@ -20,11 +19,12 @@ from codeguru_profiler_agent.codeguru_client_builder import CodeGuruClientBuilder profiling_group_name = "test-ProfilingGroup-name" -lambda_one_click_profiling_group_name = "aws-lambda-testLambdaName" -test_agent_metadata_for_lambda = AgentMetadata( - fleet_info=AWSLambda("arn:aws:lambda:us-east-1:111111111111:function:testLambdaName", "memory", "env", "agentId")) +test_lambda_profiling_group_name = "aws-lambda-testLambdaName" profile = Profile(profiling_group_name, 1.0, 0.5, current_milli_time()) +AWS_EXECUTION_ENV_KEY = "AWS_EXECUTION_ENV" +AWS_LAMBDA_FUNCTION_MEMORY_SIZE_KEY = "AWS_LAMBDA_FUNCTION_MEMORY_SIZE" +HANDLER_ENV_NAME_FOR_CODEGURU_KEY = "HANDLER_ENV_NAME_FOR_CODEGURU" class TestSdkReporter: def before(self): @@ -83,7 +83,7 @@ def test_creates_pg_if_onboarded_with_lambda_one_click_integration(self): expected_params = { 'agentProfile': ANY, 'contentType': 'application/json', - 'profilingGroupName': profiling_group_name + 'profilingGroupName': test_lambda_profiling_group_name } self.client_stubber.add_client_error('post_agent_profile', service_error_code='ResourceNotFoundException', @@ -119,13 +119,13 @@ def test_creates_pg_if_onboarded_with_lambda_one_click_integration(self): expected_params_one_click = { 'agentProfile': ANY, 'contentType': 'application/json', - 'profilingGroupName': lambda_one_click_profiling_group_name + 'profilingGroupName': test_lambda_profiling_group_name } self.client_stubber.add_response('post_agent_profile', {}, expected_params_one_click) - os.environ.__setitem__('AWS_CODEGURU_PROFILER_GROUP_NAME', lambda_one_click_profiling_group_name) - os.environ.__setitem__('HANDLER_ENV_NAME_FOR_CODEGURU', 'test-handler') - self.subject.metadata = test_agent_metadata_for_lambda + os.environ.__setitem__(HANDLER_ENV_NAME_FOR_CODEGURU_KEY, 'test-handler') + os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, 'AWS_Lambda_python3.8') + self.subject.profiling_group_name = test_lambda_profiling_group_name with self.client_stubber: assert self.subject.report(profile) is False assert self.subject.report(profile) is True @@ -215,13 +215,22 @@ def test_agent_configuration_when_configure_agent_throws_error(self): assert AgentConfiguration.get().should_profile is True assert AgentConfiguration.get().sampling_interval == timedelta(seconds=13) - def test_when_backends_sends_resource_not_found_it_stops_the_profiling(self): + def test_when_backends_sends_resource_not_found_it_stops_the_profiling_in_non_lambda_case(self): self.client_stubber.add_client_error('configure_agent', service_error_code='ResourceNotFoundException', service_message='Simulated error in configure_agent call') + os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, 'AWS_Lambda_java') with self.client_stubber: self.subject.refresh_configuration() assert AgentConfiguration.get().should_profile is False + def test_when_backends_sends_resource_not_found_it_does_not_stop_the_profiling_lambda_case(self): + self.client_stubber.add_client_error('configure_agent', service_error_code='ResourceNotFoundException', + service_message='Simulated error in configure_agent call') + os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, 'AWS_Lambda_python3.8') + with self.client_stubber: + self.subject.refresh_configuration() + assert AgentConfiguration.get().should_profile is True + def test_when_backend_sends_validation_exception_it_stops_the_profiling(self): self.client_stubber.add_client_error('configure_agent', service_error_code='ValidationException', service_message='Simulated error in configure_agent call') @@ -260,14 +269,13 @@ def test_creates_lambda_pg_if_onboarded_with_lambda_one_click_integration(self): } self.client_stubber.add_response('create_profiling_group', expected_response_create_pg) - os.environ.__setitem__('AWS_LAMBDA_FUNCTION_MEMORY_SIZE', "512") - os.environ.__setitem__('AWS_EXECUTION_ENV', "AWS_Lambda_python3.8") - os.environ.__setitem__('AWS_CODEGURU_PROFILER_GROUP_NAME', lambda_one_click_profiling_group_name) - os.environ.__setitem__('HANDLER_ENV_NAME_FOR_CODEGURU', 'test-handler') - self.subject.metadata = test_agent_metadata_for_lambda + os.environ.__setitem__(AWS_LAMBDA_FUNCTION_MEMORY_SIZE_KEY, "512") + os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, "AWS_Lambda_python3.8") + os.environ.__setitem__(HANDLER_ENV_NAME_FOR_CODEGURU_KEY, 'test-handler') + self.subject.profiling_group_name = test_lambda_profiling_group_name with self.client_stubber: self.subject.refresh_configuration() - assert self.subject.is_lambda_one_click_pg_created_during_execution is True + assert self.subject.is_profiling_group_created_during_execution is True def test_create_pg_not_invoked_in_non_lambda_case(self): self.client_stubber.add_client_error('configure_agent', @@ -277,9 +285,5 @@ def test_create_pg_not_invoked_in_non_lambda_case(self): with self.client_stubber: self.subject.refresh_configuration() - assert self.subject.is_lambda_one_click_pg_created_during_execution is False - try: - self.client_stubber.assert_no_pending_responses() - assert False - except AssertionError: - assert True + assert self.subject.is_profiling_group_created_during_execution is False + assert self.client_stubber.assert_no_pending_responses() is None \ No newline at end of file diff --git a/test/unit/test_codeguru_client_builder.py b/test/unit/test_codeguru_client_builder.py index a336850..0feba89 100644 --- a/test/unit/test_codeguru_client_builder.py +++ b/test/unit/test_codeguru_client_builder.py @@ -2,7 +2,7 @@ import boto3 import pytest -from mock import MagicMock +from unittest.mock import MagicMock from test.pytestutils import before from codeguru_profiler_agent.codeguru_client_builder import CodeGuruClientBuilder diff --git a/test/unit/test_local_aggregator.py b/test/unit/test_local_aggregator.py index d41b336..4bfbc73 100644 --- a/test/unit/test_local_aggregator.py +++ b/test/unit/test_local_aggregator.py @@ -5,7 +5,7 @@ from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration from codeguru_profiler_agent.utils.time import current_milli_time from test.pytestutils import before -from mock import MagicMock, call, patch +from unittest.mock import MagicMock, call from codeguru_profiler_agent.profiler import DEFAULT_REPORTING_INTERVAL, \ DEFAULT_MEMORY_LIMIT_BYTES, Profiler, INITIAL_MINIMUM_REPORTING_INTERVAL, DEFAULT_SAMPLING_INTERVAL diff --git a/test/unit/test_profiler.py b/test/unit/test_profiler.py index 2417362..60c06cb 100644 --- a/test/unit/test_profiler.py +++ b/test/unit/test_profiler.py @@ -1,6 +1,6 @@ import pytest from datetime import timedelta -from mock import Mock +from unittest.mock import Mock from codeguru_profiler_agent.profiler import Profiler from codeguru_profiler_agent.profiler_runner import ProfilerRunner diff --git a/test/unit/test_profiler_disabler.py b/test/unit/test_profiler_disabler.py index 4d3cacd..5f0b282 100644 --- a/test/unit/test_profiler_disabler.py +++ b/test/unit/test_profiler_disabler.py @@ -1,7 +1,7 @@ import tempfile from pathlib import Path import shutil -from mock import Mock, patch +from unittest.mock import Mock, patch import time from codeguru_profiler_agent.model.profile import Profile diff --git a/test/unit/test_profiler_runner.py b/test/unit/test_profiler_runner.py index 946b843..02fdf67 100644 --- a/test/unit/test_profiler_runner.py +++ b/test/unit/test_profiler_runner.py @@ -2,7 +2,7 @@ from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration from test.pytestutils import before from test.help_utils import wait_for -from mock import MagicMock +from unittest.mock import MagicMock from time import sleep from codeguru_profiler_agent.profiler_runner import ProfilerRunner diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index 9a8d41f..961ed89 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -1,6 +1,6 @@ from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration from test.pytestutils import before -import mock +import unittest.mock as mock from mock import create_autospec, MagicMock, ANY from codeguru_profiler_agent.sampler import Sampler diff --git a/test/unit/test_sampling_utils.py b/test/unit/test_sampling_utils.py index 015dcfc..10d5e59 100644 --- a/test/unit/test_sampling_utils.py +++ b/test/unit/test_sampling_utils.py @@ -1,5 +1,5 @@ import pytest -from mock import mock +import unittest.mock as mock import sys from test import help_utils diff --git a/test/unit/utils/test_execution_state.py b/test/unit/utils/test_execution_state.py index 886135d..ebd77e2 100644 --- a/test/unit/utils/test_execution_state.py +++ b/test/unit/utils/test_execution_state.py @@ -3,7 +3,7 @@ import datetime from test.pytestutils import before from queue import Queue, Empty -from mock import MagicMock, call +from unittest.mock import MagicMock, call from codeguru_profiler_agent.utils.execution_state import ExecutionState From ee6d377d2e1aad3298b7137ad59dea96825dd89a Mon Sep 17 00:00:00 2001 From: dangmaul-amazon Date: Tue, 30 Mar 2021 22:33:39 -0700 Subject: [PATCH 4/6] Add a note to the docstring for refresh_configuration() and report() --- codeguru_profiler_agent/sdk_reporter/sdk_reporter.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py index a80210f..6de28ff 100644 --- a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py +++ b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py @@ -54,6 +54,11 @@ def setup(self): def refresh_configuration(self): """ Refresh the agent configuration by calling the profiler backend service. + + Note: + For an agent running on AWS Lambda, if the environment variables for Profiling using + Lambda layers are set, it tries to create a Profiling Group whenever a ResourceNotFoundException + is encountered. """ try: fleet_instance_id = self.metadata.fleet_info.get_fleet_instance_id() @@ -91,6 +96,11 @@ def report(self, profile): :param profile: Profile to be encoded and reported to the profiler backend service. :return: True if profile gets reported successfully; False otherwise. + + Note: + For an agent running on AWS Lambda, if the environment variables for Profiling using + Lambda layers are set, it tries to create a Profiling Group whenever a ResourceNotFoundException + is encountered. """ global is_create_pg_called_during_submit_profile try: From 4d23afccfc80ecd8adb3b372513ae524248a7981 Mon Sep 17 00:00:00 2001 From: dangmaul-amazon Date: Wed, 31 Mar 2021 16:43:35 -0700 Subject: [PATCH 5/6] Minor fixes for 1-click integration --- .../agent_metadata/aws_lambda.py | 3 + .../aws_lambda/lambda_handler.py | 7 +- codeguru_profiler_agent/local_aggregator.py | 8 +- .../sdk_reporter/sdk_reporter.py | 53 +++++------ test/unit/sdk_reporter/test_sdk_reporter.py | 90 +++++++------------ test/unit/test_local_aggregator.py | 2 +- 6 files changed, 63 insertions(+), 100 deletions(-) diff --git a/codeguru_profiler_agent/agent_metadata/aws_lambda.py b/codeguru_profiler_agent/agent_metadata/aws_lambda.py index 33afbbf..d772e70 100644 --- a/codeguru_profiler_agent/agent_metadata/aws_lambda.py +++ b/codeguru_profiler_agent/agent_metadata/aws_lambda.py @@ -10,6 +10,9 @@ LAMBDA_MEMORY_SIZE_ENV = "AWS_LAMBDA_FUNCTION_MEMORY_SIZE" LAMBDA_EXECUTION_ENV = "AWS_EXECUTION_ENV" +HANDLER_ENV_NAME_FOR_CODEGURU_KEY = "HANDLER_ENV_NAME_FOR_CODEGURU" +LAMBDA_TASK_ROOT = "LAMBDA_TASK_ROOT" +LAMBDA_RUNTIME_DIR = "LAMBDA_RUNTIME_DIR" # Those are used for the configure agent call: # See https://docs.aws.amazon.com/codeguru/latest/profiler-api/API_ConfigureAgent.html diff --git a/codeguru_profiler_agent/aws_lambda/lambda_handler.py b/codeguru_profiler_agent/aws_lambda/lambda_handler.py index f31968c..f922465 100644 --- a/codeguru_profiler_agent/aws_lambda/lambda_handler.py +++ b/codeguru_profiler_agent/aws_lambda/lambda_handler.py @@ -1,9 +1,8 @@ import os import logging from codeguru_profiler_agent.aws_lambda.profiler_decorator import with_lambda_profiler - +from codeguru_profiler_agent.agent_metadata.aws_lambda import HANDLER_ENV_NAME_FOR_CODEGURU_KEY HANDLER_ENV_NAME = "_HANDLER" -HANDLER_ENV_NAME_FOR_CODEGURU = "HANDLER_ENV_NAME_FOR_CODEGURU" logger = logging.getLogger(__name__) @@ -11,11 +10,11 @@ def restore_handler_env(original_handler, env=os.environ): env[HANDLER_ENV_NAME] = original_handler -def load_handler(bootstrap_module, env=os.environ, original_handler_env_key=HANDLER_ENV_NAME_FOR_CODEGURU): +def load_handler(bootstrap_module, env=os.environ, original_handler_env_key=HANDLER_ENV_NAME_FOR_CODEGURU_KEY): try: original_handler_name = env.get(original_handler_env_key) if not original_handler_name: - raise ValueError("Could not find module and function name from " + HANDLER_ENV_NAME_FOR_CODEGURU + raise ValueError("Could not find module and function name from " + HANDLER_ENV_NAME_FOR_CODEGURU_KEY + " environment variable") # Delegate to the lambda code to load the customer's module. diff --git a/codeguru_profiler_agent/local_aggregator.py b/codeguru_profiler_agent/local_aggregator.py index dda3b9c..aa31ea6 100644 --- a/codeguru_profiler_agent/local_aggregator.py +++ b/codeguru_profiler_agent/local_aggregator.py @@ -101,7 +101,7 @@ def refresh_configuration(self): self.reporter.refresh_configuration() def _report_profile(self, now): - current_last_report_attempted_value = self.last_report_attempted + previous_last_report_attempted_value = self.last_report_attempted self.last_report_attempted = now self._add_overhead_metric_to_profile() logger.info("Attempting to report profile data: " + str(self.profile)) @@ -110,13 +110,13 @@ def _report_profile(self, now): return False is_reporting_successful = self.reporter.report(self.profile) ''' - If we attempt to create a PG in the report() call, we do not want to update the last_report_attempted_value + If we attempt to create a Profiling Group in the report() call, we do not want to update the last_report_attempted_value since we did not actually report a profile. - This will occur only in the case of Lambda 1-click integration. + This will occur only in the case of profiling using CodeGuru Profiler Python agent Lambda layer. ''' if SdkReporter.check_create_pg_called_during_submit_profile == True: - self.last_report_attempted = current_last_report_attempted_value + self.last_report_attempted = previous_last_report_attempted_value SdkReporter.reset_check_create_pg_called_during_submit_profile_flag() return is_reporting_successful diff --git a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py index 6de28ff..65ad1f7 100644 --- a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py +++ b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py @@ -9,9 +9,10 @@ from codeguru_profiler_agent.reporter.reporter import Reporter from codeguru_profiler_agent.metrics.with_timer import with_timer from codeguru_profiler_agent.sdk_reporter.profile_encoder import ProfileEncoder +from codeguru_profiler_agent.agent_metadata.aws_lambda import HANDLER_ENV_NAME_FOR_CODEGURU_KEY, \ + LAMBDA_TASK_ROOT, LAMBDA_RUNTIME_DIR logger = logging.getLogger(__name__) -HANDLER_ENV_NAME_FOR_CODEGURU_KEY = "HANDLER_ENV_NAME_FOR_CODEGURU" AWS_EXECUTION_ENV_KEY = "AWS_EXECUTION_ENV" class SdkReporter(Reporter): @@ -78,11 +79,12 @@ def refresh_configuration(self): if error.response['Error']['Code'] == 'ValidationException': self.agent_config_merger.disable_profiling() self._log_request_failed(operation="configure_agent", exception=error) - if error.response['Error']['Code'] == 'ResourceNotFoundException': - if self.should_autocreate_profiling_group(): - logger.info("ResourceNotFoundException for a Lambda PG in ConfigureAgent call. "+ - "Attempting to create PG") + if self.is_compute_platform_lambda(): + logger.info( + "Profiling group not found. Will try to create a profiling group" + "with name = {} and compute platform = {} and retry calling configure agent after 5 minutes" + .format(self.profiling_group_name, 'AWSLambda')) self.create_profiling_group() else: self.agent_config_merger.disable_profiling() @@ -102,7 +104,6 @@ def report(self, profile): Lambda layers are set, it tries to create a Profiling Group whenever a ResourceNotFoundException is encountered. """ - global is_create_pg_called_during_submit_profile try: profile_stream = self._encode_profile(profile) self.codeguru_client_builder.codeguru_client.post_agent_profile( @@ -114,10 +115,11 @@ def report(self, profile): return True except ClientError as error: if error.response['Error']['Code'] == 'ResourceNotFoundException': - if self.should_autocreate_profiling_group(): - is_create_pg_called_during_submit_profile = True - logger.info("ResourceNotFoundException for a Lambda PG in PostAgentProfile call. " + - "Attempting to create PG") + if self.is_compute_platform_lambda(): + self.__class__.is_create_pg_called_during_submit_profile = True + logger.info( + "Profiling group not found. Will try to create a profiling group" + "with name = {} and compute platform = {}".format(self.profiling_group_name, 'AWSLambda')) self.create_profiling_group() return False except Exception as e: @@ -127,7 +129,7 @@ def report(self, profile): @with_timer("createProfilingGroup", measurement="wall-clock-time") def create_profiling_group(self): """ - Create a PG for the Lambda function onboarded with 1-click integration + Create a Profiling Group for the AWS Lambda function. """ try: self.codeguru_client_builder.codeguru_client.create_profiling_group( @@ -136,31 +138,20 @@ def create_profiling_group(self): ) self.is_profiling_group_created_during_execution = True logger.info("Created Lambda Profiling Group with name " + str(self.profiling_group_name)) + except ClientError as error: + if error.response['Error']['Code'] == 'ConflictException': + logger.info("Profiling Group with name {} already exists. Please use a different name." + .format(self.profiling_group_name)) except Exception as e: self._log_request_failed(operation="create_profiling_group", exception=e) - def should_autocreate_profiling_group(self): - return self.is_runtime_python() and self.is_handler_name_env_variable_set() - - def is_handler_name_env_variable_set(self): - """ - Check if the ComputeType is AWSLambda and if the environment - variables for Lambda Layer Profiling are set - """ - handler_name_ = os.environ.get(HANDLER_ENV_NAME_FOR_CODEGURU_KEY) - if not handler_name_: - logger.info("Env Variables for CodeGuru Profiler Lambda Layer are not set. Cannot create Profiling Group.") - return False - return True - - def is_runtime_python(self): + def is_compute_platform_lambda(self): """ - Check if the runtime for the AWS Lambda function is Python + Check if the compute platform is AWS Lambda. """ - # https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtimes.html - python_runtime_identifier_prefix = "python" - execution_environment_env_variable = os.environ.get(AWS_EXECUTION_ENV_KEY) - return execution_environment_env_variable.find(python_runtime_identifier_prefix) != -1 + does_lambda_task_root_exist = os.environ.get(LAMBDA_TASK_ROOT) + does_lambda_runtime_dir_exist = os.environ.get(LAMBDA_RUNTIME_DIR) + return bool(does_lambda_task_root_exist) and bool(does_lambda_runtime_dir_exist) @staticmethod def _log_request_failed(operation, exception): diff --git a/test/unit/sdk_reporter/test_sdk_reporter.py b/test/unit/sdk_reporter/test_sdk_reporter.py index 1920eba..b0abbfc 100644 --- a/test/unit/sdk_reporter/test_sdk_reporter.py +++ b/test/unit/sdk_reporter/test_sdk_reporter.py @@ -12,6 +12,8 @@ from codeguru_profiler_agent.reporter.agent_configuration import AgentConfigurationMerger from codeguru_profiler_agent.agent_metadata.agent_metadata import AgentMetadata, DefaultFleetInfo +from codeguru_profiler_agent.agent_metadata.aws_lambda import LAMBDA_EXECUTION_ENV, \ + LAMBDA_MEMORY_SIZE_ENV, LAMBDA_TASK_ROOT, LAMBDA_RUNTIME_DIR from codeguru_profiler_agent.reporter.agent_configuration import AgentConfiguration from codeguru_profiler_agent.sdk_reporter.sdk_reporter import SdkReporter from codeguru_profiler_agent.sdk_reporter.profile_encoder import ProfileEncoder @@ -19,13 +21,9 @@ from codeguru_profiler_agent.codeguru_client_builder import CodeGuruClientBuilder profiling_group_name = "test-ProfilingGroup-name" -test_lambda_profiling_group_name = "aws-lambda-testLambdaName" +autocreated_test_lambda_profiling_group_name = "aws-lambda-testLambdaName" profile = Profile(profiling_group_name, 1.0, 0.5, current_milli_time()) -AWS_EXECUTION_ENV_KEY = "AWS_EXECUTION_ENV" -AWS_LAMBDA_FUNCTION_MEMORY_SIZE_KEY = "AWS_LAMBDA_FUNCTION_MEMORY_SIZE" -HANDLER_ENV_NAME_FOR_CODEGURU_KEY = "HANDLER_ENV_NAME_FOR_CODEGURU" - class TestSdkReporter: def before(self): codeguru_client_builder = CodeGuruClientBuilder(environment={ @@ -34,6 +32,8 @@ def before(self): self.client_stubber = Stubber(codeguru_client_builder.codeguru_client) + self.clear_lambda_specific_environment_variables_for_test_run() + profile_encoder = MagicMock(name="profile_encoder", spec=ProfileEncoder) profile_encoder.encode.side_effect = lambda **args: args["output_stream"].write( b"test-profile-encoder-output") @@ -56,6 +56,11 @@ def before(self): self.subject = SdkReporter(environment=self.environment) self.subject.setup() + def clear_lambda_specific_environment_variables_for_test_run(self): + keys_to_delete = [LAMBDA_TASK_ROOT, LAMBDA_RUNTIME_DIR, LAMBDA_EXECUTION_ENV, LAMBDA_MEMORY_SIZE_ENV] + for key in keys_to_delete: + if key in os.environ: + os.environ.__delitem__(key) class TestReport(TestSdkReporter): @before @@ -79,11 +84,11 @@ def test_return_false_when_report_throws_error(self): with self.client_stubber: assert self.subject.report(profile) is False - def test_creates_pg_if_onboarded_with_lambda_one_click_integration(self): + def test_create_profiling_group_called_when_pg_does_not_exist_lambda_case(self): expected_params = { 'agentProfile': ANY, 'contentType': 'application/json', - 'profilingGroupName': test_lambda_profiling_group_name + 'profilingGroupName': autocreated_test_lambda_profiling_group_name } self.client_stubber.add_client_error('post_agent_profile', service_error_code='ResourceNotFoundException', @@ -116,16 +121,17 @@ def test_creates_pg_if_onboarded_with_lambda_one_click_integration(self): } self.client_stubber.add_response('create_profiling_group', expected_response_create_pg) - expected_params_one_click = { + expected_params_post_agent_profile = { 'agentProfile': ANY, 'contentType': 'application/json', - 'profilingGroupName': test_lambda_profiling_group_name + 'profilingGroupName': autocreated_test_lambda_profiling_group_name } - self.client_stubber.add_response('post_agent_profile', {}, expected_params_one_click) + self.client_stubber.add_response('post_agent_profile', {}, expected_params_post_agent_profile) + + os.environ.__setitem__(LAMBDA_TASK_ROOT, 'test-task-root') + os.environ.__setitem__(LAMBDA_RUNTIME_DIR, 'test-dir') + self.subject.profiling_group_name = autocreated_test_lambda_profiling_group_name - os.environ.__setitem__(HANDLER_ENV_NAME_FOR_CODEGURU_KEY, 'test-handler') - os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, 'AWS_Lambda_python3.8') - self.subject.profiling_group_name = test_lambda_profiling_group_name with self.client_stubber: assert self.subject.report(profile) is False assert self.subject.report(profile) is True @@ -142,46 +148,9 @@ def test_create_pg_not_invoked_in_non_lambda_case(self): 'post_agent_profile call', expected_params=expected_params) - expected_params_post_agent_profile = { - 'agentProfile': ANY, - 'contentType': 'application/json', - 'profilingGroupName': profiling_group_name - } - self.client_stubber.add_response('post_agent_profile', {}, expected_params_post_agent_profile) - - expected_response_create_pg = { - 'profilingGroup': { - 'agentOrchestrationConfig': { - 'profilingEnabled': True - }, - 'arn': 'string', - 'computePlatform': 'AWSLambda', - 'createdAt': datetime(2015, 1, 1), - 'name': 'string', - 'profilingStatus': { - 'latestAgentOrchestratedAt': datetime(2015, 1, 1), - 'latestAgentProfileReportedAt': datetime(2015, 1, 1), - 'latestAggregatedProfile': { - 'period': 'PT5M', - 'start': datetime(2015, 1, 1) - } - }, - 'tags': { - 'string': 'string' - }, - 'updatedAt': datetime(2015, 1, 1) - } - } - self.client_stubber.add_response('create_profiling_group', expected_response_create_pg) - with self.client_stubber: assert self.subject.report(profile) is False - assert self.subject.report(profile) is True - try: - self.client_stubber.assert_no_pending_responses() - assert False - except AssertionError: - assert True + self.client_stubber.assert_no_pending_responses() class TestConfigureAgent(TestSdkReporter): @@ -218,15 +187,15 @@ def test_agent_configuration_when_configure_agent_throws_error(self): def test_when_backends_sends_resource_not_found_it_stops_the_profiling_in_non_lambda_case(self): self.client_stubber.add_client_error('configure_agent', service_error_code='ResourceNotFoundException', service_message='Simulated error in configure_agent call') - os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, 'AWS_Lambda_java') with self.client_stubber: self.subject.refresh_configuration() assert AgentConfiguration.get().should_profile is False - def test_when_backends_sends_resource_not_found_it_does_not_stop_the_profiling_lambda_case(self): + def test_when_backends_sends_resource_not_found_it_does_not_stop_the_profiling_in_lambda_case(self): self.client_stubber.add_client_error('configure_agent', service_error_code='ResourceNotFoundException', service_message='Simulated error in configure_agent call') - os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, 'AWS_Lambda_python3.8') + os.environ.__setitem__(LAMBDA_TASK_ROOT, 'test-task-root') + os.environ.__setitem__(LAMBDA_RUNTIME_DIR, 'test-dir') with self.client_stubber: self.subject.refresh_configuration() assert AgentConfiguration.get().should_profile is True @@ -238,7 +207,7 @@ def test_when_backend_sends_validation_exception_it_stops_the_profiling(self): self.subject.refresh_configuration() assert AgentConfiguration.get().should_profile is False - def test_creates_lambda_pg_if_onboarded_with_lambda_one_click_integration(self): + def test_create_profiling_group_called_when_pg_does_not_exist_in_lambda_case(self): self.client_stubber.add_client_error('configure_agent', service_error_code='ResourceNotFoundException', service_message='Simulated ResourceNotFoundException in ' @@ -269,10 +238,11 @@ def test_creates_lambda_pg_if_onboarded_with_lambda_one_click_integration(self): } self.client_stubber.add_response('create_profiling_group', expected_response_create_pg) - os.environ.__setitem__(AWS_LAMBDA_FUNCTION_MEMORY_SIZE_KEY, "512") - os.environ.__setitem__(AWS_EXECUTION_ENV_KEY, "AWS_Lambda_python3.8") - os.environ.__setitem__(HANDLER_ENV_NAME_FOR_CODEGURU_KEY, 'test-handler') - self.subject.profiling_group_name = test_lambda_profiling_group_name + os.environ.__setitem__(LAMBDA_MEMORY_SIZE_ENV, "512") + os.environ.__setitem__(LAMBDA_EXECUTION_ENV, "AWS_Lambda_python3.8") + os.environ.__setitem__(LAMBDA_TASK_ROOT, 'test-task-root') + os.environ.__setitem__(LAMBDA_RUNTIME_DIR, 'test-dir') + self.subject.profiling_group_name = autocreated_test_lambda_profiling_group_name with self.client_stubber: self.subject.refresh_configuration() assert self.subject.is_profiling_group_created_during_execution is True @@ -286,4 +256,4 @@ def test_create_pg_not_invoked_in_non_lambda_case(self): with self.client_stubber: self.subject.refresh_configuration() assert self.subject.is_profiling_group_created_during_execution is False - assert self.client_stubber.assert_no_pending_responses() is None \ No newline at end of file + self.client_stubber.assert_no_pending_responses() diff --git a/test/unit/test_local_aggregator.py b/test/unit/test_local_aggregator.py index 4bfbc73..bbe96d6 100644 --- a/test/unit/test_local_aggregator.py +++ b/test/unit/test_local_aggregator.py @@ -42,7 +42,7 @@ class TestLocalAggregator: def before(self): self.mock_reporter = MagicMock(name="reporter", spec=SdkReporter) - self.mock_reporter.lambda_one_click_pg_created_during_execution = False + self.mock_reporter.is_profiling_group_created_during_execution = False self.mock_profile = MagicMock(name="profile", spec=Profile) self.mock_profile_factory = MagicMock( name="profile_factory", From b444e4fa44bba3528751a156d80102a18974eafb Mon Sep 17 00:00:00 2001 From: dangmaul-amazon Date: Thu, 1 Apr 2021 12:06:24 -0700 Subject: [PATCH 6/6] Removing unused variable and adding a wrapper to check if profile should be autocreated --- .../sdk_reporter/sdk_reporter.py | 17 +++++++++++------ test/unit/sdk_reporter/test_sdk_reporter.py | 3 +-- test/unit/test_local_aggregator.py | 1 - 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py index 65ad1f7..a3fb6bf 100644 --- a/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py +++ b/codeguru_profiler_agent/sdk_reporter/sdk_reporter.py @@ -35,7 +35,6 @@ def __init__(self, environment): self.timer = environment.get("timer") self.metadata = environment["agent_metadata"] self.agent_config_merger = environment["agent_config_merger"] - self.is_profiling_group_created_during_execution = False def _encode_profile(self, profile): output_profile_stream = io.BytesIO() @@ -80,9 +79,9 @@ def refresh_configuration(self): self.agent_config_merger.disable_profiling() self._log_request_failed(operation="configure_agent", exception=error) if error.response['Error']['Code'] == 'ResourceNotFoundException': - if self.is_compute_platform_lambda(): + if self.should_auto_create_profiling_group(): logger.info( - "Profiling group not found. Will try to create a profiling group" + "Profiling group not found. Will try to create a profiling group " "with name = {} and compute platform = {} and retry calling configure agent after 5 minutes" .format(self.profiling_group_name, 'AWSLambda')) self.create_profiling_group() @@ -115,10 +114,10 @@ def report(self, profile): return True except ClientError as error: if error.response['Error']['Code'] == 'ResourceNotFoundException': - if self.is_compute_platform_lambda(): + if self.should_auto_create_profiling_group(): self.__class__.is_create_pg_called_during_submit_profile = True logger.info( - "Profiling group not found. Will try to create a profiling group" + "Profiling group not found. Will try to create a profiling group " "with name = {} and compute platform = {}".format(self.profiling_group_name, 'AWSLambda')) self.create_profiling_group() return False @@ -136,7 +135,6 @@ def create_profiling_group(self): profilingGroupName=self.profiling_group_name, computePlatform='AWSLambda' ) - self.is_profiling_group_created_during_execution = True logger.info("Created Lambda Profiling Group with name " + str(self.profiling_group_name)) except ClientError as error: if error.response['Error']['Code'] == 'ConflictException': @@ -145,6 +143,13 @@ def create_profiling_group(self): except Exception as e: self._log_request_failed(operation="create_profiling_group", exception=e) + def should_auto_create_profiling_group(self): + """ + Currently the only condition we check is to verify that the Compute Platform is AWS Lambda. + In future, other checks could be places inside this method. + """ + return self.is_compute_platform_lambda() + def is_compute_platform_lambda(self): """ Check if the compute platform is AWS Lambda. diff --git a/test/unit/sdk_reporter/test_sdk_reporter.py b/test/unit/sdk_reporter/test_sdk_reporter.py index b0abbfc..287894d 100644 --- a/test/unit/sdk_reporter/test_sdk_reporter.py +++ b/test/unit/sdk_reporter/test_sdk_reporter.py @@ -245,7 +245,7 @@ def test_create_profiling_group_called_when_pg_does_not_exist_in_lambda_case(sel self.subject.profiling_group_name = autocreated_test_lambda_profiling_group_name with self.client_stubber: self.subject.refresh_configuration() - assert self.subject.is_profiling_group_created_during_execution is True + self.client_stubber.assert_no_pending_responses() def test_create_pg_not_invoked_in_non_lambda_case(self): self.client_stubber.add_client_error('configure_agent', @@ -255,5 +255,4 @@ def test_create_pg_not_invoked_in_non_lambda_case(self): with self.client_stubber: self.subject.refresh_configuration() - assert self.subject.is_profiling_group_created_during_execution is False self.client_stubber.assert_no_pending_responses() diff --git a/test/unit/test_local_aggregator.py b/test/unit/test_local_aggregator.py index bbe96d6..604a3ae 100644 --- a/test/unit/test_local_aggregator.py +++ b/test/unit/test_local_aggregator.py @@ -42,7 +42,6 @@ class TestLocalAggregator: def before(self): self.mock_reporter = MagicMock(name="reporter", spec=SdkReporter) - self.mock_reporter.is_profiling_group_created_during_execution = False self.mock_profile = MagicMock(name="profile", spec=Profile) self.mock_profile_factory = MagicMock( name="profile_factory",