import boto3 import logging import os import json import re from datetime import datetime from typing import Optional from supabase import create_client, Client from utils.rekognition_utils import validate_manifest_file, start_rekognition_training from utils.logging_utils import setup_logger, reduce_logging_verbosity from models.training import poll_training_status, evaluate_model, store_training_metadata # Initialize logger logger = setup_logger() reduce_logging_verbosity() # Initialize clients s3_client = boto3.client('s3') ssm_client = boto3.client('ssm') # Initialize Supabase client supabase: Optional[Client] = None def get_ssm_parameter(param_name): """Get a parameter from SSM Parameter Store""" try: response = ssm_client.get_parameter(Name=param_name, WithDecryption=True) return response['Parameter']['Value'] except Exception as e: logger.warning(f"Failed to get SSM parameter {param_name}: {str(e)}") return None def initialize_supabase(): """Initialize Supabase client using SSM parameters""" global supabase try: supabase_url = get_ssm_parameter('/supabase/url') supabase_key = get_ssm_parameter('/supabase/anon') if supabase_url and supabase_key: supabase = create_client(supabase_url, supabase_key) logger.info("Supabase client initialized using SSM parameters") return True else: logger.warning("Supabase credentials not found in SSM, metadata storage will be skipped") return False except Exception as e: logger.error(f"Failed to initialize Supabase client: {str(e)}") return False def fix_project_arn(project_arn): """ Fix the project ARN format by ensuring it ends with the correct project ID. AWS Rekognition project ARNs must use the correct project ID: 1725357683732 """ if not project_arn: logger.error("No project ARN provided") return None logger.info(f"Checking project ARN format: {project_arn}") # Use the known working project ID project_id = "1725357683732" # Check if the ARN already has a version number pattern = r'project\/[a-zA-Z0-9_\.-]{1,255}\/[0-9]+$' if re.search(pattern, project_arn): # If it matches the pattern but has a different ID, replace it if not project_arn.endswith(f"/{project_id}"): # Extract everything up to the last slash base_arn = '/'.join(project_arn.split('/')[:-1]) fixed_arn = f"{base_arn}/{project_id}" logger.info(f"Updated project ID: {project_arn} -> {fixed_arn}") return fixed_arn else: logger.info(f"Project ARN format is already correct: {project_arn}") return project_arn # Get just the project name without any trailing slashes if '/project/' in project_arn: base_arn = project_arn.rstrip('/') # Standard case: append the correct project ID fixed_arn = f"{base_arn}/{project_id}" logger.info(f"Fixed ARN by adding correct project ID: {project_arn} -> {fixed_arn}") return fixed_arn logger.warning(f"Cannot determine correct ARN format from: {project_arn}") # Fall back to simple approach if more sophisticated extraction fails if project_arn.endswith('/'): fixed_arn = f"{project_arn}{project_id}" else: fixed_arn = f"{project_arn}/{project_id}" logger.info(f"Applied fallback ARN fix: {project_arn} -> {fixed_arn}") return fixed_arn def lambda_handler(event, context): """Lambda handler for starting the training job""" try: # Get environment variables training_bucket = os.environ.get('TRAINING_BUCKET') project_arn = os.environ.get('PROJECT_ARN') if not training_bucket or not project_arn: missing = [] if not training_bucket: missing.append('TRAINING_BUCKET') if not project_arn: missing.append('PROJECT_ARN') # Try to get from SSM if not in environment if 'TRAINING_BUCKET' in missing: training_bucket = get_ssm_parameter('/datafy-rekognition-stack/training-bucket') if training_bucket: missing.remove('TRAINING_BUCKET') if 'PROJECT_ARN' in missing: project_name = get_ssm_parameter('/datafy-rekognition-stack/project-name') if project_name: project_arn = f"arn:aws:rekognition:{os.environ.get('AWS_REGION')}:{os.environ.get('AWS_ACCOUNT_ID')}:project/{project_name}" missing.remove('PROJECT_ARN') if missing: raise ValueError(f"Missing required configuration: {', '.join(missing)}") # Fix the project ARN format project_arn = fix_project_arn(project_arn) if not project_arn: raise ValueError("Invalid project ARN format") logger.info(f"Starting training job with bucket: {training_bucket}") logger.info(f"Using project ARN: {project_arn}") # Initialize Supabase initialize_supabase() # Verify manifest exists and is valid manifest_key = "manifest.jsonl" try: # Validate manifest file if not validate_manifest_file(training_bucket, manifest_key): return { 'statusCode': 400, 'body': json.dumps({ 'message': 'Manifest validation failed', 'error': 'Invalid manifest format or content' }) } # Generate version name version_name = f"version-{datetime.now().strftime('%Y%m%d-%H%M%S')}" # Start training training_response = start_rekognition_training( training_bucket=training_bucket, version_name=version_name, project_arn=project_arn ) # Prepare training info for metadata storage training_info = { 'version_name': version_name, 'project_version_arn': training_response.get('ProjectVersionArn'), 'project_arn': project_arn, 'status': 'TRAINING', 'manifest_file': f"s3://{training_bucket}/{manifest_key}", 'bucket': training_bucket } # Store metadata in Supabase if supabase: store_training_metadata(training_info, supabase) # Poll training status status_info = poll_training_status( project_arn=project_arn, version_name=version_name ) # Update training info with status training_info.update({ 'status': status_info['status'], 'status_message': status_info.get('status_message') }) # If training completed, evaluate the model if status_info['status'] == 'TRAINING_COMPLETED': evaluation = evaluate_model( project_arn=project_arn, version_name=version_name ) # Update training info with evaluation results training_info.update({ 'evaluation': evaluation['metrics'], 'meets_criteria': evaluation['meets_criteria'] }) # Update metadata in Supabase if supabase: store_training_metadata(training_info, supabase) return { 'statusCode': 200, 'body': json.dumps({ 'message': 'Training completed successfully', 'training_info': training_info, 'status': status_info, 'evaluation': evaluation }, default=str) } else: # Update metadata in Supabase with final status if supabase: store_training_metadata(training_info, supabase) return { 'statusCode': 200, 'body': json.dumps({ 'message': 'Training job started', 'training_info': training_info, 'status': status_info }, default=str) } except Exception as e: logger.error(f"Error starting training job: {str(e)}") return { 'statusCode': 500, 'body': json.dumps({ 'message': 'Failed to start training job', 'error': str(e) }) } except ValueError as e: logger.error(f"Configuration error: {str(e)}") return { 'statusCode': 500, 'body': json.dumps({ 'message': 'Configuration error', 'error': str(e) }) }