"""Lambda function compute implementation for Parsl Ephemeral AWS Provider.
SPDX-License-Identifier: Apache-2.0
SPDX-FileCopyrightText: 2025-2026 Scott Friedman and Project Contributors
"""
import logging
import json
import time
from typing import Dict, Any, Set
from botocore.exceptions import ClientError, NoCredentialsError
from ..exceptions import ResourceCreationError, ResourceCleanupError, JobSubmissionError
from ..config import SecurityConfig
from ..security import (
CredentialManager,
CredentialConfiguration,
SecurityEventType,
SecurityEventSeverity,
SecurityEvent,
)
from ..error_handling import RobustErrorHandler, RetryConfig
from ..utils.aws import get_or_create_iam_role, resolve_manager_session
from ..constants import (
TAG_PREFIX,
TAG_MANAGED,
TAG_WORKFLOW_ID,
TAG_JOB_ID,
DEFAULT_LAMBDA_RUNTIME,
DEFAULT_LAMBDA_HANDLER,
STATUS_PENDING,
STATUS_RUNNING,
STATUS_SUCCEEDED,
STATUS_FAILED,
)
logger = logging.getLogger(__name__)
[docs]
class LambdaManager:
"""Manager for AWS Lambda compute resources."""
[docs]
def __init__(self, provider: Any) -> None:
"""Initialize the Lambda manager.
Parameters
----------
provider : EphemeralProvider
The provider instance
"""
self.provider = provider
# Initialize error handling
self.error_handler = RobustErrorHandler(
retry_config=RetryConfig(
max_attempts=5, base_delay=2.0, exponential_backoff=True, jitter=True
)
)
logger.info("Error handler initialized for Lambda operations")
# Initialize security configuration and credential management
self._setup_security_config()
# Initialize audit logging
self.audit_logger = self.security_config.get_audit_logger()
if self.audit_logger:
self.audit_logger.log_event(
SecurityEvent(
event_type=SecurityEventType.CONFIG_CHANGE,
severity=SecurityEventSeverity.INFO,
message="LambdaManager initialized",
resource_type="lambda_manager",
workflow_id=self.provider.workflow_id,
metadata={"provider_region": self.provider.region},
)
)
logger.info("Audit logging enabled for Lambda operations")
# Initialize credential manager
credential_config = self.security_config.get_credential_configuration()
# Override credential config with provider-specific settings if provided
if hasattr(provider, "aws_access_key_id") or hasattr(provider, "aws_profile"):
# Legacy credential handling - create credential config from provider settings
credential_config = self._create_credential_config_from_provider()
try:
self.credential_manager = CredentialManager(credential_config)
logger.info("Credential manager initialized successfully")
# Log successful credential initialization
if self.audit_logger:
self.audit_logger.log_credential_access(
access_type="credential_init",
identity=credential_config.role_arn or "default",
success=True,
workflow_id=self.provider.workflow_id,
)
except Exception as e:
logger.error(f"Failed to initialize credential manager: {e}")
# Log failed credential initialization
if self.audit_logger:
self.audit_logger.log_credential_access(
access_type="credential_init",
identity="unknown",
success=False,
error=str(e),
workflow_id=self.provider.workflow_id,
)
raise ResourceCreationError(f"Credential initialization failed: {e}")
# Resolve the AWS session. The caller's own session takes precedence;
# the credential manager is only a fallback for a provider that has
# none. Going straight to the credential manager discarded an
# explicitly configured session -- role credentials, a chosen profile,
# a LocalStack endpoint -- in favour of ambient environment
# credentials, so operations could land in a different account than the
# caller selected (#117).
try:
self.aws_session = resolve_manager_session(
self.provider, self.credential_manager
)
except NoCredentialsError as e:
logger.error(f"No valid AWS credentials found: {e}")
raise ResourceCreationError(f"AWS credential error: {e}")
# The legacy fallback that used to sit here was dead code that still ran:
# it built session_kwargs from four unguarded provider attributes before
# testing `if not self.aws_session`, which can never be true because
# session resolution either returns a session or raises. So the only
# thing it accomplished was raising AttributeError for any provider
# lacking aws_access_key_id/aws_secret_access_key/aws_session_token/
# aws_profile -- which EphemeralProvider does not define. Credential
# resolution now belongs entirely to resolve_manager_session() and the
# credential manager (#117).
# Initialize clients
self.lambda_client = self.aws_session.client("lambda")
self.iam_client = self.aws_session.client("iam")
# Track resources for cleanup
self.function_names: Set[str] = set()
self.role_names: Set[str] = set()
self.jobs: Dict[str, Any] = {}
def _setup_security_config(self) -> None:
"""Set up security configuration from provider settings."""
# Get security settings from provider if available
security_env = getattr(self.provider, "security_environment", "dev")
vpc_cidr = getattr(self.provider, "vpc_cidr", None)
admin_cidrs = getattr(self.provider, "admin_cidr_blocks", None)
# Use provider's security config if available, otherwise create default
if hasattr(self.provider, "security_config") and self.provider.security_config:
self.security_config = self.provider.security_config
else:
if security_env.lower() == "production":
self.security_config = SecurityConfig.create_production_config(
vpc_cidr=vpc_cidr or "10.0.0.0/16",
admin_cidrs=admin_cidrs or ["10.0.0.0/8"],
)
else:
self.security_config = SecurityConfig.create_development_config(
vpc_cidr=vpc_cidr or "10.0.0.0/16"
)
def _create_credential_config_from_provider(self) -> CredentialConfiguration:
"""Create credential configuration from provider settings.
Returns
-------
CredentialConfiguration
Credential configuration based on provider settings
"""
# Extract credential settings from provider
role_arn = getattr(self.provider, "role_arn", None)
aws_profile = getattr(self.provider, "aws_profile", None)
use_env_vars = (
hasattr(self.provider, "aws_access_key_id")
and self.provider.aws_access_key_id is not None
)
# Create credential configuration
config = CredentialConfiguration(
role_arn=role_arn,
enable_sanitization=True,
sanitize_logs=True,
use_environment_variables=use_env_vars,
use_profile=aws_profile,
auto_refresh_tokens=True,
)
# Set security-based defaults
if self.security_config.environment.value == "production":
config.use_environment_variables = False
config.use_profile = None
config.require_mfa = False
logger.info(
f"Lambda Created credential config: role_arn={bool(role_arn)}, "
f"profile={aws_profile}, use_env={use_env_vars}"
)
return config
def _create_lambda_execution_role(self) -> str:
"""Get or create an IAM role for Lambda execution (idempotent).
Returns
-------
str
ARN of the IAM role
"""
role_name = f"{TAG_PREFIX}-lambda-role-{self.provider.workflow_id}"
assume_role_policy = {
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Principal": {"Service": "lambda.amazonaws.com"},
"Action": "sts:AssumeRole",
}
],
}
role_arn = get_or_create_iam_role(
iam_client=self.iam_client,
role_name=role_name,
assume_role_policy=assume_role_policy,
policy_arns=[
"arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
],
tags=[
{"Key": TAG_MANAGED, "Value": "true"},
{"Key": TAG_WORKFLOW_ID, "Value": self.provider.workflow_id},
],
description=f"Execution role for Parsl Lambda functions ({self.provider.workflow_id})",
)
# Track role for cleanup
self.role_names.add(role_name)
# Wait for IAM propagation using role_exists waiter
try:
waiter = self.iam_client.get_waiter("role_exists")
waiter.wait(RoleName=role_name, WaiterConfig={"MaxAttempts": 10})
except Exception:
logger.debug(
"IAM waiter not available; proceeding without propagation wait"
)
return role_arn
def _create_lambda_function(self, job_id: str, command: str) -> str:
"""Create a Lambda function for job execution.
Parameters
----------
job_id : str
ID of the job
command : str
Command to execute
Returns
-------
str
Name of the Lambda function
"""
# Generate a unique function name
function_name = f"{TAG_PREFIX}-func-{self.provider.workflow_id}-{job_id}"
try:
# Create execution role if needed
role_arn = self._create_lambda_execution_role()
# Generate Lambda function code
zip_file = self._generate_lambda_code(command)
# Create Lambda function
response = self.lambda_client.create_function(
FunctionName=function_name,
Runtime=DEFAULT_LAMBDA_RUNTIME,
Role=role_arn,
Handler=DEFAULT_LAMBDA_HANDLER,
Code={"ZipFile": zip_file},
Description=f"Parsl job {job_id} for workflow {self.provider.workflow_id}",
Timeout=min(
self.provider.lambda_timeout, 900
), # Lambda max is 900s (15 min)
MemorySize=self.provider.lambda_memory,
Tags={
TAG_MANAGED: "true",
TAG_WORKFLOW_ID: self.provider.workflow_id,
TAG_JOB_ID: job_id,
},
)
# Track function for cleanup
self.function_names.add(function_name)
logger.info(f"Created Lambda function: {function_name}")
# Log successful Lambda function creation
if self.audit_logger:
self.audit_logger.log_resource_operation(
operation="create",
resource_type="lambda_function",
resource_id=function_name,
success=True,
workflow_id=self.provider.workflow_id,
job_id=job_id,
runtime=DEFAULT_LAMBDA_RUNTIME,
timeout=min(self.provider.lambda_timeout, 900),
memory=self.provider.lambda_memory,
)
return function_name
except ClientError as e:
logger.error(f"Error creating Lambda function: {e}")
# Log failed Lambda function creation
if self.audit_logger:
self.audit_logger.log_resource_operation(
operation="create",
resource_type="lambda_function",
resource_id=function_name,
success=False,
workflow_id=self.provider.workflow_id,
job_id=job_id,
error=str(e),
)
raise ResourceCreationError(f"Failed to create Lambda function: {e}")
def _generate_lambda_code(self, command: str) -> bytes:
"""Generate code for the Lambda function.
Parameters
----------
command : str
Command to execute
Returns
-------
bytes
Zip file content containing the Lambda function code
"""
import io
import zipfile
# A plain template, not an f-string: the handler is full of dict literals
# whose braces an f-string would read as replacement fields. The command
# is the only substitution, and it goes in as a JSON literal so any
# quoting in it survives.
handler_template = """
import json
import subprocess
import sys
import os
import traceback
DEFAULT_COMMAND = __COMMAND__
def main(event, context):
print("Starting Parsl job execution")
try:
# Get command from event or use the baked-in command
command = event.get('command', DEFAULT_COMMAND)
print(f"Executing command: {command}")
# Execute the command
result = subprocess.run(
command,
shell=True,
capture_output=True,
text=True
)
# Prepare response
response = {
'statusCode': 200 if result.returncode == 0 else 500,
'command': command,
'stdout': result.stdout,
'stderr': result.stderr,
'returncode': result.returncode
}
# Log results
print(f"Command completed with return code: {result.returncode}")
print(f"STDOUT: {result.stdout[:1000]}")
print(f"STDERR: {result.stderr[:1000]}")
return response
except Exception as e:
# Log the exception
print(f"Error executing command: {e}")
traceback.print_exc()
# Return error response
return {
'statusCode': 500,
'error': str(e),
'traceback': traceback.format_exc()
}
"""
handler_code = handler_template.replace("__COMMAND__", json.dumps(command))
# Create a ZIP file in memory
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
zip_file.writestr("handler.py", handler_code)
return zip_buffer.getvalue()
[docs]
def submit_job(self, job_id: str, command: str) -> Dict[str, Any]:
"""Submit a job for execution.
Parameters
----------
job_id : str
ID of the job
command : str
Command to execute
Returns
-------
Dict[str, Any]
Dictionary containing job information
"""
try:
# Create Lambda function
function_name = self._create_lambda_function(job_id, command)
# Invoke function asynchronously
response = self.lambda_client.invoke(
FunctionName=function_name,
InvocationType="Event", # Asynchronous invocation
Payload=json.dumps({"command": command, "job_id": job_id}),
)
# Verify the async invocation was accepted (202 = queued successfully)
status_code = response.get("StatusCode")
if status_code != 202:
raise JobSubmissionError(
f"Lambda async invocation failed for job {job_id}: "
f"StatusCode={status_code}"
)
# FunctionError is absent for async invocations but guard defensively
if response.get("FunctionError"):
raise JobSubmissionError(
f"Lambda invocation error for job {job_id}: "
f"{response['FunctionError']}"
)
# Get request ID from response
request_id = response.get("ResponseMetadata", {}).get("RequestId")
# Record job information
self.jobs[job_id] = {
"id": job_id,
"function_name": function_name,
"command": command,
"request_id": request_id,
"status": STATUS_PENDING,
"submitted_at": time.time(),
}
logger.info(f"Submitted job {job_id} to Lambda function {function_name}")
return {
"job_id": job_id,
"function_name": function_name,
"request_id": request_id,
}
except Exception as e:
logger.error(f"Error submitting job: {e}")
raise JobSubmissionError(f"Failed to submit job: {e}")
[docs]
def get_job_status(self, function_name: str, request_id: str) -> str:
"""Get the status of a job.
Parameters
----------
function_name : str
Name of the Lambda function
request_id : str
Request ID from the function invocation
Returns
-------
str
Job status
"""
try:
# For AWS Lambda, we can't directly query the status of an async invocation
# We'd need to implement a more complex solution, such as:
# 1. Use CloudWatch Logs to check for completion
# 2. Use a state store (DynamoDB, etc.) that the Lambda updates
# 3. Use Step Functions for workflow tracking
# For now, we'll simulate the status based on time elapsed
# In a real implementation, we'd use one of the approaches above
# Find the job
job = None
for j in self.jobs.values():
if (
j.get("function_name") == function_name
and j.get("request_id") == request_id
):
job = j
break
if not job:
return "UNKNOWN"
# If the job already has a terminal status, return it
if job["status"] in [STATUS_SUCCEEDED, STATUS_FAILED]:
return job["status"]
# Otherwise, simulate status based on time elapsed
elapsed = time.time() - job["submitted_at"]
if elapsed < 5:
status = STATUS_PENDING
elif elapsed < self.provider.lambda_timeout:
status = STATUS_RUNNING
else:
# After timeout, mark as COMPLETED.
# Real status requires CloudWatch Logs integration (v0.3.0).
#
# The job ID comes from the job record, not a `job_id` local:
# this method is called with (function_name, request_id) and
# locates the job by scanning. Interpolating a bare `job_id` here
# raised NameError, which the blanket `except` below turned into
# "UNKNOWN" -- so a job that outlived its timeout could never
# reach a terminal status and was polled forever (#111).
logger.warning(
f"Lambda job {job.get('id')} exceeded configured timeout. "
"Marking as COMPLETED. Integrate CloudWatch Logs in v0.3.0 "
"for accurate terminal status."
)
status = STATUS_SUCCEEDED
# Update job status
job["status"] = status
return status
except Exception as e:
logger.error(f"Error getting job status: {e}")
return "UNKNOWN"
[docs]
def cleanup_all_resources(self) -> None:
"""Clean up all AWS resources created by this manager."""
try:
# Delete Lambda functions
for function_name in list(self.function_names):
try:
self.lambda_client.delete_function(FunctionName=function_name)
logger.info(f"Deleted Lambda function: {function_name}")
self.function_names.remove(function_name)
except Exception as e:
logger.error(f"Error deleting Lambda function {function_name}: {e}")
# Detach and delete IAM roles
for role_name in list(self.role_names):
try:
# Detach policies
try:
self.iam_client.detach_role_policy(
RoleName=role_name,
PolicyArn="arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole",
)
except Exception as e:
logger.error(
f"Error detaching policy from role {role_name}: {e}"
)
# Delete role
self.iam_client.delete_role(RoleName=role_name)
logger.info(f"Deleted IAM role: {role_name}")
self.role_names.remove(role_name)
except Exception as e:
logger.error(f"Error deleting IAM role {role_name}: {e}")
except Exception as e:
logger.error(f"Error cleaning up resources: {e}")
raise ResourceCleanupError(f"Failed to clean up resources: {e}")