⭐openGauss資料庫原始碼解析系列文章—— AI查詢時間預測⭐

語言: CN / TW / HK

上一篇介紹了“8.5 指標採集、預測與異常檢測”的相關內容,本篇我們介紹“8.6 AI查詢時間預測”的相關精彩內容介紹。

8.6 AI查詢時間預測

在前面介紹過“慢SQL發現”特性,該特性的典型場景是新業務上線前的檢查,輸入源是提前採集到的SQL流水資料。慢SQL發現功能主要主要應用在多條SQL語句的批量檢查上,要求之前執行過SQL語句,因此給出的結果主要是定性的,在某些場景下可能難以滿足使用者對於評估精度的要求。
因此,為了彌補上述場景的不足,滿足使用者更精確的SQL時間預測需求,同時為AI優化器做鋪墊,實現了本章所述的功能。
由於實際業務場景具有複雜的特質,現有的資料庫靜態代價估計模型往往統計結果失準,從而選擇了一些執行計劃較差的路徑。因此,針對上述複雜場景,需要資料庫的代價估計模型具備自我更新的能力。本特性主要功能為基於查詢語句的歷史資料,對當前執行的SQL語句進行查詢耗時和基數的估算。

8.6.1 使用場景

AI查詢分析的前提是需要獲取執行計劃。首先需要根據使用者需求在查詢執行時收集複雜查詢實際查詢計劃(包括計劃結構、運算元型別、相關資料來源、過濾條件等)、各運算元節點實際執行時間、優化器估算代價、實際返回行數、優化器估算行數、SMP併發執行緒數、等資訊。將其記錄在資料表中,並進行持久化管理包括定期進行資料失效清理。
本功能主要分為兩個方面,一個是行數估算,一個是查詢預測,前者是後者預測好壞的前提。目前openGauss基於線上學習對執行計劃各層的結果集大小進行估算,僅起到展示作用,並未影響到執行計劃的生成。後續可幫助優化器更準確地進行結果集估算,從而獲取更優的執行計劃。
當前階段本需求會提供系統函式來進行預測,並加入到explain中進行實際比較驗證。

8.6.2 現有技術

當前學術界在AI4DB領域,對基於機器學習的行數估算和查詢時延預測有許多嘗試。

1. 傳統方法

正如資料庫優化器專家Guy Lohman在部落格Is query optimization a “solved” problem中所說,傳統資料庫查詢效能預測的“阿喀琉斯之踵”便是中間結果集大小的估算。對於行數估算傳統基於統計資訊行數估算方法主要基於三類假設。
(1) 資料獨立分佈假設。
(2) 均勻分佈假設。
(3) 主外來鍵假設。
而實際場景中資料往往存在一定的相關性和傾斜性,此時上述假設可能會被打破,導致傳統資料庫優化器在多表連線中間結果集大小估算中可能會存在數個數量級的誤差。
2000年以來,以基於取樣的估算、基於取樣的核密度函式估算、基於多列直方圖為代表的統計學方法被提出,用於解決資料相關性帶來的估算問題。然而這些方法都存在一個共性問題,就是模型無法進行增量維護,而收集這些額外的統計資訊會增加巨大的資料庫維護開銷,雖然在一些特定的問題場景(如多列Range條件選擇率)取得了很大的準確率提升,但並沒有被各大資料庫廠商廣泛採用。
傳統效能預測方法主要依賴代價模型,在以下幾個方面存在明顯劣勢。
(1) 準確性:隨著底層硬體架構和優化技術不斷演進,實際效能預測模型的複雜度遠不可以用線性模型來建模。
(2) 可擴充套件性:代價模型的開發成本較高,不能面面俱到地對使用者具體場景進行優化。
(3) 可校準性:代價模型靈活性僅侷限於各資源維度線性相加時使用的係數,以及部分懲罰代價,靈活性較差,使用者實際使用時難以校準。
(4) 時效性:代價模型依賴統計資訊的收集和使用,目前缺乏增量維護方法,導致資料流動性較大的場景下統計資訊長期處於失效狀態。

