Update app.py

app.py

@@ -21,10 +21,14 @@ HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
-storage_client = storage.Client.from_service_account_info(credentials_info)
-bucket = storage_client.bucket(GCS_BUCKET_NAME)
-logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
+try:
+    credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
+    storage_client = storage.Client.from_service_account_info(credentials_info)
+    bucket = storage_client.bucket(GCS_BUCKET_NAME)
+    logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
+except Exception as e:
+    logger.error(f"Error al cargar las credenciales o bucket: {e}")
+    raise RuntimeError(f"Error al cargar las credenciales o bucket: {e}")
 
 app = FastAPI()
 
@@ -42,65 +46,142 @@ class GCSHandler:
 
     def create_folder_if_not_exists(self, folder_name):
         if not self.file_exists(folder_name):
+            logger.debug(f"Creando carpeta {folder_name} en GCS.")
             self.bucket.blob(folder_name + "/").upload_from_string("")
 
     def upload_file(self, blob_name, file_stream):
         self.create_folder_if_not_exists(os.path.dirname(blob_name))
         blob = self.bucket.blob(blob_name)
-        blob.upload_from_file(file_stream)
+        try:
+            blob.upload_from_file(file_stream)
+            logger.info(f"Archivo '{blob_name}' subido exitosamente a GCS.")
+        except Exception as e:
+            logger.error(f"Error subiendo el archivo '{blob_name}' a GCS: {e}")
+            raise HTTPException(status_code=500, detail=f"Error subiendo archivo '{blob_name}' a GCS")
 
     def download_file(self, blob_name):
         blob = self.bucket.blob(blob_name)
         if not blob.exists():
+            logger.error(f"Archivo '{blob_name}' no encontrado en GCS.")
             raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
-        return BytesIO(blob.download_as_bytes())
+        return blob.open("rb")
+
+    def generate_signed_url(self, blob_name, expiration=3600):
+        blob = self.bucket.blob(blob_name)
+        return blob.generate_signed_url(expiration=expiration)
 
 def download_model_from_huggingface(model_name):
-    base_url = f"https://huggingface.co/{model_name}/resolve/main/"
+    url = f"https://huggingface.co/{model_name}/tree/main"
+    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
+
+    try:
+        logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
+        response = requests.get(url, headers=headers)
+        if response.status_code == 200:
+            model_files = [
+                "pytorch_model.bin",
+                "config.json",
+                "tokenizer.json",
+                "model.safetensors",
+            ]
+            for file_name in model_files:
+                file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
+                file_content = requests.get(file_url).content
+                blob_name = f"lilmeaty_garca/{model_name}/{file_name}"
+                bucket.blob(blob_name).upload_from_string(file_content)
+                logger.info(f"Archivo '{file_name}' subido exitosamente al bucket GCS.")
+        else:
+            logger.error(f"Error al acceder al árbol de archivos de Hugging Face para '{model_name}'.")
+            raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
+    except Exception as e:
+        logger.error(f"Error descargando archivos de Hugging Face: {e}")
+        raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
+
+def download_and_verify_model(model_name):
     model_files = [
-        "pytorch_model.bin", "config.json", "tokenizer.json", "model.safetensors",
-        "pytorch_model.bin.index.json", "tokenizer_config.json",
-        "special_tokens_map.json", "vocab.json", "merges.txt"
+        "pytorch_model.bin",
+        "config.json",
+        "tokenizer.json",
+        "model.safetensors",
     ]
-    for filename in model_files:
-        try:
-            url = base_url + filename
-            response = requests.get(url, stream=True, headers={"Authorization": f"Bearer {HF_API_TOKEN}"})
-            response.raise_for_status()
-            blob_name = f"lilmeaty_garca/{model_name}/{filename}"
-            gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-            gcs_handler.upload_file(blob_name, response.raw)
-        except requests.exceptions.RequestException as e:
-            logger.warning(f"No se pudo descargar {filename} para {model_name}: {e}")
+    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+    model_files_exist = all(gcs_handler.file_exists(f"lilmeaty_garca/{model_name}/{file}") for file in model_files)
+    if not model_files_exist:
+        download_model_from_huggingface(model_name)
 
 def load_model_from_gcs(model_name):
+    model_files = [
+        "pytorch_model.bin",
+        "config.json",
+        "tokenizer.json",
+        "model.safetensors",
+    ]
+    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+    model_files_streams = {}
+
+    for file in model_files:
+        file_path = f"lilmeaty_garca/{model_name}/{file}"
+        if gcs_handler.file_exists(file_path):
+            model_files_streams[file] = gcs_handler.download_file(file_path)
+        else:
+            logger.error(f"Archivo '{file}' no encontrado en GCS.")
+            raise HTTPException(status_code=500, detail=f"Archivo '{file}' no encontrado.")
+
+    model_stream = model_files_streams.get("pytorch_model.bin") or model_files_streams.get("model.safetensors")
+    tokenizer_stream = model_files_streams.get("tokenizer.json")
+    config_stream = model_files_streams.get("config.json")
+
+    model = AutoModelForCausalLM.from_pretrained(model_stream, config=config_stream)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
+
+    return model, tokenizer
+
+def load_model(model_name):
     gcs_handler = GCSHandler(GCS_BUCKET_NAME)
     try:
