diff --git a/sagemaker-core/docs/index.rst b/sagemaker-core/docs/index.rst index ac905bd49d..84fbf42dcc 100644 --- a/sagemaker-core/docs/index.rst +++ b/sagemaker-core/docs/index.rst @@ -14,5 +14,14 @@ SageMaker Core Shapes ######################## .. automodule:: sagemaker.core.shapes + :members: + :noindex: + + +############################## +SageMaker Core Lambda Helper +############################## + +.. automodule:: sagemaker.core.lambda_helper :members: :noindex: \ No newline at end of file diff --git a/sagemaker-core/src/sagemaker/core/lambda_helper.py b/sagemaker-core/src/sagemaker/core/lambda_helper.py index 7c1a4c26e7..51c9ae630c 100644 --- a/sagemaker-core/src/sagemaker/core/lambda_helper.py +++ b/sagemaker-core/src/sagemaker/core/lambda_helper.py @@ -171,7 +171,7 @@ def update(self): # get function name to be used in S3 upload path if self.function_arn: - versioned_function_name = self.function_arn.split("funtion:")[-1] + versioned_function_name = self.function_arn.split("function:")[-1] if ":" in versioned_function_name: function_name_for_s3 = versioned_function_name.split(":")[0] else: @@ -202,6 +202,52 @@ def update(self): else: raise ValueError(error) + def update_configuration(self): + """Method to update a Lambda function's configuration. + + Updates configuration properties such as timeout, memory_size, runtime, + handler, execution_role_arn, vpc_config, environment, and layers. + + Returns: boto3 response from Lambda's update_function_configuration method. + """ + lambda_client = _get_lambda_client(self.session) + function_identifier = self.function_name or self.function_arn + + kwargs = {"FunctionName": function_identifier} + + if self.handler is not None: + kwargs["Handler"] = self.handler + if self.runtime is not None: + kwargs["Runtime"] = self.runtime + if self.execution_role_arn is not None: + kwargs["Role"] = self.execution_role_arn + if self.timeout is not None: + kwargs["Timeout"] = self.timeout + if self.memory_size is not None: + kwargs["MemorySize"] = self.memory_size + if self.vpc_config: + kwargs["VpcConfig"] = self.vpc_config + if self.environment: + kwargs["Environment"] = self.environment + if self.layers: + kwargs["Layers"] = self.layers + + retry_attempts = 7 + for i in range(retry_attempts): + try: + response = lambda_client.update_function_configuration(**kwargs) + return response + except ClientError as e: + error = e.response["Error"] + code = error["Code"] + if code == "ResourceConflictException": + if i == retry_attempts - 1: + raise ValueError(error) + time.sleep(2**i) + else: + raise ValueError(error) + + def upsert(self): """Method to create a lambda function or update it if it already exists diff --git a/sagemaker-core/tests/unit/test_lambda_helper.py b/sagemaker-core/tests/unit/test_lambda_helper.py index cc0a52be1e..d917749392 100644 --- a/sagemaker-core/tests/unit/test_lambda_helper.py +++ b/sagemaker-core/tests/unit/test_lambda_helper.py @@ -336,6 +336,139 @@ def test_upsert_updates_existing_function(self, mock_update, mock_create): mock_update.assert_called_once() +class TestLambdaUpdateConfiguration: + """Test Lambda.update_configuration method.""" + + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_update_configuration_success(self, mock_get_client): + """Test updating Lambda function configuration.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + mock_client.update_function_configuration.return_value = { + "FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:my-function", + "Timeout": 300, + "MemorySize": 256, + } + + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + script="/path/to/script.py", + handler="script.handler", + timeout=300, + memory_size=256, + runtime="python3.12", + environment={"Variables": {"KEY": "value"}}, + layers=["arn:aws:lambda:us-west-2:123456789012:layer:my-layer:1"], + ) + result = lambda_obj.update_configuration() + + assert result["Timeout"] == 300 + assert result["MemorySize"] == 256 + mock_client.update_function_configuration.assert_called_once_with( + FunctionName="my-function", + Handler="script.handler", + Runtime="python3.12", + Role="arn:aws:iam::123456789012:role/my-role", + Timeout=300, + MemorySize=256, + Environment={"Variables": {"KEY": "value"}}, + Layers=["arn:aws:lambda:us-west-2:123456789012:layer:my-layer:1"], + ) + + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_update_configuration_with_function_arn(self, mock_get_client): + """Test updating configuration using function ARN.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + mock_client.update_function_configuration.return_value = {"Timeout": 60} + + lambda_obj = Lambda( + function_arn="arn:aws:lambda:us-west-2:123456789012:function:my-function", + timeout=60, + ) + result = lambda_obj.update_configuration() + + assert result["Timeout"] == 60 + call_kwargs = mock_client.update_function_configuration.call_args[1] + assert call_kwargs["FunctionName"] == "arn:aws:lambda:us-west-2:123456789012:function:my-function" + assert call_kwargs["Timeout"] == 60 + + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_update_configuration_retry_on_resource_conflict(self, mock_get_client): + """Test update_configuration retries on ResourceConflictException.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + error = ClientError( + {"Error": {"Code": "ResourceConflictException", "Message": "Resource in use"}}, + "UpdateFunctionConfiguration", + ) + mock_client.update_function_configuration.side_effect = [ + error, + {"Timeout": 300}, + ] + + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + script="/path/to/script.py", + handler="script.handler", + timeout=300, + ) + + with patch("time.sleep"): + result = lambda_obj.update_configuration() + + assert result["Timeout"] == 300 + assert mock_client.update_function_configuration.call_count == 2 + + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_update_configuration_max_retries_exceeded(self, mock_get_client): + """Test update_configuration fails after max retries.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + error = ClientError( + {"Error": {"Code": "ResourceConflictException", "Message": "Resource in use"}}, + "UpdateFunctionConfiguration", + ) + mock_client.update_function_configuration.side_effect = error + + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + script="/path/to/script.py", + handler="script.handler", + ) + + with patch("time.sleep"): + with pytest.raises(ValueError): + lambda_obj.update_configuration() + + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_update_configuration_non_retryable_error(self, mock_get_client): + """Test update_configuration raises on non-retryable errors.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + error = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": "Function not found"}}, + "UpdateFunctionConfiguration", + ) + mock_client.update_function_configuration.side_effect = error + + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + script="/path/to/script.py", + handler="script.handler", + ) + + with pytest.raises(ValueError): + lambda_obj.update_configuration() + + class TestLambdaInvoke: """Test Lambda.invoke method."""