2. 機器學習方法

機器學習模型在模型複雜度、可校準性、可增量維護性幾個維度的優勢能夠彌補傳統優化器代價模型的不足,基於機器學習的查詢效能預測逐漸成為資料庫學術界和產業界的主流研究方向之一。
除前文8.3節慢SQL發現部分介紹過相關方法外,清華大學的Learned Cost Estimator模型基於Multi-task Learning和字元條件的Word-Embedding方法進一步提升了預測準確率。
至此,機器學習方法雖然從實驗效果上看達到了較高的準確率,但現實業務場景持續性的資料分佈變化對模型的線上學習能力提出了要求。openGauss採用了資料驅動的線上學習模式,通過核心不斷收集歷史作業效能資訊,並在AI Engine側使用了R-LSTM(recursive long short term memory,遞迴長短期記憶網路)模型對運算元級查詢時延和中間結果集大小進行預測。

8.6.3 實現原理

在這裡插入圖片描述

圖8-15 AI查詢效能預測架構示意圖

總體而言,查詢效能預測由資料庫核心側和AI Engine側兩個部分組成,如圖8-15所示。
(1) 資料庫核心側除提供資料庫基本功能外還需要對歷史資料進行收集和持久化管理,並通過curl向AI Engine側傳送HTTPS請求。
(2) AI Engine提供模型訓練、執行預測、模型管理等介面,基於Flask框架的服務端接受HTTPS請求,該流程如圖8-16所示。
在這裡插入圖片描述

圖 8-16 資料庫核心和AI Engine程序關係示意圖

開啟資料收集相關引數後(其對效能可能有5%左右的影響,取決於實際業務負載情況),歷史性能資料被持久化收集在資料庫的系統表中,用於模型的訓練。
模型訓練之前,使用者需要對模型引數進行配置(詳見8.6.5使用示例)。使用者訓練指令下發之後,核心程序會向AI Engine側傳送configure請求,用於初始化機器學習模型。configure流程時序如圖8-17所示。
在這裡插入圖片描述

圖8-17 configure流程時序圖

模型配置成功後,核心程序向AI Engine側傳送train請求,觸發訓練,該流程如圖8-18所示。
在這裡插入圖片描述

圖8-18 train流程時序圖 模型訓練之後,使用者下發預測指令,資料庫會先向AI Engine側傳送setup請求,用於模型載入,載入成功後傳送predict請求得到預測結果,如圖8-19所示。 圖8-19 模型預測完整流程時序圖,分為setup和predict兩個階段 本特性架構上支援多模型,目前已實現R- LSTM模型,該模型架構如圖8-20所示。 計劃中,運算元間的執行順序也會影響運算元的效能。基於這種特性,我們使用了LSTM神經網路模型來學習計劃中運算元間這種有意義的依賴關係,並根據行數/時間預測的場景對模型的結構、損失函式、優化演算法等方面進行鍼對性的優化,提高此場景下學習和預測的準確率。 輸入:查詢計劃樹,各節點上的運算元型別,對應表名列名以及過濾條件。 輸出:行數、startup time、total time、Peak Memory。 在編碼(encoding)階段,每個計劃節點(plan node)被編碼成固定長度,連線成序列作為輸入LSTM神經網路的特徵值。 LSTM具有多個重複神經網路模組組成的鏈式網路,在每個模組中都有三個函式來決定歷史時序中的哪些資訊將被傳遞到下一個時序的網路模組中。最後一個模組的輸出值h_t即為模型返回的預測結果。

