Компьютеры быстры, но вы этого не знаете

Люди чертовски плохо разбираются в величинах, особенно в тех, которые не могут воспринимать биологически. Например, мы интуитивно понимаем, насколько тяжелее предмет массой 10 кг предмета массой 1 кг.

Ощущение величин можно улучшить, преобразовав их каким-нибудь образом в сигналы, знакомые мозгу.

Смотрели ли вы эти видео?

1. Сравнение размера Вселенной в 3D
2. Измеряем богатство Джеффа Безоса в рисе

Второе нравится мне больше всего. Ежедневно я съедаю по чашке риса, так что измеряю состояние Джеффа не только наглядно, но и своим желудком.

Совсем недавно я внёс несколько оптимизаций в код, что помогло мне интуитивно понять, насколько быстро может работать компьютер. И я решил, что этим нужно поделиться.

Что мы оптимизируем?

Функция выглядит примерно так:

# учтите, что это сокращённый вид настоящей функции
def aggregate(input_df):
    weights = initialize_weights(len(input_df)) 

    # input_df содержит три столбца: timestamp, score и id
    output_df = group_by_id(input_df) 

    # выполняется сортировка в группе
    output_df = sort_by_time(output_df) 
    output_df = sort_by_score(output_df, desc=True)

    # выполняется ранжирование в группе
    output_df = rank_within_id(output_df)

    # добавляется столбец "results"
    output_df = multiply_weights(output_df, weights) 
    results = results_sum_in_each_group(output_df)

    return results # список

Это функция агрегации показателей, используемая одним из моих сервисов машинного обучения (Machine Learning, ML). Эта функция критически важна для того, чтобы модель работала. Нужно учесть, что вывод этой функции должен быть вычислен менее чем за 500 мс, чтобы она хотя бы попала в продакшен (где нам нужно менее чем за 500 мс выполнить 400-800 вызовов этой функции). Меня попросили её оптимизировать.

Давайте измерим. насколько быстро работает эта функция при 1000 вызовах (предположим, что каждые входящие данные имеют максимум 10 элементов).

Для выполнения 1000 вызовов понадобилось около 8 секунд. Очень плохо. Я хотел попробовать сделать 1 миллион вызовов, но для этого бы понадобилось более 20 минут. Посмотрим, как можно улучшить ситуацию.

Оптимизация 1: пишем алгоритм без Pandas + тривиальные улучшения алгоритма

Библиотека Pandas языка Python отлично подходит для экспериментов с данными, но ужасна в продакшене. Если оказалось, что вы используете её в продакшене, то настало время повзрослеть. [Примечание: не поймите меня неверно. Pandas достаточно быстра для типичного массива данных, но не для обработки, которая в моём случае замедляла pandas. Медленным может быть сам процесс создания объектов Pandas. Если вам нужно сгруппировать + отсортировать 1 миллион строк, то Pandas будет быстрее, чем чистый Python. Но если вашему сервису нужно реагировать менее чем за 500 мс, то вы прочувствуете на себе влияние каждой строки кода Pandas (и любые излишние траты, накладываемые абстрагированием).] Я заменил Pandas простыми списками Python и вручную реализовал алгоритм группировки и сортировки.

Кроме того, в каждом вызове функция вычисления весов инициализировалась заново, но единственное, что она делала — инициализировала ту же последовательность весов для какого-то размера массива. Это оказалось бонусом для меня и я начал вычислять веса предварительно.

Вот как выглядела функция:

WEIGHTS = initialize_weights(99999) 
def aggregate_efficient(input_lists):
    global WEIGHTS

    # input_lists содержит 2 списка
    output_lists = algorithm_wizardry(input_lists) 

    # добавляем столбец "results"
    output_lists = multiply_weights(output_lists, WEIGHTS) 
    results = results_sum_in_each_group(output_lists)

    return results # список
aggregate_efficient(ip_list)

И теперь 1 миллион вызовов занимает вот столько:

Мы снизили время с 20 минут до 12 секунд! Скорость приблизительно увеличилась на 9900%.

Этого достаточно, чтобы функция попала в продакшен. Но зачем останавливаться на этом?

Оптимизация 2: улучшаем функции с помощью Cython

Один из простейших трюков для ускорения функции на Python заключается в том, чтобы просто написать её на Cython. Вот как это сделать:

