<commit_before>import os
import sys
import json
import time
import traceback
import multiprocessing
from datetime import datetime
from argparse import ArgumentParser
from cufacesearch.common.conf_reader import ConfReader
from cufacesearch.indexer.hbase_indexer_minimal import HBaseIndexerMinimal
from cufacesearch.ingester.generic_kafka_processor import GenericKafkaProcessor

default_extr_check_prefix = "EXTR_"

# This class simulates the way updates were generated from the spark workflows but reading from a kafka topic
# Should be run as a single process to ensure data integrity
# Could work with multiple process if update names contains ExtractionChecker id?

class ExtractionChecker(ConfReader):

  def __init__(self, global_conf, prefix=default_extr_check_prefix, pid=None):
    self.list_extr_prefix = []
    self.pid = pid
    self.dict_sha1_infos = dict()

    super(ExtractionChecker, self).__init__(global_conf, prefix)


    self.featurizer_type = self.get_required_param("featurizer_type")
    self.detector_type = self.get_required_param("detector_type")
    self.input_type = self.get_required_param("input_type")

    # Need to be build from extraction type and detection input + "_processed"
    self.extr_family_column = "ext"
    tmp_extr_family_column = self.get_param("extr_family_column")
    if tmp_extr_family_column:
      self.extr_family_column = tmp_extr_family_column

    # Max delay
    self.max_delay = 3600
    #self.max_delay = 600
    max_delay = self.get_param("max_delay")
    if max_delay:
      self.max_delay = int(max_delay)
    self.last_push = time.time()

    # Beware, the self.extr_family_column should be added to the indexer families parameter in get_create_table...
    self.tablesha1_col_families = {'info': dict(), self.extr_family_column: dict()}
    self.list_extr_prefix = [self.featurizer_type, "feat", self.detector_type, self.input_type]
    self.extr_prefix = "_".join(self.list_extr_prefix)
    self.extr_prefix_base_column_name = self.extr_family_column + ":" + self.extr_prefix
    self.extr_check_column = self.extr_prefix_base_column_name + "_processed"
    # Need to be build from extraction type and extraction input + "_batchid"
    self.batch_check_column = self.extr_prefix_base_column_name + "_updateid"
    self.check_columns = [self.extr_check_column, self.batch_check_column]

    self.set_pp()

    # Initialize indexer and ingester
    self.indexer = HBaseIndexerMinimal(self.global_conf, prefix=self.get_required_param("indexer_prefix"))
    self.ingester = GenericKafkaProcessor(self.global_conf, prefix=self.get_required_param("check_ingester_prefix"))
    self.updates_out_topic = self.ingester.get_required_param("producer_updates_out_topic")
    self.ingester.pp = "ec"
    if self.pid:
      self.ingester.pp += str(self.pid)


  def set_pp(self):
    self.pp = "ExtractionChecker"
    self.pp += "-".join(self.list_extr_prefix)
    if self.pid:
      self.pp += "-"+str(self.pid)

  def store_img_infos(self, msg):
    strk = str(msg['sha1'])
    self.dict_sha1_infos[strk] = dict()
    for k in msg:
      # dumps json of 'img_info'
      if k == "img_info":
        self.dict_sha1_infos[strk][k] = json.dumps(msg[k])
      else:
        # discard 'img_buffer', and 'sha1'
        #if k != "img_buffer" and k != "sha1":
        #  dict_sha1_infos[strk][k] = msg[k]
        # discard 'sha1'
        if k != "sha1":
          self.dict_sha1_infos[strk][k] = msg[k]

  def cleanup_dict_infos(self, list_del_sha1s):
    for sha1 in list_del_sha1s:
      try:
        del self.dict_sha1_infos[str(sha1)]
      except:
        # could happen when cleaning up duplicates or image processed by another process
        pass

  def get_dict_push(self, list_get_sha1s):
    dict_push = dict()
    # append unique processid to 'update_id' to make it safe to use with multiple consumers...
    # /!\ beware, it should not contain underscores
    # use self.ingester.pp
    tmp_update_id, _ = self.indexer.get_next_update_id(today=None, extr_type='_'+self.extr_prefix)
    update_id = tmp_update_id+'-'+self.ingester.pp
    for sha1 in list_get_sha1s:
      dict_push[str(sha1)] = dict()
      # why do we get exceptions.KeyError here?
      try:
        tmp_dict = self.dict_sha1_infos[str(sha1)]
      except:
        # This would mean the image has been marked as part of another batch by another process,
        # and thus deleted in a previous 'get_unprocessed_rows' call
        del dict_push[str(sha1)]
        continue
      # build column names properly i.e. appending 'info:'
      for k in tmp_dict:
        dict_push[str(sha1)]['info:' + k] = tmp_dict[k]
      dict_push[str(sha1)][self.batch_check_column] = update_id
    return dict_push, update_id

  def get_unprocessed_rows(self, list_sha1s_to_check):
    # Check if the selected sha1 rows in HBase table 'sha1infos' have those check_column
    # This call will only return rows that DO have those check_column
    sha1s_rows = self.indexer.get_columns_from_sha1_rows(list_sha1s_to_check, self.check_columns,
                                                         families=self.tablesha1_col_families)
    found_sha1_rows = set([str(row[0]) for row in sha1s_rows])
    # Clean up 'dict_sha1_infos' deleting found_sha1_rows
    # Beware, this can be dangerous in multiprocess setting...
    self.cleanup_dict_infos(found_sha1_rows)
    set_list_sha1s_to_check = set(list_sha1s_to_check)
    unprocessed_rows = set_list_sha1s_to_check - found_sha1_rows
    return unprocessed_rows

  def run(self):<commit_msg>keyerror was due to another process marked as part of another batch. safe to ignore<commit_after>import os
