aws-samples / aws-samples/aws-ai-phi-deidentification
State machine design flaw
- Dominant language
- JavaScript
- Stars
- 20
- Forks
- 4
- PR merge metrics
- No merged PRs in 30d
Description
The lambda idp-init-phi-detection fails due to the s3 files not being ready - I've plugged the hole with the following update to the code for that lambda, but a better design for the statemachine (i.e. to trigger init-phi-detection when the files are added to the s3bucket) would be better:
```python
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
import boto3
import os
import json
import logging
import time
import random
import traceback
from typing import Dict, Any
# Configure logging
logger = logging.getLogger(__name__)
# AWS clients
comp_med = boto3.client('comprehendmedical', region_name='ca-central-1')
s3 = boto3.client('s3')
ddb = boto3.client('dynamodb')
# Environment variables
role = os.environ.get('IAM_ROLE')
def configure_logging() -> None:
"""Configure logging level based on environment variable."""
log_level = os.environ.get('LOG_LEVEL', 'INFO')
logger.setLevel(log_level)
# Add handler if not already added
if not logger.handlers:
handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.info(f"Logging configured with level: {log_level}")
def get_environment_variables() -> Dict[str, str]:
"""Get all environment variables."""
logger.debug("Retrieving environment variables")
env_vars = {}
for name, value in os.environ.items():
env_vars[name] = value
return env_vars
def update_error_state(env_vars: Dict[str, str], workflow_id: str) -> None:
"""Update workflow status to failed in DynamoDB.
Args:
env_vars: Dictionary containing environment variables
workflow_id: The ID of the workflow
"""
logger.info(f"Updating de-identification status to failed for workflow {workflow_id}")
try:
wf_update = f"UPDATE \"{env_vars['IDP_TABLE']}\" SET de_identification_status=? WHERE part_key=? AND sort_key=?"
ddb.execute_statement(Statement=wf_update, Parameters=[
{'S': 'failed'},
{'S': workflow_id},
{'S': f"input/{workflow_id}/"}
])
logger.info(f"Successfully updated status to failed for workflow {workflow_id}")
except Exception as e:
logger.error(f"Failed to update error state in DynamoDB: {str(e)}")
logger.error(f"Error type: {type(e).__name__}")
logger.error(traceback.format_exc())
def wait_for_directory_content(bucket: str, directory: str) -> bool:
"""Wait for S3 directory to have content using exponential backoff with jitter.
Args:
bucket: S3 bucket name
directory: S3 directory path/prefix
Returns:
bool: True if directory has content, False otherwise
Raises:
Exception: If directory is empty after maximum wait time
"""
logger.info(f"Checking for files in directory s3://{bucket}/{directory}")
max_total_wait = 600 # 10 minutes
base_delay = 5 # seconds
max_delay = 60 # seconds
total_wait = 0
attempt = 0
while total_wait < max_total_wait:
# Check for directory content
response = s3.list_objects_v2(Bucket=bucket, Prefix=directory)
# Correctly check for any object that is not a directory placeholder
if 'Contents' in response:
# Filter out the directory placeholder object itself
logger.info(f'Found {", ".join([obj["Key"] for obj in response.get("Contents", [])])} in s3://{bucket}/{directory}')
if len(response['Contents']) > 0:
logger.info(f"Returning true, directory s3://{bucket}/{directory} has content.")
return True
# Calculate next delay with exponential backoff and jitter
attempt += 1
delay = min(max_delay, base_delay * (2 ** attempt))
jitter = random.randint(0, min(5, delay // 2))
wait_time = delay + jitter
# Ensure we don't wait past the max_total_wait
if total_wait + wait_time > max_total_wait:
wait_time = max_total_wait - total_wait
if wait_time <= 0:
break
logger.info(
f"Directory not ready, waiting {wait_time} seconds before retry (attempt {attempt}, total wait so far: {total_wait}s)")
time.sleep(wait_time)
total_wait += wait_time
logger.warning(f"Directory s3://{bucket}/{directory} is empty after maximum wait time of {max_total_wait} seconds.")
raise Exception(f"Input directory is empty after maximum wait time: s3://{bucket}/{directory}")
def start_phi_detection_job(bucket: str, input_dir: str, output_dir: str, workflow_id: str) -> str:
"""Start a PHI detection job.
Args:
bucket: S3 bucket name
input_dir: S3 input directory
output_dir: S3 output directory
workflow_id: Workflow ID
Returns:
str: PHI job ID
"""
logger.info(f"Starting PHI detection job for s3://{bucket}/{input_dir}")
try:
response = comp_med.start_phi_detection_job(
InputDataConfig={
'S3Bucket': bucket,
'S3Key': input_dir
},
OutputDataConfig={
'S3Bucket': bucket,
'S3Key': output_dir
},
DataAccessRoleArn=role,
JobName=f'phi-job-{workflow_id}',
LanguageCode='en'
)
phi_job_id = response['JobId']
logger.info(f"PHI detection job started successfully with job ID: {phi_job_id}")
logger.debug(f"ComprehendMedical response: {json.dumps(response, default=str)}")
return phi_job_id
except Exception as e:
logger.error(f"Failed to start PHI detection job: {str(e)}")
logger.error(f"Error type: {type(e).__name__}")
logger.error(traceback.format_exc())
raise
def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
"""Lambda handler function for initiating PHI detection.
Args:
event: Lambda event
context: Lambda context
Returns:
Dict containing workflow details and PHI job ID
"""
configure_logging()
logger.info(f"Processing event: {json.dumps(event, default=str)}")
# Get environment variables
env_vars = get_environment_variables()
# Extract event parameters
workflow_id = event["workflow_id"]
phi_input_dir = event['phi_input_dir']
if not phi_input_dir.endswith('/'):
phi_input_dir += '/'
bucket = event["bucket"]
phi_output_dir = f'public/phi-output/{workflow_id}/'
phi_job_id = None
try:
# Wait for directory to have content
wait_for_directory_content(bucket, phi_input_dir)
# Start PHI detection job
phi_job_id = start_phi_detection_job(bucket, phi_input_dir, phi_output_dir, workflow_id)
except Exception as e:
logger.error(f"Error in PHI detection workflow for {workflow_id}: {str(e)}")
logger.error(f"Error type: {type(e).__name__}")
logger.error(traceback.format_exc())
update_error_state(env_vars, workflow_id)
# Return result
result = {
"workflow_id": workflow_id,
"bucket": bucket,
"phi_job_id": phi_job_id,
"phi_output_dir": phi_output_dir
}
logger.info(f"Lambda execution completed with result: {json.dumps(result, default=str)}")
return result
```
Contributor guide
Research direction
The issue identifies the idp-init-phi-detection Lambda, S3 input files, and the state machine as the relevant entry points; begin by tracing the state machine invocation and how files are added to the bucket. Define how detection should start only after the input files are available. Done means the workflow no longer depends on the Lambda's polling workaround and reliably starts the PHI detection job.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- backend, cloud
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100