faiss_index_bq_dataset.py

import argparse
import logging
import os
import pkgutil
import sys
from urllib.parse import urlsplit

import apache_beam as beam
import faiss
from apache_beam.options.pipeline_options import PipelineOptions


def parse_d6w_config(argv=None):
    """Parse d6w config.
    :param argv: d6w config
    :return: dictionary containing d6w config
    """

    parser = argparse.ArgumentParser(
        description="See https://docbird.twitter.biz/d6w/model.html for any parameters inherited from d6w job config"
    )
    parser.add_argument(
        "--job_name", dest="job_name", required=True, help="d6w attribute"
    )
    parser.add_argument(
        "--project", dest="project", required=True, help="d6w attribute"
    )
    parser.add_argument(
        "--staging_location",
        dest="staging_location",
        required=True,
        help="d6w attribute",
    )
    parser.add_argument(
        "--temp_location", dest="temp_location", required=True, help="d6w attribute"
    )
    parser.add_argument(
        "--output_location",
        dest="output_location",
        required=True,
        help="GCS bucket and path where resulting artifacts are uploaded",
    )
    parser.add_argument(
        "--service_account_email",
        dest="service_account_email",
        required=True,
        help="d6w attribute",
    )
    parser.add_argument(
        "--factory_string",
        dest="factory_string",
        required=False,
        help="FAISS factory string describing index to build. See https://github.com/facebookresearch/faiss/wiki/The-index-factory",
    )
    parser.add_argument(
        "--metric",
        dest="metric",
        required=True,
        help="Metric used to compute distance between embeddings. Valid values are 'l2', 'ip', 'l1', 'linf'",
    )
    parser.add_argument(
        "--use_gpu",
        dest="gpu",
        required=True,
        help="--use_gpu=yes if you want to use GPU during index building",
    )

    known_args, unknown_args = parser.parse_known_args(argv)
    d6w_config = vars(known_args)
    d6w_config["gpu"] = d6w_config["gpu"].lower() == "yes"
    d6w_config["metric"] = parse_metric(d6w_config)

    """
  WARNING: Currently, d6w (a Twitter tool used to deploy Dataflow jobs to GCP) and
  PipelineOptions.for_dataflow_runner (a helper method in twitter.ml.common.apache_beam) do not
  play nicely together. The helper method will overwrite some of the config specified in the d6w
  file using the defaults in https://sourcegraph.twitter.biz/git.twitter.biz/source/-/blob/src/python/twitter/ml/common/apache_beam/__init__.py?L24.'
  However, the d6w output message will still report that the config specified in the d6w file was used.
  """
    logging.warning(
        f"The following d6w config parameters will be overwritten by the defaults in "
        f"https://sourcegraph.twitter.biz/git.twitter.biz/source/-/blob/src/python/twitter/ml/common/apache_beam/__init__.py?L24\n"
        f"{str(unknown_args)}"
    )
    return d6w_config


def get_bq_query():
    """
    Query is expected to return rows with unique entityId
    """
    return pkgutil.get_data(__name__, "bq.sql").decode("utf-8")


def parse_metric(config):
    metric_str = config["metric"].lower()
    if metric_str == "l2":
        return faiss.METRIC_L2
    elif metric_str == "ip":
        return faiss.METRIC_INNER_PRODUCT
    elif metric_str == "l1":
        return faiss.METRIC_L1
    elif metric_str == "linf":
        return faiss.METRIC_Linf
    else:
        raise Exception(f"Unknown metric: {metric_str}")


