336 lines
13 KiB
Python
336 lines
13 KiB
Python
import argparse
|
|
import json
|
|
import logging
|
|
import logging.handlers
|
|
import os
|
|
import random
|
|
import time
|
|
from multiprocessing import Pool, Queue, current_process
|
|
from urllib.parse import urljoin
|
|
|
|
import requests
|
|
from dotenv import load_dotenv
|
|
|
|
# --- 전역 변수 및 로깅 시스템 ---
|
|
|
|
worker_log_queue = None
|
|
completed_tasks_count = 0
|
|
total_tasks = 0
|
|
|
|
|
|
def setup_main_logger(log_queue):
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
formatter = logging.Formatter(
|
|
"%(asctime)s - %(processName)s - %(levelname)s - %(message)s"
|
|
)
|
|
console_handler = logging.StreamHandler()
|
|
console_handler.setFormatter(formatter)
|
|
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
|
log_filename = f"{timestamp}_robust_run.log"
|
|
workspace_dir = os.path.dirname(os.path.abspath(__file__)) or "."
|
|
log_filepath = os.path.join(workspace_dir, log_filename)
|
|
file_handler = logging.FileHandler(log_filepath, encoding="utf-8")
|
|
file_handler.setFormatter(formatter)
|
|
queue_listener = logging.handlers.QueueListener(
|
|
log_queue, console_handler, file_handler
|
|
)
|
|
queue_listener.start()
|
|
|
|
bootstrap_logger = logging.getLogger("bootstrap")
|
|
bootstrap_logger.setLevel(logging.INFO)
|
|
bootstrap_logger.addHandler(console_handler)
|
|
bootstrap_logger.addHandler(file_handler)
|
|
bootstrap_logger.propagate = False
|
|
|
|
bootstrap_logger.info(f"메인 로거 설정 완료. 로그 파일: '{log_filepath}'")
|
|
bootstrap_logger.info("=" * 20 + " 스크립트 시작 " + "=" * 20)
|
|
return queue_listener, bootstrap_logger
|
|
|
|
|
|
def init_worker(log_queue):
|
|
global worker_log_queue
|
|
worker_log_queue = log_queue
|
|
|
|
|
|
def get_worker_logger():
|
|
global worker_log_queue
|
|
worker_logger = logging.getLogger(current_process().name)
|
|
if not worker_logger.handlers:
|
|
worker_logger.setLevel(logging.INFO)
|
|
queue_handler = logging.handlers.QueueHandler(worker_log_queue)
|
|
worker_logger.addHandler(queue_handler)
|
|
return worker_logger
|
|
|
|
|
|
# --- API 요청 함수 (변경 없음) ---
|
|
def start_extraction(post_url, file_path, filename, headers, model_name=None):
|
|
worker_logger = get_worker_logger()
|
|
worker_logger.info(f"[{filename}] 파일 업로드 및 추출 요청 시작...")
|
|
# ... (내부 로직은 이전과 동일)
|
|
start_time = time.time()
|
|
try:
|
|
with open(file_path, "rb") as input_f:
|
|
files_to_upload = {"input_file": (filename, input_f)}
|
|
data_payload = {}
|
|
if model_name:
|
|
data_payload["model"] = model_name
|
|
response = requests.post(
|
|
post_url,
|
|
files=files_to_upload,
|
|
data=data_payload,
|
|
headers=headers,
|
|
timeout=300,
|
|
)
|
|
end_time = time.time()
|
|
worker_logger.info(
|
|
f"[{filename}] 추출 요청 완료. 소요 시간: {end_time - start_time:.2f}초, 상태 코드: {response.status_code}"
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except requests.exceptions.RequestException as e:
|
|
end_time = time.time()
|
|
worker_logger.exception(
|
|
f"[{filename}] 추출 요청 중 치명적 오류 발생. 소요 시간: {end_time - start_time:.2f}초. 오류: {e}"
|
|
)
|
|
return None
|
|
|
|
|
|
def check_progress(base_url, progress_path, filename, headers):
|
|
worker_logger = get_worker_logger()
|
|
get_url = urljoin(base_url + "/", progress_path.lstrip("/"))
|
|
# ... (내부 로직은 이전과 동일)
|
|
base_delay = 5
|
|
max_delay = 60
|
|
attempt = 0
|
|
last_status = ""
|
|
while True:
|
|
delay = min(max_delay, base_delay * (2**attempt))
|
|
jitter = delay * 0.2
|
|
sleep_duration = delay + random.uniform(-jitter, jitter)
|
|
worker_logger.info(
|
|
f"[{filename}] {sleep_duration:.2f}초 후 상태 확인 (시도 {attempt + 1})..."
|
|
)
|
|
time.sleep(sleep_duration)
|
|
start_time = time.time()
|
|
try:
|
|
response = requests.get(get_url, headers=headers, timeout=60)
|
|
end_time = time.time()
|
|
if response.status_code != 200:
|
|
worker_logger.warning(
|
|
f"[{filename}] 상태 확인 응답 코드: {response.status_code}. 소요 시간: {end_time - start_time:.2f}초"
|
|
)
|
|
if response.status_code == 404:
|
|
if attempt < 5:
|
|
worker_logger.warning(
|
|
f"[{filename}] 작업을 찾을 수 없음(404). 재시도합니다."
|
|
)
|
|
attempt += 1
|
|
continue
|
|
else:
|
|
worker_logger.error(
|
|
f"[{filename}] 재시도 횟수 초과 후에도 작업을 찾을 수 없음(404). 처리 중단."
|
|
)
|
|
return None
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
if "final_result" in data and data.get("final_result") is not None:
|
|
worker_logger.info(f"[{filename}] 처리 완료.")
|
|
return data["final_result"]
|
|
if "progress_logs" in data and data["progress_logs"]:
|
|
status_message = data["progress_logs"][-1].get(
|
|
"status", "상태 확인 중..."
|
|
)
|
|
if status_message != last_status:
|
|
last_status = status_message
|
|
worker_logger.info(f"[{filename}] 진행 상태: {last_status}")
|
|
attempt += 1
|
|
except requests.exceptions.ReadTimeout:
|
|
end_time = time.time()
|
|
worker_logger.warning(
|
|
f"[{filename}] 상태 확인 타임아웃. 소요 시간: {end_time - start_time:.2f}초. 재시도..."
|
|
)
|
|
attempt += 1
|
|
except requests.exceptions.RequestException as e:
|
|
end_time = time.time()
|
|
worker_logger.exception(
|
|
f"[{filename}] 상태 확인 중 치명적 오류 발생. 소요 시간: {end_time - start_time:.2f}초. 오류: {e}"
|
|
)
|
|
return None
|
|
|
|
|
|
# --- 워커 함수 및 콜백 함수 ---
|
|
|
|
|
|
# <--- 핵심 변경 1: 워커 함수가 순번(index)과 전체 개수(total)를 인자로 받음 ---
|
|
def process_file_worker(index, total, file_path, config):
|
|
"""단일 파일을 처리하는 워커 함수. 이제 순번과 전체 개수를 인자로 받습니다."""
|
|
worker_logger = get_worker_logger()
|
|
filename = os.path.basename(file_path)
|
|
|
|
# <--- 핵심 변경 2: 로그 출력 시 (순번/전체) 형식 사용 ---
|
|
worker_logger.info(f"--- ({index}/{total}) 처리 시작: {filename} ---")
|
|
|
|
initial_response = start_extraction(
|
|
config["post_url"], file_path, filename, config["headers"], config["model_name"]
|
|
)
|
|
if not initial_response:
|
|
worker_logger.error(f"[{filename}] 파일 처리 실패 (추출 시작 단계)")
|
|
return index, filename, False # <--- 핵심 변경 3: 결과 반환 시 순번도 함께 전달
|
|
|
|
request_id = initial_response.get("request_id")
|
|
status_check_url = initial_response.get("status_check_url")
|
|
|
|
if not request_id or not status_check_url:
|
|
worker_logger.error(
|
|
f"[{filename}] 초기 응답이 잘못되었습니다: {initial_response}"
|
|
)
|
|
return index, filename, False
|
|
|
|
worker_logger.info(f"[{filename}] 작업 요청 성공. Request ID: {request_id}")
|
|
|
|
final_result = check_progress(
|
|
config["base_url"], status_check_url, filename, config["headers"]
|
|
)
|
|
|
|
if final_result:
|
|
output_path = os.path.join(
|
|
config["output_dir"], f"{os.path.splitext(filename)[0]}.json"
|
|
)
|
|
try:
|
|
with open(output_path, "w", encoding="utf-8") as f:
|
|
json.dump(final_result, f, indent=2, ensure_ascii=False)
|
|
worker_logger.info(f"[{filename}] 결과 저장 완료: {output_path}")
|
|
return index, filename, True
|
|
except IOError:
|
|
worker_logger.exception(f"[{filename}] 파일 저장 중 오류 발생")
|
|
return index, filename, False
|
|
else:
|
|
worker_logger.error(f"[{filename}] 파일 처리 실패 (결과 확인 단계)")
|
|
return index, filename, False
|
|
|
|
|
|
def update_progress(result):
|
|
"""진행 상황을 업데이트하고 출력하는 콜백 함수."""
|
|
global completed_tasks_count, total_tasks
|
|
completed_tasks_count += 1
|
|
|
|
# <--- 핵심 변경 4: 콜백 함수가 순번을 함께 받아서 로그에 표시 ---
|
|
index, filename, success = result
|
|
status = "성공" if success else "실패"
|
|
logging.getLogger("bootstrap").info(
|
|
f"--- ({completed_tasks_count}/{total_tasks}) 완료: [{index}번째 파일] {filename} (상태: {status}) ---"
|
|
)
|
|
|
|
|
|
# --- 메인 실행 로직 ---
|
|
def main():
|
|
global total_tasks
|
|
# ... (인자 파싱 및 설정 부분은 이전과 동일) ...
|
|
dotenv_path = os.path.join(os.path.dirname(__file__), ".env")
|
|
load_dotenv(dotenv_path=dotenv_path)
|
|
BASE_URL = os.getenv("BASE_URL")
|
|
API_KEY = os.getenv("API_KEY")
|
|
parser = argparse.ArgumentParser(
|
|
description="문서 정보 추출 자동화 스크립트 (병렬 처리)"
|
|
)
|
|
parser.add_argument("input_dir", help="입력 디렉터리 경로")
|
|
parser.add_argument(
|
|
"-o", "--output_dir", default="results", help="출력 디렉터리 경로"
|
|
)
|
|
parser.add_argument(
|
|
"--endpoint",
|
|
choices=["i18n", "d6c"],
|
|
default="i18n",
|
|
help="추출 API 엔드포인트 선택",
|
|
)
|
|
parser.add_argument("--model", dest="model_name", help="사용할 LLM 모델 이름")
|
|
parser.add_argument(
|
|
"-w", "--workers", type=int, default=4, help="동시에 실행할 워커 프로세스 수"
|
|
)
|
|
args = parser.parse_args()
|
|
log_queue = Queue()
|
|
queue_listener, logger = setup_main_logger(log_queue)
|
|
if not BASE_URL or not API_KEY:
|
|
logger.error("환경 변수(BASE_URL, API_KEY)가 .env 파일에 설정되지 않았습니다.")
|
|
queue_listener.stop()
|
|
return
|
|
if not os.path.isdir(args.input_dir):
|
|
logger.error(f"입력 디렉터리를 찾을 수 없습니다 - {args.input_dir}")
|
|
queue_listener.stop()
|
|
return
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
headers = {"X-API-KEY": API_KEY}
|
|
post_url = f"{BASE_URL}/extract/inner/{args.endpoint}"
|
|
logger.info(f"API 서버: {BASE_URL}")
|
|
logger.info(f"요청 API: {post_url}")
|
|
logger.info(f"입력 디렉터리: {args.input_dir}")
|
|
logger.info(f"출력 디렉터리: {args.output_dir}")
|
|
logger.info(f"동시 작업 수 (워커): {args.workers}")
|
|
files_to_process = [
|
|
os.path.join(args.input_dir, f)
|
|
for f in sorted(os.listdir(args.input_dir))
|
|
if os.path.isfile(os.path.join(args.input_dir, f))
|
|
]
|
|
total_tasks = len(files_to_process)
|
|
logger.info(f"총 {total_tasks}개의 파일을 처리합니다.")
|
|
config = {
|
|
"base_url": BASE_URL,
|
|
"post_url": post_url,
|
|
"headers": headers,
|
|
"model_name": args.model_name,
|
|
"output_dir": args.output_dir,
|
|
}
|
|
|
|
results = []
|
|
pool = None
|
|
try:
|
|
pool = Pool(
|
|
processes=args.workers, initializer=init_worker, initargs=(log_queue,)
|
|
)
|
|
async_results = []
|
|
|
|
# <--- 핵심 변경 5: for 루프에서 enumerate를 사용하여 순번(i)을 함께 전달 ---
|
|
for i, file_path in enumerate(files_to_process):
|
|
# 작업 인자에 (순번, 전체 개수, 파일 경로, 설정)을 담음
|
|
task_args = (i + 1, total_tasks, file_path, config)
|
|
res = pool.apply_async(
|
|
process_file_worker, args=task_args, callback=update_progress
|
|
)
|
|
async_results.append(res)
|
|
|
|
pool.close()
|
|
pool.join()
|
|
|
|
results = [res.get() for res in async_results]
|
|
|
|
except KeyboardInterrupt:
|
|
logger.warning("사용자에 의해 작업이 중단되었습니다.")
|
|
if pool:
|
|
pool.terminate()
|
|
pool.join()
|
|
except Exception:
|
|
logging.getLogger(__name__).exception(
|
|
"메인 프로세스에서 예기치 않은 오류가 발생했습니다."
|
|
)
|
|
if pool:
|
|
pool.terminate()
|
|
pool.join()
|
|
|
|
# <--- 핵심 변경 6: 최종 결과 요약 시 순번을 제외하고 파일명만 사용 ---
|
|
successful_files = [filename for index, filename, success in results if success]
|
|
failed_files = [filename for index, filename, success in results if not success]
|
|
|
|
logger.info("=" * 20 + " 모든 작업 완료 " + "=" * 20)
|
|
logger.info(
|
|
f"총 {total_tasks}개 파일 중 {len(successful_files)}개 성공, {len(failed_files)}개 실패."
|
|
)
|
|
if failed_files:
|
|
logger.warning(f"실패한 파일 목록: {failed_files}")
|
|
|
|
queue_listener.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|