<commit_before>#!/usr/bin/env python3
import argparse
import asyncio
import json
import logging
import os.path
import sys
from collections import defaultdict
from enum import Enum
from typing import Any
from typing import DefaultDict
from typing import Dict
from typing import List
from typing import NamedTuple
from typing import Optional
from typing import Set
from typing import Tuple

import yaml
from kazoo.client import KazooClient
from kazoo.exceptions import NoNodeError

logger = logging.getLogger("check_orphans")

PREFIX = "/smartstack/global/"
CHUNK_SIZE = 50  # How many concurrent xinetd connections
DEFAULT_ZK_DISCOVERY_PATH = "/nail/etc/zookeeper_discovery/infrastructure/local.yaml"
DEFAULT_NERVE_XINETD_PORT = 8735


class ExitCode(Enum):
    OK = 0
    ORPHANS = 1
    COLLISIONS = 2


def get_zk_hosts(path: str) -> List[str]:
    with open(path) as f:
        x = yaml.load(f)
    return [f"{host}:{port}" for host, port in x]


SmartstackData = Dict[str, Dict[str, Any]]


def get_zk_data(blacklisted_services: Set[str]) -> SmartstackData:
    logger.info(f"using {DEFAULT_ZK_DISCOVERY_PATH} for zookeeper")
    zk_hosts = get_zk_hosts(DEFAULT_ZK_DISCOVERY_PATH)

    logger.debug(f"connecting to zk hosts {zk_hosts}")
    zk = KazooClient(hosts=zk_hosts)
    zk.start()

    logger.debug(f"pulling smartstack data from zookeeper")
    zk_data = {}
    services = zk.get_children(PREFIX)
    for service in services:
        if service in blacklisted_services:
            continue
        service_instances = zk.get_children(os.path.join(PREFIX, service))
        instances_data = {}
        for instance in service_instances:
            try:
                instance_node = zk.get(os.path.join(PREFIX, service, instance))<commit_msg>Fix `check_orphans.py` uses unsafe yaml loader<commit_after>#!/usr/bin/env python3
import argparse
import asyncio
import json
import logging
import os.path
import sys
from collections import defaultdict
from enum import Enum
from typing import Any
from typing import DefaultDict
from typing import Dict
from typing import List
from typing import NamedTuple
from typing import Optional
from typing import Set
from typing import Tuple

import yaml
from kazoo.client import KazooClient
from kazoo.exceptions import NoNodeError

logger = logging.getLogger("check_orphans")

PREFIX = "/smartstack/global/"
CHUNK_SIZE = 50  # How many concurrent xinetd connections
DEFAULT_ZK_DISCOVERY_PATH = "/nail/etc/zookeeper_discovery/infrastructure/local.yaml"
DEFAULT_NERVE_XINETD_PORT = 8735


class ExitCode(Enum):
    OK = 0
    ORPHANS = 1
    COLLISIONS = 2


def get_zk_hosts(path: str) -> List[str]:
    with open(path) as f:
        x = yaml.safe_load(f)
    return [f"{host}:{port}" for host, port in x]


SmartstackData = Dict[str, Dict[str, Any]]


def get_zk_data(blacklisted_services: Set[str]) -> SmartstackData:
    logger.info(f"using {DEFAULT_ZK_DISCOVERY_PATH} for zookeeper")
    zk_hosts = get_zk_hosts(DEFAULT_ZK_DISCOVERY_PATH)

    logger.debug(f"connecting to zk hosts {zk_hosts}")
    zk = KazooClient(hosts=zk_hosts)
    zk.start()

    logger.debug(f"pulling smartstack data from zookeeper")
    zk_data = {}
    services = zk.get_children(PREFIX)
    for service in services:
        if service in blacklisted_services:
            continue
        service_instances = zk.get_children(os.path.join(PREFIX, service))
        instances_data = {}
        for instance in service_instances:
            try:
                instance_node = zk.get(os.path.join(PREFIX, service, instance))
{'fraction_non_alphanumeric': 0.08131682611142375, 'fraction_numerical': 0.005627462014631401, 'mean_word_length': 3.9306518723994452, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3503230', 'n_tokens_mistral': 1178, 'n_tokens_neox': 1122, 'n_words': 316}