在這裡插入圖片描述
在這裡插入圖片描述
其中,Хt是當前時序模組的輸入,ht﹣₁是當前時序模組的輸入,是前一個時序的輸出資訊,使用sigmoid(σ)函式得到當前細胞狀態中將要輸出的部分Οt;Ct表示所有歷史時序保留的資訊,通過tanh函式處理後與當前狀態輸出資訊Οt相乘得到此狀態的輸出ht,將具有三個元素的一維向量 [startup time, total time, cardinality] 的預測結果同真實資料進行比較,使用ratio-error計算模型的損失函式。
在這裡插入圖片描述

圖8-20 模型架構圖

8.6.4 關鍵原始碼解析

1. 專案結構

AI Engine側涉及的主要檔案路徑為openGauss-server/src/gausskernel/dbmind/tools/predictor,其檔案結構如表8-13所示。

表8-13 AI Engine檔案結構

檔案結構

說明

install

部署所需檔案路徑

install/ca_ext.txt

證書配置檔案

install/requirements-gpu.txt

使用GPU(graphics processing unit,圖形處理器)訓練依賴庫列表

install/requirements.txt

使用CPU訓練依賴庫列表

install/ssl.sh

證書生成指令碼

python

專案程式碼路徑

python/certs.py

加密通訊

python/e_log

系統日誌路徑

python/log

模型訓練日誌路徑

python/log.conf

配置檔案

python/model.py

機器學習模型

python/run.py

服務端主函式

python/saved_models

模型訓練checkpoint

python/settings.py

工程配置檔案

python/uploads

Curl傳輸的檔案存放路徑

核心側主要涉及的檔案路徑為openGauss-server/src/gausskernel/optimizer/util/learn,其檔案結構如表8-14所示。

表8-14 核心端主要檔案結構

檔案結構

說明

comm.cpp

通訊層程式碼實現

encoding.cpp

資料編碼

ml_model.cpp

通用模型呼叫介面

plan_tree_model.cpp

樹狀模型呼叫介面

2. 訓練流程

核心側的模型訓練介面通過ModelTrainInternal函式實現,該函式的關鍵部分如下:

static void ModelTrainInternal(const char* templateName, const char* modelName, ModelAccuracy** mAcc)
{
  …
    /* 對於樹形模型呼叫對應的訓練介面 */
    char* trainResultJson = TreeModelTrain(modelinfo, labels);
    /* 解析返回結果 */
    …
    ModelTrainInfo* info = GetModelTrainInfo(jsonObj);
    cJSON_Delete(jsonObj);
    /* 更新模型資訊 */
    Relation modelRel = heap_open(OptModelRelationId, RowExclusiveLock);
   …
    UpdateTrainRes(values, datumsMax, datumsAcc, nLabel, mAcc, info, labels);

    HeapTuple modelTuple = SearchSysCache1(OPTMODEL, CStringGetDatum(modelName));
   …
    HeapTuple newTuple = heap_modify_tuple(modelTuple, RelationGetDescr(modelRel), values, nulls, replaces);
    simple_heap_update(modelRel, &newTuple->t_self, newTuple);
CatalogUpdateIndexes(modelRel, newTuple);
…
}

核心側的樹狀模型訓練介面通過TreeModelTrain函式實現,核心程式碼如下:

char* TreeModelTrain(Form_gs_opt_model modelinfo, char* labels)
{
    char* filename = (char*)palloc0(sizeof(char) * MAX_LEN_TEXT);
    char* buf = NULL;
    /* configure階段 */
    ConfigureModel(modelinfo, labels, &filename);

    /* 將編碼好的資料寫入臨時檔案 */
    SaveDataToFile(filename);

    /* Train階段 */
    buf = TrainModel(modelinfo, filename);
    return buf;
}