-        model = AutoModelForCausalLM.from_pretrained(f"gs://{GCS_BUCKET_NAME}/lilmeaty_garca/{model_name}")
-        tokenizer = AutoTokenizer.from_pretrained(f"gs://{GCS_BUCKET_NAME}/lilmeaty_garca/{model_name}")
+        model, tokenizer = load_model_from_gcs(model_name)
+        logger.info(f"Modelo '{model_name}' cargado exitosamente desde GCS.")
+        return model, tokenizer
+    except HTTPException:
+        logger.warning(f"Modelo '{model_name}' no encontrado en GCS. Intentando como model_id...")
+
+    try:
+        download_and_verify_model(model_name)
+        model, tokenizer = load_model_from_gcs(model_name)
+        logger.info(f"Modelo '{model_name}' cargado exitosamente desde Hugging Face.")
         return model, tokenizer
     except Exception as e:
-        logger.error(f"Error al cargar el modelo '{model_name}' desde GCS: {e}")
+        logger.error(f"Error al intentar cargar el modelo '{model_name}': {e}")
         raise HTTPException(status_code=500, detail=f"Error al cargar el modelo '{model_name}': {e}")
 
 @app.on_event("startup")
 async def startup():
-    def download_all_models_in_background():
-        models_url = "https://huggingface.co/api/models?full=true&limit=100"
-        try:
-            while models_url:
-                response = requests.get(models_url)
-                response.raise_for_status()
-                models_data = response.json()
-                for model in models_data["models"]:  # Corrected: Access 'models' list
-                    model_name = model["id"]
-                    download_model_from_huggingface(model_name)
-
-                models_url = models_data.get("next") #removed , None because its not necessary
-        except Exception as e:
-            logger.error(f"Error al descargar modelos en segundo plano: {e}")
-
-    threading.Thread(target=download_all_models_in_background, daemon=True).start()
+    try:
+        logger.info("Iniciando la descarga de modelos en segundo plano...")
+        run_in_background()
+        gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+        blobs = list(bucket.list_blobs(prefix="lilmeaty_garca/"))
+        model_names = set([blob.name.split("/")[1] for blob in blobs])
+
+        def download_model_thread(model_name):
+            try:
+                download_and_verify_model(model_name)
+            except Exception as e:
+                logger.error(f"Error descargando modelo '{model_name}': {e}")
+
+        threads = []
+        for model_name in model_names:
+            thread = threading.Thread(target=download_model_thread, args=(model_name,))
+            thread.start()
+            threads.append(thread)
+
+        for thread in threads:
+            thread.join()
+
+        logger.info("Todos los modelos se descargaron correctamente o ya estaban presentes.")
+    except Exception as e:
+        logger.error(f"Error durante la descarga de modelos al iniciar: {e}")
+        raise HTTPException(status_code=500, detail=f"Error durante la descarga de modelos: {e}")
 
 @app.post("/predict/")
 async def predict(request: DownloadModelRequest):
@@ -108,13 +189,34 @@ async def predict(request: DownloadModelRequest):
         model_name = request.model_name
         pipeline_task = request.pipeline_task
         input_text = request.input_text
-        model, tokenizer = load_model_from_gcs(model_name)
-        pipe = pipeline(pipeline_task, model=model, tokenizer=tokenizer, device=0 if os.getenv("USE_GPU") else -1)
+
+        model, tokenizer = load_model(model_name)
+        pipe = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
         result = pipe(input_text)
+
         return {"result": result}
     except Exception as e:
         logger.error(f"Error procesando la solicitud: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+def download_all_models_in_background():
+    models_url = "https://huggingface.co/api/models"
+    try:
+        response = requests.get(models_url)
+        if response.status_code != 200:
+            logger.error("Error al obtener la lista de modelos de Hugging Face.")
+            raise HTTPException(status_code=500, detail="Error al obtener la lista de modelos.")
+
+        models = response.json()
+        for model in models:
+            model_name = model["id"]
+            download_model_from_huggingface(model_name)
+    except Exception as e:
+        logger.error(f"Error al descargar modelos en segundo plano: {e}")
+        raise HTTPException(status_code=500, detail="Error al descargar modelos en segundo plano.")
+
+def run_in_background():
+    threading.Thread(target=download_all_models_in_background, daemon=True).start()
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)