import os import boto3 import torch from fastapi import FastAPI, HTTPException from pydantic import BaseModel import safetensors.torch from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline from fastapi.responses import StreamingResponse import io import requests import uvicorn from dotenv import load_dotenv import re from tqdm import tqdm from huggingface_hub import hf_hub_download # 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 response['Body'] # Devolver el cuerpo directamente para el StreamingResponse 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_stream(self, model_name): try: print(f"[INFO] Cargando el modelo {model_name} desde S3...") model_prefix = model_name.lower() model_files = self.get_model_file_parts(model_prefix) # Verificar si el modelo está en S3 if model_files: print(f"[INFO] Modelo encontrado en S3, cargando directamente...") else: print(f"[INFO] Modelo no encontrado en S3, descargando desde Hugging Face...") self.download_and_upload_from_huggingface(model_name) model_files = self.get_model_file_parts(model_prefix) # Leer y cargar todos los archivos del modelo model_streams = [] for model_file in 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 = config_stream.read().decode("utf-8") # Cargar el modelo dependiendo del tipo de archivo (torch o safetensors) if model_files[0].endswith("model.safetensors"): 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_stream(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(files.get('Contents', []), desc="Verificando archivos", unit="archivo"): key = obj['Key'] if re.match(rf"{model_name}/model(-\d+-of-\d+)?", key) or key.endswith("model.safetensors"): 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 los archivos necesarios de Hugging Face usando huggingface_hub files_to_download = [ "pytorch_model.bin", "config.json", "tokenizer.json", ] for file in files_to_download: # Descargar cada archivo desde Hugging Face y subirlo a S3 file_path = hf_hub_download(repo_id=model_name, filename=file, use_auth_token=HUGGINGFACE_TOKEN) s3_key = f"{model_name}/{file}" self.upload_file_to_s3(file_path, 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) 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_stream(model_request.model_name) tokenizer = streamer.load_tokenizer_from_stream(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-image", "text-to-audio", "text-to-video"]: raise HTTPException(status_code=400, detail="Pipeline task no soportado") nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, max_length=2046) input_text = model_request.input_text print(f"[INFO] Ejecutando tarea {task} con el texto de entrada...") outputs = nlp_pipeline(input_text) # Eliminación de archivo local después de subir a S3 if task == "text-to-speech": s3_key = f"{model_request.model_name}/generated_audio.wav" return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="audio/wav") elif task == "text-to-image": s3_key = f"{model_request.model_name}/generated_image.png" return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png") elif task == "text-to-video": s3_key = f"{model_request.model_name}/generated_video.mp4" return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4") return {"output": outputs} except Exception as e: print(f"[ERROR] Error en el proceso de predicción: {str(e)}") raise HTTPException(status_code=500, detail="Error al realizar la predicción.") if __name__ == "__main__": print("Iniciando servidor de predicciones en localhost:8000") uvicorn.run(app, host="0.0.0.0", port=8000)