1. Создаём функцию .pyx
   Чем больше cdef мы сможем в неё поместить, тем лучше оптимизация.

   def aggregate_efficient_cyth(double[:] score_array, 
                           double[:] time_array, 
                           double[:] weights): 
   
       results = algorithm_wizardry(score_array, 
                                    time_array, weights)
       return results

2. Определяем файл setup.py

   Добавляем флаги компилятора, чтобы всё работало быстро. Некоторые говорят, что флаг -O3 опасен, но мы пойдём таким путём.

   from distutils.core import setup
   from Cython.Build import cythonize
   from distutils.extension import Extension
   from Cython.Distutils import build_ext
   
   
   ext_modules = [
           Extension("agg_cython",
               ["agg_cython.pyx"],
               libraries=["m"],
               extra_compile_args = ["-O3", "-ffast-math", "-march=native", "-fopenmp" ],
               extra_link_args=['-fopenmp'],
               language="c++")
           ]
   setup(name="agg_cyth_pure",cmdclass = {'build_ext': build_ext}, ext_modules=ext_modules,)

И вот как быстро выполняется 1 миллион вызовов:

Около 6,59 секунды, то есть увеличение скорости примерно на 82%. Нам получилось почти вдвое сократить время по сравнению с предыдущей оптимизацией. Потрясающе!

Но на этом мы не остановимся. Давайте начнём творить полное безумие.

Оптимизация 3: пишем функцию на чистом C++

Вот тут-то и начинается веселье. Это один из самых важных навыков, добавленных мной в свой технический инструментарий (благодаря Рагсу), и если производительность очень важна, то вам поможет это:

1. Реализуем функцию на чистом C++

   #include "agg_cyth_fast.hpp"
   
   using namespace std;
   
   double agg_efficient(double score_array[], 
                       long time_array[],
                       double weight_lookup[],
                       int N){
   
       vector<double> results = algorithm_wizardry(score_array,
                                      time_array, weight_lookup, N);
       return results;
   }

2. Подготовим файл заголовка

   #ifndef AGG_H
   #define AGG_H
   
   #include <iostream>
   #include <map>
   #include <vector>
   #include <algorithm>
   
   double agg_efficient(double[], long[], double[], int);
   #endif

3. Напишем файл .pyx для общения с C++

   import numpy as np
   from math import exp 
   from libc.math cimport exp as c_exp
   from cython.parallel import prange
   cimport cython
   
   
   cdef extern from "agg_cyth_fast.hpp" nogil:
       double agg_efficient(double[], long[], double[], int)
   
   def agg_efficient_fs(double[:] score_array, 
                         long[:] time_array,  
                          double[:] weight_lookup):
       cdef int N = len(time_array)
       cdef double Y;
       Y = agg_efficient(&score_array[0], 
                         &time_array[0], 
                          &weight_lookup[0], N);
       return Y

Далее мы используем файл setup.py, похожий на файл из предыдущей оптимизации, и соберём нашу функцию. По сути, это позволит нам передавать объекты Python функции на чистом C++.

Вот какой получается скорость 1 миллиона вызовов:

Чистый C++ может быть безумно быстрым. Мы ускорили вычисления примерно на 119%!

Но мы ещё не закончили. Часть кода файла .pyx в этом разделе сокрыта, но внимательный наблюдатель может догадаться о следующей оптимизации по импортам.

Оптимизация 4: настало время уничтожить Global Interpreter Lock (GIL)

В Python есть раздражающая вещь под названием GIL, которая не позволяет выполнять многопоточный код на Python быстрее, чем однопоточный.

Если уж вы хотите по-настоящему напрячь компьютер, то зачем использовать только одно ядро?

Теперь мы используем prange для параллелизации наших вычислений.

Просто добавим следующий код в файл .pyx:

@cython.boundscheck(False)
@cython.wraparound(False)
def agg_efficient_fs_batch(double[:,:] score_array, 
                            long[:,:] time_array, 
                            double[:] weight_lookup):
    cdef int M = len(score_array)
    cdef double[:] Y = np.zeros(M)
    cdef int i, N
    for i in prange(M, nogil=True):
        N = len(time_array[i])
        Y[i] = agg_efficient(&score_array[i][0], 
                             &time_array[i][0], 
                             &weight_lookup[0], N);
    return Y

Вот сколько нужно времени, чтобы выполнить параллельно тот же 1 миллион вызовов на машине с четырьмя ядрами:

Ускорение приблизительно на 237% по сравнению с предыдущей оптимизацией.

Подведём итог для 1 миллиона вызовов:

Думаю, именно настолько быстрыми могут быть компьютеры.