AI Engine側配置的Web服務的URI是/configure,訓練階段的URI是/train.下面的程式碼段展示了訓練過程。

  def fit(self, filename):
        keras.backend.clear_session()
        set_session(self.session)
        with self.graph.as_default():
            # 根據模型入參和出參維度變化情況,判斷是否需要初始化模型
            feature, label, need_init = self.parse(filename) 
            os.environ['CUDA_VISIBLE_DEVICES'] = '0'
            epsilon = self.model_info.make_epsilon()
            if need_init: # 冷啟動訓練
                epoch_start = 0
                self.model = self._build_model(epsilon)
            else: # 增量訓練
                epoch_start = int(self.model_info.last_epoch)
                ratio_error = ratio_error_loss_wrapper(epsilon)
                ratio_acc_2 = ratio_error_acc_wrapper(epsilon, 2)
                self.model = load_model(self.model_info.model_path,
                                        custom_objects={'ratio_error': ratio_error, 'ratio_acc': ratio_acc_2})
            self.model_info.last_epoch = int(self.model_info.max_epoch) + epoch_start
            self.model_info.dump_dict()
            log_path = os.path.join(settings.PATH_LOG, self.model_info.model_name + '_log.json')
            if not os.path.exists(log_path):
                os.mknod(log_path, mode=0o600)
            # 訓練日誌記錄回撥函式
            json_logging_callback = LossHistory(log_path, self.model_info.model_name, self.model_info.last_epoch)
            # 資料分割
            X_train, X_val, y_train, y_val = \
                train_test_split(feature, label, test_size=0.1)
            # 模型訓練
            self.model.fit(X_train, y_train, epochs=self.model_info.last_epoch,
                           batch_size=int(self.model_info.batch_size), validation_data=(X_val, y_val),
                           verbose=0, initial_epoch=epoch_start, callbacks=[json_logging_callback])
            # 記錄模型checkpoint
            self.model.save(self.model_info.model_path)
            val_pred = self.model.predict(X_val)
            val_re = get_ratio_errors_general(val_pred, y_val, epsilon)
            self.model_logger.debug(val_re)
            del self.model
            return val_re

3. 預測流程

核心側的模型預測過程主要通過ModelPredictInternal函式實現。樹狀模型預測過程通過TreeModelPredict函式實現。核心側的樹狀模型預測過程會佔用一些與AI Engine進行通訊的信令,該通訊過程如下:

char* TreeModelPredict(const char* modelName, char* filepath, const char* ip, int port)
{
    …
    if (!TryConnectRemoteServer(conninfo, &buf)) {
        DestroyConnInfo(conninfo);
        ParseResBuf(buf, filepath, "AI engine connection failed.");
        return buf;
    }

    switch (buf[0]) {
        case '0': {
            ereport(NOTICE, (errmodule(MOD_OPT_AI), errmsg("Model setup successfully.")));
            break;
        }
        case 'M': {
            ParseResBuf(buf, filepath, "Internal error: missing compulsory key.");
            break;
        }
…
    }
    /* Predict階段 */
    …
    if (!TryConnectRemoteServer(conninfo, &buf)) {
        ParseResBuf(buf, filepath, "AI engine connection failed.");
        return buf;
    }
    switch (buf[0]) {
        case 'M': {
            ParseResBuf(buf, filepath, "Internal error: fail to load the file to predict.");
            break;
        }
        case 'S': {
            ParseResBuf(buf, filepath, "Internal error: session is not loaded, model setup required.");
            break;
        }
        default: {
            break;
        }
    }
    return buf;
}

AI Engine側的Setup過程的Web介面是/model_setup,預測階段的Web介面是/predict,他們的協議都是Post。

4. 資料編碼

資料編碼分為以下兩個維度。
(1) 運算元維度:包括每個執行計劃運算元的屬性,如表8-15所示。

表8-15 運算元維度

屬性名

含義

編碼策略

Optname

運算元型別

One-hot

Orientation

返回元組儲存格式

One-hot

Strategy

邏輯屬性

One-hot

Options

物理屬性

One-hot

Quals

謂詞

hash

Projection

返回投影列

hash