def run_pipeline(argv=[]):
    config = parse_d6w_config(argv)
    argv_with_extras = argv
    if config["gpu"]:
        argv_with_extras.extend(["--experiments", "use_runner_v2"])
        argv_with_extras.extend(
            [
                "--experiments",
                "worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver",
            ]
        )
        argv_with_extras.extend(
            [
                "--worker_harness_container_image",
                "gcr.io/twttr-recos-ml-prod/dataflow-gpu/beam2_39_0_py3_7",
            ]
        )

    options = PipelineOptions(argv_with_extras)
    output_bucket_name = urlsplit(config["output_location"]).netloc

    with beam.Pipeline(options=options) as p:
        input_data = p | "Read from BigQuery" >> beam.io.ReadFromBigQuery(
            method=beam.io.ReadFromBigQuery.Method.DIRECT_READ,
            query=get_bq_query(),
            use_standard_sql=True,
        )

        index_built = input_data | "Build and upload index" >> beam.CombineGlobally(
            MergeAndBuildIndex(
                output_bucket_name,
                config["output_location"],
                config["factory_string"],
                config["metric"],
                config["gpu"],
            )
        )

        # Make linter happy
        index_built


class MergeAndBuildIndex(beam.CombineFn):
    def __init__(self, bucket_name, gcs_output_path, factory_string, metric, gpu):
        self.bucket_name = bucket_name
        self.gcs_output_path = gcs_output_path
        self.factory_string = factory_string
        self.metric = metric
        self.gpu = gpu

    def create_accumulator(self):
        return []

    def add_input(self, accumulator, element):
        accumulator.append(element)
        return accumulator

    def merge_accumulators(self, accumulators):
        merged = []
        for accum in accumulators:
            merged.extend(accum)
        return merged

    def extract_output(self, rows):
        # Reimports are needed on workers
        import glob
        import subprocess

        import faiss
        import numpy as np
        from google.cloud import storage

        client = storage.Client()
        bucket = client.get_bucket(self.bucket_name)

        logging.info("Building FAISS index")
        logging.info(f"There are {len(rows)} rows")

        ids = np.array([x["entityId"] for x in rows]).astype("long")
        embeds = np.array([x["embedding"] for x in rows]).astype("float32")
        dimensions = len(embeds[0])
        N = ids.shape[0]
        logging.info(f"There are {dimensions} dimensions")

        if self.factory_string is None:
            M = 48

            divideable_dimensions = (dimensions // M) * M
            if divideable_dimensions != dimensions:
                opq_prefix = f"OPQ{M}_{divideable_dimensions}"
            else:
                opq_prefix = f"OPQ{M}"

            clusters = N // 20
            self.factory_string = f"{opq_prefix},IVF{clusters},PQ{M}"

        logging.info(f"Factory string is {self.factory_string}, metric={self.metric}")

        if self.gpu:
            logging.info("Using GPU")

            res = faiss.StandardGpuResources()
            cpu_index = faiss.index_factory(
                dimensions, self.factory_string, self.metric
            )
            cpu_index = faiss.IndexIDMap(cpu_index)
            gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
            gpu_index.train(embeds)
            gpu_index.add_with_ids(embeds, ids)
            cpu_index = faiss.index_gpu_to_cpu(gpu_index)
        else:
            logging.info("Using CPU")

            cpu_index = faiss.index_factory(
                dimensions, self.factory_string, self.metric
            )
            cpu_index = faiss.IndexIDMap(cpu_index)
            cpu_index.train(embeds)
            cpu_index.add_with_ids(embeds, ids)

        logging.info("Built faiss index")

        local_path = "/indices"
        logging.info(f"Writing indices to local {local_path}")
        subprocess.run(f"mkdir -p {local_path}".strip().split())
        local_index_path = os.path.join(local_path, "result.index")

        faiss.write_index(cpu_index, local_index_path)
        logging.info(f"Done writing indices to local {local_path}")

        logging.info(f"Uploading to GCS with path {self.gcs_output_path}")
        assert os.path.isdir(local_path)
        for local_file in glob.glob(local_path + "/*"):
            remote_path = os.path.join(
                self.gcs_output_path.split("/")[-1], local_file[1 + len(local_path) :]
            )
            blob = bucket.blob(remote_path)
            blob.upload_from_filename(local_file)


if __name__ == "__main__":
    logging.getLogger().setLevel(logging.INFO)
    run_pipeline(sys.argv)