Hwang
61c72a8a4c
변경 사항 ---- - _test_kiwoom_condition_list.py: 키움 웹소켓 조건검색 '목록조회' 기능을 단독으로 테스트하는 스크립트 추가 - _test_kiwoom_condition_realtime.py: 'momentum' 조건식을 실시간으로 등록하고 초기 매칭 종목 리스트 및 실시간 편입/이탈을 수신하는 테스트 스크립트 추가 - _verify_columnar_bitid.py, _verify_shared_e2e_breakout.py, _verify_shared_e2e.py: 공유 메모리 및 dict 간의 데이터 일관성을 검증하는 테스트 추가 영향 ---- - 신규 테스트 스크립트 추가로 키움 웹소켓 API의 기능 검증 및 안정성을 높임 - 기존 기능에 대한 영향 없음 Co-authored-by: Cursor <cursoragent@cursor.com>
712 lines
25 KiB
Python
712 lines
25 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
kis_trader/backtest/param_search_pool.py — 파라서치 ProcessPool 공통 안전장치
|
||
============================================================================
|
||
- PR_SET_PDEATHSIG + PPID 감시: 부모 종료·터미널 끊김 후에도 워커(PPID=1)가 CPU를 잡지 않게 함
|
||
- managed_process_pool: Ctrl+C / SIGTERM 시 cancel + terminate + kill 순 정리
|
||
- ParamSearchRunLock: 동일 전략 파라서치 중복 실행 방지
|
||
- ParamSearchSharedPayload: 캔들·유니버스 1회 파일 → 워커 init 1회 로드 (task마다 pickle 방지)
|
||
- param_search_max_workers: CPU 80%(고정) + 가용 RAM × PARAM_SEARCH_MEM_FRAC(기본 80%) 중 min
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import atexit
|
||
import os
|
||
import pickle
|
||
import signal
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
from concurrent.futures import (
|
||
FIRST_COMPLETED,
|
||
Future,
|
||
ProcessPoolExecutor,
|
||
as_completed,
|
||
wait,
|
||
)
|
||
from contextlib import contextmanager
|
||
from itertools import product
|
||
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar
|
||
|
||
from kis_trader.utils.env import get_env_float, get_env_int
|
||
|
||
# 워커 프로세스 전용 — fork/spawn 직후 worker_init·shared init 에서 기록
|
||
_EXPECTED_PPID: int = 0
|
||
_WORKER_SHARED: Optional[Dict[str, Any]] = None
|
||
|
||
T = TypeVar("T")
|
||
|
||
|
||
def worker_init() -> None:
|
||
"""ProcessPoolExecutor initializer — 워커 생성 직후 한 번 호출."""
|
||
global _EXPECTED_PPID
|
||
_EXPECTED_PPID = os.getppid()
|
||
try:
|
||
if sys.platform.startswith("linux"):
|
||
import ctypes
|
||
|
||
PR_SET_PDEATHSIG = 1
|
||
libc = ctypes.CDLL("libc.so.6", use_errno=True)
|
||
libc.prctl(PR_SET_PDEATHSIG, signal.SIGTERM, 0, 0, 0)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def worker_shared_get() -> Dict[str, Any]:
|
||
"""워커: ParamSearchSharedPayload 로 로드된 캔들·유니버스 등."""
|
||
return _WORKER_SHARED if isinstance(_WORKER_SHARED, dict) else {}
|
||
|
||
|
||
def _worker_init_from_shared_path(payload_path: str) -> None:
|
||
"""워커 1회: 디스크 pickle → _WORKER_SHARED (task 인자로 대용량 dict 전달 금지)."""
|
||
global _WORKER_SHARED
|
||
worker_init()
|
||
try:
|
||
with open(payload_path, "rb") as f:
|
||
_WORKER_SHARED = pickle.load(f)
|
||
except Exception:
|
||
_WORKER_SHARED = {}
|
||
|
||
|
||
def linux_mem_available_bytes() -> int:
|
||
"""Linux MemAvailable (없으면 MemFree)."""
|
||
try:
|
||
mem: Dict[str, int] = {}
|
||
with open("/proc/meminfo", "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
parts = line.split()
|
||
if len(parts) >= 2 and parts[0].rstrip(":") in (
|
||
"MemAvailable", "MemFree", "SwapFree",
|
||
):
|
||
mem[parts[0].rstrip(":")] = int(parts[1]) * 1024
|
||
if "MemAvailable" in mem:
|
||
return mem["MemAvailable"]
|
||
if "MemFree" in mem:
|
||
return mem["MemFree"]
|
||
except OSError:
|
||
pass
|
||
return 2 * 1024 ** 3
|
||
|
||
|
||
def estimate_pickle_bytes(obj: Any) -> int:
|
||
"""대용량 payload 예상 크기 (pickle 직렬화 길이)."""
|
||
try:
|
||
return len(pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL))
|
||
except Exception:
|
||
return 0
|
||
|
||
|
||
def downsample_combos_uniform(
|
||
combos: List[Any],
|
||
max_combos: int,
|
||
) -> Tuple[List[Any], int]:
|
||
"""그리드 조합이 많을 때 균등 간격 샘플링 (tail/momentum fast 공통)."""
|
||
cap = int(max_combos)
|
||
total = len(combos)
|
||
if cap <= 0 or total <= cap:
|
||
return combos, 0
|
||
if cap == 1:
|
||
return [combos[0]], total - 1
|
||
last = total - 1
|
||
picked: List[Any] = []
|
||
for i in range(cap):
|
||
idx = int(round(i * last / (cap - 1)))
|
||
picked.append(combos[idx])
|
||
return picked, max(0, total - len(picked))
|
||
|
||
|
||
def resolve_param_search_max_combos(
|
||
mode: str,
|
||
*,
|
||
strategy_env_prefix: str,
|
||
default_fast: int,
|
||
default_other: int = 0,
|
||
max_combos_override: Optional[int] = None,
|
||
) -> int:
|
||
"""
|
||
fast 모드 기본: 균등 샘플 상한.
|
||
우선순위: CLI override → {PREFIX}_FAST_MAX_COMBOS → PARAM_SEARCH_FAST_MAX_COMBOS.
|
||
"""
|
||
if max_combos_override is not None:
|
||
return max(0, int(max_combos_override))
|
||
m = (mode or "fast").strip().lower()
|
||
if m == "fast":
|
||
specific = get_env_int(f"{strategy_env_prefix}_FAST_MAX_COMBOS", 0)
|
||
if specific > 0:
|
||
return specific
|
||
return get_env_int("PARAM_SEARCH_FAST_MAX_COMBOS", default_fast)
|
||
cap = get_env_int(f"{strategy_env_prefix}_PARAM_SEARCH_MAX_COMBOS", default_other)
|
||
return max(0, cap)
|
||
|
||
|
||
def cap_combos_uniform(
|
||
dict_combos: List[Any],
|
||
mode: str,
|
||
*,
|
||
strategy_env_prefix: str,
|
||
default_fast: int,
|
||
default_other: int = 0,
|
||
max_combos_override: Optional[int] = None,
|
||
) -> Tuple[List[Any], int, int, int]:
|
||
"""그리드 전체 → (균등 샘플 리스트, 그리드 크기, 적용 cap, 제외 개수)."""
|
||
total_grid = len(dict_combos)
|
||
cap = resolve_param_search_max_combos(
|
||
mode,
|
||
strategy_env_prefix=strategy_env_prefix,
|
||
default_fast=default_fast,
|
||
default_other=default_other,
|
||
max_combos_override=max_combos_override,
|
||
)
|
||
dropped = 0
|
||
sampled = dict_combos
|
||
if cap > 0 and total_grid > cap:
|
||
sampled, dropped = downsample_combos_uniform(dict_combos, cap)
|
||
return sampled, total_grid, cap, dropped
|
||
|
||
|
||
def grid_total_combinations(axes: List[List[Any]]) -> int:
|
||
"""축별 값 개수의 곱 — 데카르트곱을 **펼치지 않고** 전체 조합 수만 계산.
|
||
|
||
(list(product(...)) 로 수천만 개를 RAM 에 펼치다 OOM 나는 것을 막기 위함)
|
||
"""
|
||
total = 1
|
||
for vals in axes:
|
||
n = len(vals)
|
||
if n <= 0:
|
||
return 0
|
||
total *= n
|
||
return total
|
||
|
||
|
||
def _decode_combo_by_index(axes: List[List[Any]], idx: int) -> Tuple[Any, ...]:
|
||
"""혼합진법 디코딩: 정수 idx → 각 축의 값 (마지막 축이 최하위 자리).
|
||
|
||
itertools.product 와 동일한 순서의 idx→조합 매핑이라, '전체 리스트를 안 만들고'
|
||
원하는 자리(idx)의 조합만 바로 꺼낼 수 있다. 초대형 그리드 균등샘플 핵심.
|
||
"""
|
||
out: List[Any] = [None] * len(axes)
|
||
for ax in range(len(axes) - 1, -1, -1):
|
||
n = len(axes[ax])
|
||
out[ax] = axes[ax][idx % n]
|
||
idx //= n
|
||
return tuple(out)
|
||
|
||
|
||
def _uniform_indices(total: int, count: int) -> List[int]:
|
||
"""0..total-1 에서 균등 간격 정수 인덱스 count개 (stage1 균등샘플과 동일한 간격)."""
|
||
if total <= 0 or count <= 0:
|
||
return []
|
||
if count == 1:
|
||
return [0]
|
||
if count >= total:
|
||
return list(range(total))
|
||
last = total - 1
|
||
seen: set = set()
|
||
out: List[int] = []
|
||
for i in range(count):
|
||
idx = int(round(i * last / (count - 1)))
|
||
if idx not in seen:
|
||
seen.add(idx)
|
||
out.append(idx)
|
||
return out
|
||
|
||
|
||
def cap_combos_uniform_lazy(
|
||
keys: List[str],
|
||
axes: List[List[Any]],
|
||
mode: str,
|
||
*,
|
||
strategy_env_prefix: str,
|
||
default_fast: int,
|
||
default_other: int = 0,
|
||
max_combos_override: Optional[int] = None,
|
||
valid_fn: Optional[Callable[[Dict[str, Any]], bool]] = None,
|
||
max_materialize: Optional[int] = None,
|
||
) -> Tuple[List[Dict[str, Any]], int, int, int]:
|
||
"""데카르트곱을 통째로 펼치지 않고 (cap 균등샘플 / 전수) dict 조합 생성.
|
||
|
||
- 전체 곱이 PARAM_SEARCH_MAX_MATERIALIZE(기본 200만) **이하**면 기존 경로 그대로
|
||
(materialize → 유효성 필터 → 균등 cap). 기존 동작·결과 100% 동일.
|
||
- 그보다 **크면** 인덱스 균등샘플: 균등 간격 인덱스를 혼합진법으로 디코딩해
|
||
그 조합만 생성한다. 수천만 개를 펼치지 않으므로 OOM 없음 + 균등 커버리지 유지.
|
||
|
||
반환 형식은 cap_combos_uniform 과 동일: (dict_combos, total_grid, cap, dropped)
|
||
"""
|
||
cap = resolve_param_search_max_combos(
|
||
mode,
|
||
strategy_env_prefix=strategy_env_prefix,
|
||
default_fast=default_fast,
|
||
default_other=default_other,
|
||
max_combos_override=max_combos_override,
|
||
)
|
||
if max_materialize is None:
|
||
max_materialize = get_env_int("PARAM_SEARCH_MAX_MATERIALIZE", 2_000_000)
|
||
raw_total = grid_total_combinations(axes)
|
||
|
||
# ── 작은 그리드: 기존 경로 그대로 (동작·결과 동일) ──
|
||
if raw_total <= max(1, max_materialize):
|
||
dict_combos = [dict(zip(keys, c)) for c in product(*axes)]
|
||
if valid_fn is not None:
|
||
dict_combos = [c for c in dict_combos if valid_fn(c)]
|
||
total_grid = len(dict_combos)
|
||
dropped = 0
|
||
if cap > 0 and total_grid > cap:
|
||
dict_combos, dropped = downsample_combos_uniform(dict_combos, cap)
|
||
return dict_combos, total_grid, cap, dropped
|
||
|
||
# ── 초대형 그리드: 펼치지 않고 인덱스 균등샘플 ──
|
||
target = cap if cap > 0 else max_materialize
|
||
# 유효성 필터로 일부 빠질 수 있어 넉넉히 뽑은 뒤(oversample) 통과분을 다시 균등 cap.
|
||
oversample_mult = max(1, get_env_int("PARAM_SEARCH_OVERSAMPLE_MULT", 4))
|
||
n_idx = min(raw_total, max(target, target * oversample_mult))
|
||
picked: List[Dict[str, Any]] = []
|
||
for idx in _uniform_indices(raw_total, n_idx):
|
||
c = dict(zip(keys, _decode_combo_by_index(axes, idx)))
|
||
if valid_fn is None or valid_fn(c):
|
||
picked.append(c)
|
||
if cap > 0 and len(picked) > cap:
|
||
picked, _ = downsample_combos_uniform(picked, cap)
|
||
dropped = max(0, raw_total - len(picked))
|
||
return picked, raw_total, cap, dropped
|
||
|
||
|
||
def param_search_workers_by_memory(payload_bytes: int) -> int:
|
||
"""
|
||
가용 RAM × PARAM_SEARCH_MEM_FRAC(기본 0.8) 안에서 워커 수 추정.
|
||
부모 1×payload + 워커당 ~(payload + WORKER_EXTRA) 보수적 가정.
|
||
"""
|
||
avail = linux_mem_available_bytes()
|
||
frac = get_env_float("PARAM_SEARCH_MEM_FRAC", 0.8)
|
||
budget = max(0, int(avail * frac))
|
||
parent_extra = get_env_int("PARAM_SEARCH_PARENT_EXTRA_MB", 512) * 1024 ** 2
|
||
worker_extra = get_env_int("PARAM_SEARCH_WORKER_EXTRA_MB", 384) * 1024 ** 2
|
||
fallback = get_env_int("PARAM_SEARCH_MB_PER_WORKER", 768) * 1024 ** 2
|
||
|
||
parent_need = (max(0, payload_bytes) + parent_extra) if payload_bytes > 0 else parent_extra + fallback
|
||
per_worker = (max(0, payload_bytes) + worker_extra) if payload_bytes > 0 else fallback
|
||
|
||
remain = budget - parent_need
|
||
if remain <= 0:
|
||
return 1
|
||
return max(1, remain // max(per_worker, 1))
|
||
|
||
|
||
def param_search_max_workers(payload_bytes: int = 0) -> int:
|
||
"""CPU 코어 × PARAM_SEARCH_CPU_FRAC(기본 0.8=80% 고정) vs 메모리 예산 — min."""
|
||
n_cpu = os.cpu_count() or 4
|
||
cpu_frac = get_env_float("PARAM_SEARCH_CPU_FRAC", 0.8)
|
||
cpu_workers = max(1, int(n_cpu * cpu_frac))
|
||
cap = get_env_int("PARAM_SEARCH_MAX_WORKERS", 0)
|
||
if cap > 0:
|
||
cpu_workers = min(cpu_workers, cap)
|
||
|
||
mem_workers = param_search_workers_by_memory(payload_bytes)
|
||
return max(1, min(cpu_workers, mem_workers))
|
||
|
||
|
||
def param_search_worker_budget_line(payload_bytes: int) -> str:
|
||
"""run_search 로그용 — 가용 메모리·워커 산출 근거."""
|
||
avail_gb = linux_mem_available_bytes() / (1024 ** 3)
|
||
frac = get_env_float("PARAM_SEARCH_MEM_FRAC", 0.8)
|
||
cpu_w = max(1, int((os.cpu_count() or 4) * get_env_float("PARAM_SEARCH_CPU_FRAC", 0.8)))
|
||
mem_w = param_search_workers_by_memory(payload_bytes)
|
||
workers = param_search_max_workers(payload_bytes)
|
||
payload_mb = payload_bytes / (1024 ** 2) if payload_bytes > 0 else 0.0
|
||
return (
|
||
f"💾 메모리: 가용 {avail_gb:.1f}GB × {frac * 100:.0f}% 예산 | "
|
||
f"payload ~{payload_mb:.0f}MB | CPU워커≤{cpu_w} RAM워커≤{mem_w} → 실제 {workers}"
|
||
)
|
||
|
||
|
||
def param_search_chunk_size(total: int, max_workers: int) -> int:
|
||
"""
|
||
청크 크기 — 조합 수가 적어도 워커 수 이상으로 나눠 CPU 포화.
|
||
(fast 240조합 × 청크 50 → 5청크만 → 5코어만 100% 사용되는 문제 방지)
|
||
"""
|
||
if total <= 0:
|
||
return 50
|
||
chunk_max = get_env_int("PARAM_SEARCH_CHUNK_MAX", 300)
|
||
mult = max(1, get_env_int("PARAM_SEARCH_CHUNK_MULTIPLIER", 4))
|
||
target_chunks = min(total, max(max_workers, max_workers * mult))
|
||
return max(1, min(chunk_max, -(-total // target_chunks)))
|
||
|
||
|
||
def param_search_chunk_plan(
|
||
total: int,
|
||
payload_bytes: int = 0,
|
||
) -> Tuple[int, int, int]:
|
||
"""(max_workers, chunk_size, num_chunks) — run_search 공통."""
|
||
workers = param_search_max_workers(payload_bytes)
|
||
size = param_search_chunk_size(total, workers)
|
||
n_chunks = max(1, -(-total // size))
|
||
return workers, size, n_chunks
|
||
|
||
|
||
class ParamSearchProgressETA:
|
||
"""멀티프로세스 파도식 청크 완료에 맞춘 ETA 추정기.
|
||
|
||
워커 N개가 동시에 청크를 처리하면 완료가 파도(배치) 단위로 몰린다.
|
||
청크 간격 기반 ETA는 0↔파도통째로 요동하므로, 워밍업 1파도 제외 누적 평균 사용.
|
||
"""
|
||
|
||
def __init__(self, total_chunks: int, max_workers: int) -> None:
|
||
self._total_chunks = max(1, total_chunks)
|
||
self._warmup_chunks = min(max(1, max_workers), self._total_chunks)
|
||
self._warmup_end_elapsed: Optional[float] = None
|
||
|
||
def remaining_sec(self, processed: int, elapsed_so_far: float) -> float:
|
||
"""남은 예상 초 — 워밍업 파도 이후 누적 평균 × 남은 청크."""
|
||
processed = max(1, processed)
|
||
remaining_chunks = self._total_chunks - processed
|
||
if remaining_chunks <= 0:
|
||
return 0.0
|
||
if processed == self._warmup_chunks:
|
||
self._warmup_end_elapsed = elapsed_so_far
|
||
if self._warmup_end_elapsed is not None and processed > self._warmup_chunks:
|
||
steady_elapsed = elapsed_so_far - self._warmup_end_elapsed
|
||
steady_done = processed - self._warmup_chunks
|
||
per_chunk = steady_elapsed / max(1, steady_done)
|
||
else:
|
||
per_chunk = elapsed_so_far / processed
|
||
return per_chunk * remaining_chunks
|
||
|
||
@staticmethod
|
||
def format_sec(sec: float) -> str:
|
||
"""초 → 'N분 N초' / 'N시간 N분 N초'."""
|
||
eta_m, eta_s = divmod(max(0, int(sec)), 60)
|
||
eta_h, eta_m = divmod(eta_m, 60)
|
||
if eta_h > 0:
|
||
return f"{eta_h}시간 {eta_m}분 {eta_s}초"
|
||
if eta_m > 0:
|
||
return f"{eta_m}분 {eta_s}초"
|
||
return f"{eta_s}초"
|
||
|
||
@staticmethod
|
||
def format_elapsed(elapsed_so_far: float) -> str:
|
||
"""경과 초 → 'N분 N초'."""
|
||
elapsed_m, elapsed_s = divmod(max(0, int(elapsed_so_far)), 60)
|
||
return f"{elapsed_m}분 {elapsed_s}초"
|
||
|
||
@staticmethod
|
||
def render_bar(frac: float, width: int = 20) -> str:
|
||
"""진행 바: 채움 ■ / 빈칸 □ (시각적 진행 표시용)."""
|
||
frac = max(0.0, min(1.0, frac))
|
||
filled = int(round(frac * width))
|
||
return "■" * filled + "□" * (width - filled)
|
||
|
||
|
||
def param_search_max_inflight(max_workers: int) -> int:
|
||
"""동시 pending Future 상한 (큐에 대용량 인자 적재 방지)."""
|
||
extra = get_env_int("PARAM_SEARCH_INFLIGHT_EXTRA", 1)
|
||
cap = get_env_int("PARAM_SEARCH_MAX_INFLIGHT", 0)
|
||
if cap > 0:
|
||
return max(1, cap)
|
||
return max(1, max_workers + max(0, extra))
|
||
|
||
|
||
class ParamSearchSharedPayload:
|
||
"""
|
||
부모: 캔들·유니버스 등 1회 pickle 파일.
|
||
워커: init 시 1회 로드 → evaluate 에서 worker_shared_get() 사용.
|
||
"""
|
||
|
||
def __init__(self, data: Dict[str, Any]) -> None:
|
||
self.data = dict(data)
|
||
self.path: str = ""
|
||
self._prepared = False
|
||
|
||
def estimate_bytes(self) -> int:
|
||
return estimate_pickle_bytes(self.data)
|
||
|
||
def prepare(self) -> str:
|
||
if self._prepared and self.path and os.path.isfile(self.path):
|
||
return self.path
|
||
fd, path = tempfile.mkstemp(prefix="kis_param_search_", suffix=".pkl")
|
||
os.close(fd)
|
||
with open(path, "wb") as f:
|
||
pickle.dump(self.data, f, protocol=pickle.HIGHEST_PROTOCOL)
|
||
self.path = path
|
||
self._prepared = True
|
||
return self.path
|
||
|
||
def cleanup(self) -> None:
|
||
if self.path and os.path.isfile(self.path):
|
||
try:
|
||
os.remove(self.path)
|
||
except OSError:
|
||
pass
|
||
self.path = ""
|
||
self._prepared = False
|
||
|
||
|
||
def make_process_pool(
|
||
max_workers: int,
|
||
*,
|
||
shared_payload_path: Optional[str] = None,
|
||
) -> ProcessPoolExecutor:
|
||
"""ProcessPoolExecutor 생성 (max_tasks_per_child 로 장시간 워커 누수 완화).
|
||
|
||
shared_payload_path: spawn 시 pickle 가능한 모듈 함수 + initargs 로 전달 (nested closure 금지).
|
||
"""
|
||
max_tasks = get_env_int("PARAM_SEARCH_MAX_TASKS_PER_CHILD", 50)
|
||
if shared_payload_path:
|
||
init_fn = _worker_init_from_shared_path
|
||
init_args: Tuple[Any, ...] = (shared_payload_path,)
|
||
else:
|
||
init_fn = worker_init
|
||
init_args = ()
|
||
kwargs: dict = {
|
||
"max_workers": max_workers,
|
||
"initializer": init_fn,
|
||
"initargs": init_args,
|
||
}
|
||
if max_tasks > 0:
|
||
kwargs["max_tasks_per_child"] = max_tasks
|
||
return ProcessPoolExecutor(**kwargs)
|
||
|
||
|
||
def shutdown_executor_hard(
|
||
executor: Optional[ProcessPoolExecutor],
|
||
*,
|
||
wait_sec: Optional[float] = None,
|
||
) -> None:
|
||
"""Executor 워커: cancel → shutdown(wait=False) → SIGTERM → SIGKILL."""
|
||
if executor is None:
|
||
return
|
||
if wait_sec is None:
|
||
wait_sec = get_env_float("PARAM_SEARCH_SHUTDOWN_WAIT_SEC", 8.0)
|
||
|
||
try:
|
||
executor.shutdown(wait=False, cancel_futures=True)
|
||
except TypeError:
|
||
executor.shutdown(wait=False)
|
||
except Exception:
|
||
pass
|
||
|
||
procs = getattr(executor, "_processes", None)
|
||
if not procs:
|
||
return
|
||
|
||
deadline = time.time() + wait_sec
|
||
for proc in list(procs.values()):
|
||
try:
|
||
if proc.is_alive():
|
||
proc.terminate()
|
||
except Exception:
|
||
pass
|
||
|
||
while time.time() < deadline:
|
||
if not any(p.is_alive() for p in procs.values()):
|
||
break
|
||
time.sleep(0.2)
|
||
|
||
for proc in procs.values():
|
||
try:
|
||
if proc.is_alive():
|
||
proc.kill()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
@contextmanager
|
||
def managed_process_pool(
|
||
max_workers: int,
|
||
*,
|
||
shared_payload: Optional[ParamSearchSharedPayload] = None,
|
||
) -> Iterator[ProcessPoolExecutor]:
|
||
"""정상 종료는 wait=True, KeyboardInterrupt 등은 hard shutdown."""
|
||
payload_path = shared_payload.prepare() if shared_payload else None
|
||
executor = make_process_pool(max_workers, shared_payload_path=payload_path)
|
||
try:
|
||
yield executor
|
||
except (KeyboardInterrupt, SystemExit):
|
||
shutdown_executor_hard(executor)
|
||
raise
|
||
except BaseException:
|
||
shutdown_executor_hard(executor)
|
||
raise
|
||
else:
|
||
try:
|
||
executor.shutdown(wait=True, cancel_futures=False)
|
||
except TypeError:
|
||
executor.shutdown(wait=True)
|
||
except Exception:
|
||
shutdown_executor_hard(executor)
|
||
finally:
|
||
if shared_payload is not None:
|
||
shared_payload.cleanup()
|
||
|
||
|
||
def iter_pool_chunk_results(
|
||
executor: ProcessPoolExecutor,
|
||
chunks: List[Any],
|
||
submit_fn: Callable[[Any], Future],
|
||
*,
|
||
max_workers: int = 1,
|
||
max_inflight: Optional[int] = None,
|
||
heartbeat_sec: Optional[float] = None,
|
||
on_heartbeat: Optional[Callable[[int, int, float, int], None]] = None,
|
||
) -> Iterator[Any]:
|
||
"""
|
||
청크별 Future — pending 수를 max_inflight 로 제한 (메모리·큐 폭주 방지).
|
||
submit_fn(chunk) → Future (대용량 dict 는 ParamSearchSharedPayload 사용).
|
||
|
||
완료 청크가 안 나오는 '대기 구간'마다(첫 청크 워밍업 포함) 하트비트를 찍어
|
||
화면이 멈춘 것처럼 보이지 않게 한다. (사람이 진행을 눈으로 확인하고 안심하도록)
|
||
- on_heartbeat(done_count, total, elapsed_sec, running) 콜백을 주면 그걸 호출
|
||
(호출자가 진행 바·경과·마지막 ETA 를 같은 줄에 다시 그림 — 권장).
|
||
- 콜백이 없으면 내장 하트비트(진행 바 + 완료 수 + 경과)를 직접 출력.
|
||
터미널이면 ~1초마다, 파일(nohup)이면 heartbeat_sec(기본 10초)마다 갱신.
|
||
"""
|
||
if not chunks:
|
||
return
|
||
inflight = max_inflight if max_inflight is not None else param_search_max_inflight(max_workers)
|
||
if heartbeat_sec is None:
|
||
heartbeat_sec = get_env_float("PARAM_SEARCH_HEARTBEAT_SEC", 10.0)
|
||
use_cr = sys.stdout.isatty()
|
||
# 터미널은 자주(1초) 갱신해 진행 바가 살아있게, 파일은 너무 잦은 줄 방지로 길게.
|
||
beat_iv = 1.0 if use_cr else max(1.0, heartbeat_sec)
|
||
total = len(chunks)
|
||
pending: Dict[Future, Any] = {}
|
||
chunk_iter = iter(chunks)
|
||
|
||
def _fill() -> None:
|
||
while len(pending) < inflight:
|
||
try:
|
||
ch = next(chunk_iter)
|
||
except StopIteration:
|
||
break
|
||
pending[submit_fn(ch)] = ch
|
||
|
||
_fill()
|
||
start = time.time()
|
||
done_count = 0
|
||
while pending:
|
||
# 완료된 청크가 생길 때까지 beat_iv 단위로 끊어 대기 → 그 사이 하트비트 출력.
|
||
done, _ = wait(list(pending.keys()), timeout=beat_iv,
|
||
return_when=FIRST_COMPLETED)
|
||
if not done:
|
||
if heartbeat_sec <= 0:
|
||
continue
|
||
el = time.time() - start
|
||
running = len(pending)
|
||
if on_heartbeat is not None:
|
||
on_heartbeat(done_count, total, el, running)
|
||
else:
|
||
bar = ParamSearchProgressETA.render_bar(done_count / max(1, total))
|
||
if done_count == 0:
|
||
label = f"🔥 워밍업·첫 청크 진행 중(워커 {running})"
|
||
else:
|
||
label = f"⚙️ 처리 중(진행 {running})"
|
||
msg = (f"{label} {bar} {done_count}/{total} | "
|
||
f"경과 {ParamSearchProgressETA.format_elapsed(el)}")
|
||
if use_cr:
|
||
sys.stdout.write(f"\r{msg} ")
|
||
sys.stdout.flush()
|
||
else:
|
||
print(msg, flush=True)
|
||
continue
|
||
for fut in done:
|
||
pending.pop(fut, None)
|
||
done_count += 1
|
||
yield fut.result()
|
||
_fill()
|
||
|
||
|
||
def _is_pid_alive(pid: int) -> bool:
|
||
if pid <= 0:
|
||
return False
|
||
try:
|
||
os.kill(pid, 0)
|
||
except ProcessLookupError:
|
||
return False
|
||
except PermissionError:
|
||
return True
|
||
else:
|
||
return True
|
||
|
||
|
||
class ParamSearchRunLock:
|
||
"""동일 이름 파라서치 중복 실행 방지 (PID 락 파일)."""
|
||
|
||
def __init__(self, name: str) -> None:
|
||
self.name = name
|
||
lock_dir = os.path.join(os.path.expanduser("~"), ".kis_bot_locks")
|
||
os.makedirs(lock_dir, exist_ok=True)
|
||
self.path = os.path.join(lock_dir, f"{name}.lock")
|
||
self._acquired = False
|
||
|
||
def acquire(self) -> bool:
|
||
if os.path.isfile(self.path):
|
||
old_pid = 0
|
||
try:
|
||
with open(self.path, "r", encoding="utf-8") as f:
|
||
old_pid = int(f.read().strip().split()[0])
|
||
except (ValueError, OSError):
|
||
old_pid = 0
|
||
if _is_pid_alive(old_pid):
|
||
return False
|
||
try:
|
||
os.remove(self.path)
|
||
except OSError:
|
||
pass
|
||
|
||
try:
|
||
fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)
|
||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||
f.write(f"{os.getpid()}\n")
|
||
self._acquired = True
|
||
atexit.register(self.release)
|
||
return True
|
||
except FileExistsError:
|
||
return False
|
||
|
||
def release(self) -> None:
|
||
if not self._acquired:
|
||
return
|
||
try:
|
||
if os.path.isfile(self.path):
|
||
with open(self.path, "r", encoding="utf-8") as f:
|
||
content = f.read().strip()
|
||
if content.startswith(str(os.getpid())):
|
||
os.remove(self.path)
|
||
except OSError:
|
||
pass
|
||
self._acquired = False
|
||
|
||
def __enter__(self) -> "ParamSearchRunLock":
|
||
if not self.acquire():
|
||
raise RuntimeError(
|
||
f"이미 실행 중인 {self.name} 파라서치가 있습니다.\n"
|
||
f" ps -ef | grep {self.name}\n"
|
||
f" pkill -f '{self.name}' 또는 rm {self.path}"
|
||
)
|
||
return self
|
||
|
||
def __exit__(self, *_args) -> None:
|
||
self.release()
|
||
|
||
|
||
def try_acquire_run_lock(name: str) -> Optional[ParamSearchRunLock]:
|
||
"""락 획득 실패 시 None (main 에서 메시지 출력용)."""
|
||
lock = ParamSearchRunLock(name)
|
||
if lock.acquire():
|
||
return lock
|
||
return None
|
||
|
||
|
||
def assert_parent_alive() -> None:
|
||
"""
|
||
워커 청크 루프에서 호출.
|
||
부모가 죽어 init(PPID=1)에 붙었거나 예상 부모와 다르면 즉시 종료.
|
||
"""
|
||
ppid = os.getppid()
|
||
if ppid == 1:
|
||
os._exit(0)
|
||
if _EXPECTED_PPID > 1 and ppid != _EXPECTED_PPID:
|
||
os._exit(0)
|