(2) 計劃維度。
對於每個運算元,在其固有屬性之外,openGauss還對query id,plan node id和parent node id進行了記錄,在訓練/預測階段,使用這些資訊將運算元資訊重建為樹狀計劃結構,且可以遞迴構建子計劃樹來進行資料增強,從而提升模型泛化能力。樹狀資料結構如圖8-21所示。
在這裡插入圖片描述

圖8-21 樹狀資料結構示意圖 核心側的樹狀資料編碼通過GetOPTEncoding函式實現。

5. 模型結構

AI Engine的模型解析、訓練和預測見8.6.4章節,下面的程式碼展示了模型的結構。

class RnnModel():
    def _build_model(self, epsilon):
        model = Sequential()
        model.add(LSTM(units=int(self.model_info.hidden_units), return_sequences=True, input_shape=(None, int(self.model_info.feature_length))))
        model.add(LSTM(units=int(self.model_info.hidden_units), return_sequences=False))
        model.add(Dense(units=int(self.model_info.hidden_units), activation='relu'))
        model.add(Dense(units=int(self.model_info.hidden_units), activation='relu'))
        model.add(Dense(units=int(self.model_info.label_length), activation='sigmoid'))
        optimizer = keras.optimizers.Adadelta(lr=float(self.model_info.learning_rate), rho=0.95)
        ratio_error = ratio_error_loss_wrapper(epsilon)
        ratio_acc_2 = ratio_error_acc_wrapper(epsilon, 2)
        model.compile(loss=ratio_error, metrics=[ratio_acc_2], optimizer=optimizer)
        return model

AI Engine的損失函式使用ratio error(部分文獻中使用qerror代稱),該損失函式相較於MRE和MSE的優勢在於其能夠等價地懲罰高估和低估兩種情況,公式為:
在這裡插入圖片描述
ε宣告為效能預測值的無窮小值,防止分母為0的情況發生。

8.6.5 使用示例

AI查詢時間預測功能使用示例如下。
① 定義效能預測模型,程式碼如下:

INSERT INTO gs_opt_model VALUES(‘rlstm’, ‘model_name’, ‘host_ip’, ‘port’);

② 通過GUC引數開啟資料收集,配置的引數列表,程式碼如下:

enable_resource_track = on;
enable_resource_record = on;

③ 編碼訓練資料,程式碼如下:

SELECT gather_encoding_info('db_name');

④ 校準模型,程式碼如下:

SELECT model_train_opt('template_name', 'model_name');

⑤ 監控訓練狀態,程式碼如下:

SELECT track_train_process('host_ip', 'port');

⑥ 通過explain + SQL語句來預測SQL查詢的效能,程式碼如下:

EXPLAIN (..., predictor 'model_name') SELECT ...

獲得結果,其中,“p-time”列為標籤預測值。

Row Adapter  (cost=110481.35..110481.35 rows=100 p-time=99..182 width=100) (actual time=375.158..375.160 rows=2 loops=1)

8.6.6 演進路線

目前模型的泛化能力不足,依賴外接的AI Engine元件,且深度學習網路比較重,這會為部署造成困難;模型需要資料進行訓練,冷啟動階段的銜接不夠順暢,後續從以下幾個方面演進。
(1) 加入不同複雜度模型,並支援多模型融合分析,提供更健壯的模型預測結果和置信度。
(2) AI Engine考慮加入任務佇列,目前僅支援單併發預測/訓練,可以考慮建立多個服務端進行併發業務。
(3) 基於線上學習/遷移學習的增強,考慮對損失函式加入錨定懲罰代價來避免災難遺忘問題,同時優化資料管理模式,考慮data score機制,根據資料時效性賦權。
(4) 將本功能與優化器深度結合,探索基於AI的路徑選擇方法。

感謝大家學習第8章 AI技術中“8.6 AI查詢時間預測”的精彩內容,下一篇我們開啟“8.7 DeepSQL”的相關內容的介紹。
敬請期待。

「其他文章」