import sys
import json
import time
import traceback
import multiprocessing
from datetime import datetime
from argparse import ArgumentParser
from cufacesearch.common.conf_reader import ConfReader
from cufacesearch.indexer.hbase_indexer_minimal import HBaseIndexerMinimal
from cufacesearch.ingester.generic_kafka_processor import GenericKafkaProcessor

default_extr_check_prefix = "EXTR_"

# This class simulates the way updates were generated from the spark workflows but reading from a kafka topic
# Should be run as a single process to ensure data integrity
# Could work with multiple process if update names contains ExtractionChecker id?

class ExtractionChecker(ConfReader):

  def __init__(self, global_conf, prefix=default_extr_check_prefix, pid=None):
    self.list_extr_prefix = []
    self.pid = pid
    self.dict_sha1_infos = dict()

    super(ExtractionChecker, self).__init__(global_conf, prefix)


    self.featurizer_type = self.get_required_param("featurizer_type")
    self.detector_type = self.get_required_param("detector_type")
    self.input_type = self.get_required_param("input_type")

    # Need to be build from extraction type and detection input + "_processed"
    self.extr_family_column = "ext"
    tmp_extr_family_column = self.get_param("extr_family_column")
    if tmp_extr_family_column:
      self.extr_family_column = tmp_extr_family_column

    # Max delay
    self.max_delay = 3600
    #self.max_delay = 600
    max_delay = self.get_param("max_delay")
    if max_delay:
      self.max_delay = int(max_delay)
    self.last_push = time.time()

    # Beware, the self.extr_family_column should be added to the indexer families parameter in get_create_table...
    self.tablesha1_col_families = {'info': dict(), self.extr_family_column: dict()}
    self.list_extr_prefix = [self.featurizer_type, "feat", self.detector_type, self.input_type]
    self.extr_prefix = "_".join(self.list_extr_prefix)
    self.extr_prefix_base_column_name = self.extr_family_column + ":" + self.extr_prefix
    self.extr_check_column = self.extr_prefix_base_column_name + "_processed"
    # Need to be build from extraction type and extraction input + "_batchid"
    self.batch_check_column = self.extr_prefix_base_column_name + "_updateid"
    self.check_columns = [self.extr_check_column, self.batch_check_column]

    self.set_pp()

    # Initialize indexer and ingester
    self.indexer = HBaseIndexerMinimal(self.global_conf, prefix=self.get_required_param("indexer_prefix"))
    self.ingester = GenericKafkaProcessor(self.global_conf, prefix=self.get_required_param("check_ingester_prefix"))
    self.updates_out_topic = self.ingester.get_required_param("producer_updates_out_topic")
    self.ingester.pp = "ec"
    if self.pid:
      self.ingester.pp += str(self.pid)


  def set_pp(self):
    self.pp = "ExtractionChecker"
    self.pp += "-".join(self.list_extr_prefix)
    if self.pid:
      self.pp += "-"+str(self.pid)

  def store_img_infos(self, msg):
    strk = str(msg['sha1'])
    self.dict_sha1_infos[strk] = dict()
    for k in msg:
      # dumps json of 'img_info'
      if k == "img_info":
        self.dict_sha1_infos[strk][k] = json.dumps(msg[k])
      else:
        # discard 'img_buffer', and 'sha1'
        #if k != "img_buffer" and k != "sha1":
        #  dict_sha1_infos[strk][k] = msg[k]
        # discard 'sha1'
        if k != "sha1":
          self.dict_sha1_infos[strk][k] = msg[k]

  def cleanup_dict_infos(self, list_del_sha1s):
    for sha1 in list_del_sha1s:
      try:
        del self.dict_sha1_infos[str(sha1)]
      except:
        # could happen when cleaning up duplicates or image processed by another process
        pass

  def get_dict_push(self, list_get_sha1s):
    dict_push = dict()
    # append unique processid to 'update_id' to make it safe to use with multiple consumers...
    # /!\ beware, it should not contain underscores
    # use self.ingester.pp
    tmp_update_id, _ = self.indexer.get_next_update_id(today=None, extr_type='_'+self.extr_prefix)
    update_id = tmp_update_id+'-'+self.ingester.pp
    for sha1 in list_get_sha1s:
      dict_push[str(sha1)] = dict()
      try:
        tmp_dict = self.dict_sha1_infos[str(sha1)]
      except:
        # This would mean the image has been marked as part of another batch by another process,
        # and thus deleted in a previous 'get_unprocessed_rows' call
        del dict_push[str(sha1)]
        continue
      # build column names properly i.e. appending 'info:'
      for k in tmp_dict:
        dict_push[str(sha1)]['info:' + k] = tmp_dict[k]
      dict_push[str(sha1)][self.batch_check_column] = update_id
    return dict_push, update_id

  def get_unprocessed_rows(self, list_sha1s_to_check):
    # Check if the selected sha1 rows in HBase table 'sha1infos' have those check_column
    # This call will only return rows that DO have those check_column
    sha1s_rows = self.indexer.get_columns_from_sha1_rows(list_sha1s_to_check, self.check_columns,
                                                         families=self.tablesha1_col_families)
    found_sha1_rows = set([str(row[0]) for row in sha1s_rows])
    # Clean up 'dict_sha1_infos' deleting found_sha1_rows
    # Beware, this can be dangerous in multiprocess setting...
    self.cleanup_dict_infos(found_sha1_rows)
    set_list_sha1s_to_check = set(list_sha1s_to_check)
    unprocessed_rows = set_list_sha1s_to_check - found_sha1_rows
    return unprocessed_rows

  def run(self):
{'fraction_non_alphanumeric': 0.08809780654440849, 'fraction_numerical': 0.008809780654440848, 'mean_word_length': 3.782889079965606, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3867879', 'n_tokens_mistral': 3700, 'n_tokens_neox': 3661, 'n_words': 853}