Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion sagemaker-train/src/sagemaker/ai_registry/evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ def _create_lambda_function(cls, name: str, source_file: str, role: Optional[str
# Create Lambda function
lambda_client = boto3.client("lambda")
function_name = f"SageMaker-evaluator-{name}-{datetime.now().strftime('%Y%m%d_%H%M%S')}"
handler_name = f"{os.path.splitext(os.path.basename(source_file))[0]}.lambda_handler"
handler_name = "lambda_function.lambda_handler"
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fix itself makes sense if the zip archive always packages the source file as lambda_function.py regardless of the original filename. However, this introduces a subtle coupling — the hardcoded string "lambda_function.lambda_handler" must match whatever entry name is used when building the zip. Consider extracting this as a module-level constant (e.g., LAMBDA_HANDLER_NAME = "lambda_function.lambda_handler") and adding a comment explaining why it's hardcoded (i.e., the zip archive always renames the source to lambda_function.py). This will prevent future contributors from "fixing" it back to the dynamic version.

Also, as a pre-existing issue: this method calls boto3.client("lambda") directly (line 383). Per V3 architecture tenets, all AWS API interactions should go through sagemaker-core rather than calling boto3 directly. This is out of scope for this PR, but worth tracking.


try:
lambda_response = lambda_client.create_function(
Expand Down
16 changes: 16 additions & 0 deletions sagemaker-train/tests/integ/ai_registry/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,22 @@ def sample_jsonl_file():
os.unlink(f.name)


@pytest.fixture
def sample_lambda_py_file():
"""Create a raw Python Lambda file with a non-default filename to test handler derivation."""
code = '''import json
def lambda_handler(event, context):
return {"statusCode": 200, "body": json.dumps({"score": 0.9})}
'''
with tempfile.NamedTemporaryFile(mode='w', suffix='.py', prefix='my_custom_evaluator_', delete=False) as f:
f.write(code)
f.flush()
os.fsync(f.fileno())
fname = f.name
yield fname
os.unlink(fname)


@pytest.fixture
def sample_lambda_code():
"""Create sample Lambda function code as zip."""
Expand Down
41 changes: 41 additions & 0 deletions sagemaker-train/tests/integ/ai_registry/test_evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,47 @@ def test_create_reward_function_from_local_code(self, unique_name, sample_lambda
assert evaluator.method == EvaluatorMethod.BYOC
assert evaluator.reference is not None

def test_create_reward_function_from_local_py_file_and_invoke(
self, unique_name, sample_lambda_py_file, test_role, cleanup_list
):
"""End-to-end test: create evaluator from a raw .py file with non-default name and invoke it.

Regression test for the handler name bug where the Lambda was created with an incorrect
handler derived from the source filename instead of 'lambda_function.lambda_handler'.
"""
import json
import boto3

evaluator = Evaluator.create(
name=unique_name,
type=REWARD_FUNCTION,
source=sample_lambda_py_file,
role=test_role,
wait=True, # wait for Lambda to be active
)
cleanup_list.append(evaluator)
assert evaluator.method == EvaluatorMethod.BYOC
assert evaluator.reference is not None

# Wait for Lambda to become Active before invoking
lambda_client = boto3.client("lambda")
waiter = lambda_client.get_waiter("function_active_v2")
waiter.wait(FunctionName=evaluator.reference)

# Invoke the Lambda directly to verify the handler is correct
lambda_client = boto3.client("lambda")
response = lambda_client.invoke(
FunctionName=evaluator.reference,
InvocationType="RequestResponse",
Payload=json.dumps({"input": "test"}).encode(),
)
assert response["StatusCode"] == 200
assert "FunctionError" not in response, (
f"Lambda invocation failed with error: {response.get('FunctionError')}"
)
result = json.loads(response["Payload"].read())
assert result.get("statusCode") == 200

def test_get_evaluator(self, unique_name, sample_prompt_file, cleanup_list):
"""Test retrieving evaluator by name."""
try:
Expand Down
3 changes: 3 additions & 0 deletions sagemaker-train/tests/unit/ai_registry/test_evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,9 @@ def test_create_with_byoc(self, mock_air_hub, mock_boto3):

assert evaluator.method == EvaluatorMethod.BYOC
mock_air_hub.upload_to_s3.assert_called_once()
mock_lambda_client.create_function.assert_called_once()
call_kwargs = mock_lambda_client.create_function.call_args[1]
assert call_kwargs["Handler"] == "lambda_function.lambda_handler"

@patch('sagemaker.ai_registry.evaluator.AIRHub')
def test_get_all(self, mock_air_hub):
Expand Down
Loading