import os import json import boto3 import uvicorn from fastapi import FastAPI, HTTPException from pydantic import BaseModel from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline from huggingface_hub import hf_hub_download from io import BytesIO import torch import safetensors from dotenv import load_dotenv import tqdm import re # Cargar las variables de entorno desde el archivo .env load_dotenv() # Configuración de AWS y Hugging Face AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID") AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY") AWS_REGION = os.getenv("AWS_REGION") S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME") # Nombre del bucket S3 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN") # Token de Hugging Face # Cliente S3 de Amazon s3_client = boto3.client( 's3', aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY, region_name=AWS_REGION ) app = FastAPI() # Pydantic Model para el cuerpo de la solicitud del endpoint /predict/ class DownloadModelRequest(BaseModel): model_name: str pipeline_task: str input_text: str class S3DirectStream: def __init__(self, bucket_name): self.s3_client = boto3.client( 's3', aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY, region_name=AWS_REGION ) self.bucket_name = bucket_name def stream_from_s3(self, key): try: print(f"[INFO] Descargando archivo {key} desde S3...") response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key) return BytesIO(response['Body'].read()) # Devolver el cuerpo como BytesIO except self.s3_client.exceptions.NoSuchKey: raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.") except Exception as e: print(f"[ERROR] Error al descargar {key}: {str(e)}") raise HTTPException(status_code=500, detail=f"Error al descargar archivo {key} desde S3.") def file_exists_in_s3(self, key): try: self.s3_client.head_object(Bucket=self.bucket_name, Key=key) return True except self.s3_client.exceptions.ClientError: return False def load_model_from_s3(self, model_name): try: print(f"[INFO] Cargando el modelo {model_name} desde S3...") # Verificar si el modelo existe en S3 model_prefix = model_name.lower() model_files = self.get_model_file_parts(model_prefix) if not model_files: print(f"[INFO] El modelo {model_name} no está en S3. Procediendo a descargar desde Hugging Face...") self.download_and_upload_from_huggingface(model_name) model_files = self.get_model_file_parts(model_prefix) if not model_files: raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.") # Cargar todos los archivos del modelo desde S3 model_streams = [] for model_file in tqdm.tqdm(model_files, desc="Cargando archivos del modelo", unit="archivo"): model_streams.append(self.stream_from_s3(model_file)) config_stream = self.stream_from_s3(f"{model_prefix}/config.json") config_data = json.loads(config_stream.read().decode("utf-8")) # Cargar el modelo dependiendo del tipo de archivo (torch o safetensors) if any(file.endswith("model.safetensors") for file in model_files): print("[INFO] Cargando el modelo como safetensor...") model = AutoModelForCausalLM.from_config(config_data) model.load_state_dict(safetensors.torch.load_stream(model_streams[0])) else: print("[INFO] Cargando el modelo como archivo binario de PyTorch...") model = AutoModelForCausalLM.from_config(config_data) model.load_state_dict(torch.load(model_streams[0], map_location="cpu")) print("[INFO] Modelo cargado con éxito.") return model except Exception as e: print(f"[ERROR] Error al cargar el modelo desde S3: {e}") raise HTTPException(status_code=500, detail="Error al cargar el modelo desde S3.") def load_tokenizer_from_s3(self, model_name): try: print(f"[INFO] Cargando el tokenizer {model_name} desde S3...") tokenizer_stream = self.stream_from_s3(f"{model_name}/tokenizer.json") tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream) return tokenizer except Exception as e: print(f"[ERROR] Error al cargar el tokenizer desde S3: {e}") raise HTTPException(status_code=500, detail="Error al cargar el tokenizer desde S3.") def get_model_file_parts(self, model_name): print(f"[INFO] Listando archivos del modelo en S3 con prefijo {model_name}...") files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name) model_files = [] for obj in tqdm.tqdm(files.get('Contents', []), desc="Verificando archivos", unit="archivo"): key = obj['Key'] # Verificar si es un archivo relevante del modelo if re.match(rf"{model_name}/.*", key): model_files.append(key) if not model_files: print(f"[WARNING] No se encontraron archivos para el modelo {model_name}.") return model_files def download_and_upload_from_huggingface(self, model_name): try: print(f"[INFO] Descargando {model_name} desde Hugging Face...") # Descargar todos los archivos del modelo files_to_download = hf_hub_download(repo_id=model_name, use_auth_token=HUGGINGFACE_TOKEN) # Subir a S3 los archivos descargados for file in files_to_download: file_name = os.path.basename(file) s3_key = f"{model_name}/{file_name}" if not self.file_exists_in_s3(s3_key): self.upload_file_to_s3(file, s3_key) except Exception as e: print(f"[ERROR] Error al descargar y subir modelo desde Hugging Face: {e}") raise HTTPException(status_code=500, detail="Error al descargar y subir modelo desde Hugging Face.") def upload_file_to_s3(self, file_path, s3_key): try: print(f"[INFO] Subiendo archivo {file_path} a S3 con key {s3_key}...") with open(file_path, 'rb') as data: self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=data) # Eliminar archivo local después de ser subido os.remove(file_path) except Exception as e: print(f"[ERROR] Error al subir archivo a S3: {e}") raise HTTPException(status_code=500, detail="Error al subir archivo a S3.") @app.post("/predict/") async def predict(model_request: DownloadModelRequest): try: print(f"[INFO] Recibiendo solicitud para predecir con el modelo {model_request.model_name}...") streamer = S3DirectStream(S3_BUCKET_NAME) model = streamer.load_model_from_s3(model_request.model_name) tokenizer = streamer.load_tokenizer_from_s3(model_request.model_name) task = model_request.pipeline_task if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", "text-to-speech", "text-to-video", "text-to-image"]: raise HTTPException(status_code=400, detail="Pipeline task no soportado") # Configurar el pipeline de transformers según la tarea nlp_pipeline = None if task == "text-generation": nlp_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer) elif task == "sentiment-analysis": nlp_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer) elif task == "translation": nlp_pipeline = pipeline("translation", model=model, tokenizer=tokenizer) elif task == "fill-mask": nlp_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer) elif task == "question-answering": nlp_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer) elif task == "text-to-speech": nlp_pipeline = pipeline("text-to-speech", model=model, tokenizer=tokenizer) elif task == "text-to-video": nlp_pipeline = pipeline("text-to-video", model=model, tokenizer=tokenizer) elif task == "text-to-image": nlp_pipeline = pipeline("text-to-image", model=model, tokenizer=tokenizer) result = nlp_pipeline(model_request.input_text) return {"result": result} except Exception as e: print(f"[ERROR] Error en el proceso de predicción: {str(e)}") raise HTTPException(status_code=500, detail="Error en el proceso de predicción") # Ejecutar la app con Uvicorn if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)