145 lines
5.7 KiB
Python
145 lines
5.7 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from typing import List
|
|
|
|
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
if REPO_ROOT not in sys.path:
|
|
sys.path.insert(0, REPO_ROOT)
|
|
|
|
from analytics import AnalyticsRepository, AnalyticsService, _is_model_traffic_line, _parse_traffic_size
|
|
from config import ANALYTICS_TRAFFIC_ROOT, MODEL_TRAFFIC_THRESHOLD, MONITORING_DB_PATH
|
|
|
|
|
|
def _split_packages(raw: str) -> List[str]:
|
|
return [item.strip() for item in str(raw or "").split(",") if item.strip()]
|
|
|
|
|
|
def _load_all_packages(service: AnalyticsService, tag: str = "", limit: int = 0) -> List[str]:
|
|
normalized_tag = str(tag or "").strip()
|
|
with service.repo._connect() as connection:
|
|
summaries = service.repo._list_catalog_summaries(connection)
|
|
packages = []
|
|
for item in sorted(summaries, key=lambda value: str(value.get("package_name") or "")):
|
|
if normalized_tag and normalized_tag not in (item.get("incremental_batch_tags") or []):
|
|
continue
|
|
package_name = str(item.get("package_name") or "").strip()
|
|
if package_name:
|
|
packages.append(package_name)
|
|
return packages[:limit] if limit > 0 else packages
|
|
|
|
|
|
def _update_model_stats(repo: AnalyticsRepository, package_name: str, flow_count: int, traffic_bytes: int) -> bool:
|
|
with repo._write_lock, repo._connect() as connection:
|
|
latest = connection.execute(
|
|
"""
|
|
SELECT package_name, batch_tag, run_kind, attempt
|
|
FROM collection_task
|
|
WHERE package_name = ?
|
|
ORDER BY
|
|
datetime(COALESCE(NULLIF(completed_at, ''), NULLIF(started_at, ''), created_at)) DESC,
|
|
attempt DESC,
|
|
batch_tag DESC,
|
|
run_kind DESC
|
|
LIMIT 1
|
|
""",
|
|
(package_name,),
|
|
).fetchone()
|
|
if not latest:
|
|
return False
|
|
connection.execute(
|
|
"""
|
|
UPDATE collection_task
|
|
SET model_flow_count = ?,
|
|
model_traffic_bytes = ?
|
|
WHERE package_name = ?
|
|
AND batch_tag = ?
|
|
AND run_kind = ?
|
|
AND attempt = ?
|
|
""",
|
|
(
|
|
flow_count,
|
|
traffic_bytes,
|
|
latest["package_name"],
|
|
latest["batch_tag"],
|
|
latest["run_kind"],
|
|
latest["attempt"],
|
|
),
|
|
)
|
|
return True
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Backfill model traffic stats for all or selected packages."
|
|
)
|
|
parser.add_argument("--db-path", default=MONITORING_DB_PATH)
|
|
parser.add_argument("--traffic-root", default=ANALYTICS_TRAFFIC_ROOT)
|
|
parser.add_argument("--packages", default="", help="Comma-separated package names. Defaults to all packages.")
|
|
parser.add_argument("--tag", default="", help="Only process packages whose incremental_batch_tag contains this tag.")
|
|
parser.add_argument("--limit", type=int, default=0)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
parser.add_argument("--auto-promote", action="store_true",
|
|
help="Auto-set model_eligible=1 and collection_status=pending if model_flow_count > 0.")
|
|
args = parser.parse_args()
|
|
|
|
scanned = 0
|
|
matched = 0
|
|
backfilled = 0
|
|
promoted = 0
|
|
missing_task = 0
|
|
service = AnalyticsService(
|
|
db_path=args.db_path,
|
|
traffic_root=args.traffic_root,
|
|
start_worker=False,
|
|
)
|
|
try:
|
|
packages = _split_packages(args.packages) or _load_all_packages(service, args.tag, args.limit)
|
|
|
|
for package_name in packages:
|
|
scanned += 1
|
|
traffic_files = service._find_traffic_files(package_name)
|
|
if not traffic_files:
|
|
continue
|
|
model_flow_count = 0
|
|
model_traffic_bytes = 0
|
|
for file_path in traffic_files:
|
|
try:
|
|
with open(file_path, "r", encoding="utf-8", errors="ignore") as handle:
|
|
for raw_line in handle:
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
fields = [item.strip() for item in line.split(",")]
|
|
if not fields or fields[0].strip() != package_name:
|
|
continue
|
|
if _is_model_traffic_line(fields):
|
|
model_flow_count += 1
|
|
model_traffic_bytes += _parse_traffic_size(fields[6])
|
|
except OSError:
|
|
continue
|
|
if model_flow_count == 0:
|
|
continue
|
|
matched += 1
|
|
eligible = 1 if 0 < model_flow_count <= MODEL_TRAFFIC_THRESHOLD else 0
|
|
if args.dry_run:
|
|
print(f"DRY-RUN {package_name}: model_flow={model_flow_count} model_bytes={model_traffic_bytes} eligible={eligible}")
|
|
continue
|
|
if not _update_model_stats(service.repo, package_name, model_flow_count, model_traffic_bytes):
|
|
missing_task += 1
|
|
continue
|
|
backfilled += 1
|
|
if args.auto_promote and eligible:
|
|
promoted += service.set_packages_pending([package_name], reason="model_traffic_backfill")
|
|
print(f"backfilled {package_name}: model_flow={model_flow_count} model_bytes={model_traffic_bytes} eligible={eligible}")
|
|
finally:
|
|
service.close()
|
|
|
|
print(f"done scanned={scanned} matched={matched} backfilled={backfilled} promoted={promoted} missing_task={missing_task}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|