۷ تغییر برای ۸ برابر شدن سرعت همگام‌سازی دیتابیس جستجوی تصویری

در این پست فرایند افزایش سرعت ۸ برابری جاب بروزرسانی بردارهای موجود در Qdrant را بررسی می‌کنیم. به صورت خلاصه به موارد زیر می‌پردازیم:

• گام اول: افزایش سرعت پیش پردازش تصاویر
• گام دوم: افزایش سرعت مدل امبدینگ
• گام سوم: افزایش سرعت دانلود تصاویر
• گام چهارم: استفاده از Ray برای پردازش توزیع‌شده
• گام پنجم: پخش کردن دیپلویمنت‌های Ray بین GPUهای مختلف
• گام ششم: بهبود سرعت جابه‌جایی داده بین دیپلویمنت‌ها
• گام هفتم: دیپلوی کردن Qdrant برروی GPU

مقدمه

در ترب گزینه جستجو تصویری وجود دارد که با آپلود یک تصویر توسط کاربر، ابتدا بردار امبدینگ (Embedding Vector) آن توسط یک مدل پردازش تصویر ترنسفورمری (ViT) استخراج میشود و با مقایسه آن با بردارهای موجود در پایگاه داده برداری (Vector Database)، محصولات مشابه بازیابی و به کاربر نمایش داده می‌شوند.

این پایگاه داده برداری فقط به منظور جستجو تصویری استفاده می‌شود و از پایگاه داده اصلی جستجو ترب که مبتنی بر ElasticSearch است جدا می‌باشد، در نتیجه این نیاز وجود دارد که در یک فرآیند پیوسته در حال اجرا، اطلاعات موجود در پایگاه داده اصلی بررسی شوند و تغییرات payload (شامل اطلاعاتی که برای فیلتر کردن استفاده می‌شوند) و content (شامل اطلاعاتی که برای مشابهت‌یابی استفاده می‌شود، به صورت مشخص تصویر محصول) با استفاده از ریکوست‌های update payload یا update vector در پایگاه داده برداری اعمال شوند. پایگاه داده‌ برداری مورد استفاده در ترب Qdrant می‌باشد.

چرا بهینه‌سازی و افزایش سرعت؟

در ابتدای توسعه جستجو تصویری، به دلیل تمرکز بر پوشاک مدلی که برای تولید بردارهای امبدینگ استفاده شد مدل Marqo-FashionCLIP بود که یک نسخه تیون شده از OpenaiCLIP است و به صورت مشخص معماری vision tower این مدل ViT-B-16 میباشد که ۸۶.۵ میلیون پارامتر دارد و حجم آن ۳۳۰ مگابایت است که در مجموع مدل کوچکی به حساب می‌آید. اما برای افزایش کیفیت جستجو تصویری، بخصوص در دسته‌هایی غیر از پوشاک نیاز به مدل‌های بزرگتری بود. این در حالی است که ظرفیت جاب فعلی برابر با ایندکس کردن ۶۰ محصول در ثانیه (۵ میلیون محصول در روز) می‌باشد. مدلی که برای جایگزین کردن در نظر داشتیم مدل SigLIP2-so400m می‌باشد که ۵ برابر بزرگتر از مدل فعلی است. در صورت جایگزین کردن این مدل سرعت جاب فعلی پایین می‌آید و این کاهش سرعت عملا باعث بی‌مصرف شدن جاب بروزرسانی به دلیل کندی زیاد آن میشد. علاوه بر آن ظرفیت ایندکس ۶۰ محصول در ثانیه در مقیاس چند ده میلیونی سرعت پایینی به حساب می‌آید و سرعت توسعه و کیفیت جستجو را کاهش می‌دهد.

گام اول: افزایش سرعت پیش‌پردازش تصاویر

برای تولید امبدینگ تصاویر با مدل Marqo-FashionCLIP از فریمورک open-clip-torch استفاده می‌شود. این فریمورک در استفاده از مدل‌ها و آموزش‌ها آنها مفید است اما برای استفاده از پروداکشن احتمالا گزینه مناسبی نمی‌باشد. با پروفایل کردن زمان پردازش بخش‌های مختلف متوجه شدیم که نسبت زمان پیش‌پردازش و اجرای forward path مدل به صورت ۷۰ به ۳۰ می‌باشد. این مساله به این دلیل است که به صورت پیش‌فرض پیش‌پردازش تصاویر برروی cpu و به صورت ترتیبی اجرا می‌شود.

در ابتدا این پیش‌پردازش را با torchvision بازنویسی کرده و با قرار دادن آن در یک nn.module و استفاده از gpu و پردازش به صورت بچ افزایش سرعت حدودا ۹۰ برابری را تجربه کردیم. مدتی بعد گزینه use_fast در نسخه ۴.۴۹.۰ کتابخانه transformers متعلق به HuggingFace در imageProcessorها نیز قرار گرفت که پیاده‌سازی مشابهی داشت و افزایش سرعت مشابهی را داشت. با مقایسه سه روش به جدول زیر می‌رسیم:

و در نهایت از پیش‌پردازش‌کننده Fast استفاده می‌کنیم.

گام دوم: افزایش سرعت مدل امبدینگ

برای اجرای سریعتر مدل از روش‌های مختلف quantization می‌توان استفاده کرد. فریمورک onnx یک شروع خوب برای فرآیند تسریع مدل می‌باشد اما یک گزینه در دسترس و نسبتا راحت استفاده از TensorRT است. با استفاده از trtexec برای تبدیل مدل به engine با دقت fp16 و استفاده از Polygraphy برای اجرا مدل در بچ‌سایز ۳۲ زمان اجرا نسبت به حالت پیش‌فرض (fp32 و pytorch) حدود ۴.۳ برابر کاهش می‌یابد.

گام سوم: افزایش سرعت دانلود تصاویر

در بخش مربوط به دانلود تصاویر، به استفاده از aiohttp یک سشن باز کرده و به صورت بچ چندین تصویر به صورت موازی از سرویس تصاویر ترب دانلود و با استفاده از PIL و BytesIO تصویر با فرمت PIL آماده و به بخش پیش‌پردازش وارد می‌شود. سرویس تصاویر ترب قابلیت تعیین اندازه‌ی تصویر دریافتی را دارد که تا کنون با اندازه ۶۴۰ در ۶۴۰ تصاویر دریافت می‌شدند.

برای افزایش سرعت این بخش تغییر اندازه تصاویر در مبدا اتفاق افتاد. یعنی سرویس تصاویر ترب به درخواست کلاینت دانلود، تصاویر را در سایز ۳۸۴ در ۳۸۴ ارسال میکرد. این تغییر باعث افزایش سرعت ۳ برابری دانلود تصاویر شد.

برای تبدیل تصاویر دانلود شده از فرمت بایت به np.array می‌توان PIL را کنار گذاشت و از cv2 استفاده کرد. در جاب فعلی روش اولی به کار گرفته شده بود که با تغییر آن و استفاده از cv2 سرعت بخش دانلود ۲۵ درصد سریعتر شد. هر دوی این روش‌ها در فرآیند دانلود یک بچ از تصاویر به صورت موازی به کار گرفته و مقایسه شده‌اند.

گام چهارم: استفاده از Ray برای پردازش توزیع‌شده

در مرحله بعد، بخش‌های مختلف جاب، شامل دانلود تصاویر، پیش‌پردازش، تولید امبدینگ و قرار دادن در Qdrant به دیپلویمنت‌های مختلفی در یک اپلیکیشن ray serve تبدیل شدند. این تغییر به ما اجازه می‌داد تا هر مرحله‌ای که bottleneck باشد را با افزایش تعداد replica با بخش‌های دیگر هم‌سرعت کنیم.

در ترب برای سرو مدل‌های ماشین لرنینگ از ray serve استفاده می‌کنیم و بردن بخش های مختلف جاب باعث شد سربار رفت و برگشت داده، بخصوص داده حجیم امبدینگ برای بچ‌سایزهای بزرگ، کاهش پیدا کند و سربار serialization و ارسال از طریق http حذف گردد. Ray برای جابه‌جایی داده بین پراسس‌ها از shared memory داخلی استفاده می‌کند. با استفاده از این روش و اجرا فرآیند به صورت موازی در دو کارت گرافیک، سرعت نسبت به حالت غیرموازی دو برابر شد.

گام پنجم: پخش کردن دیپلویمنت‌ها بین GPUهای مختلف

در هنگام تبدیل مدل‌های امبدینگ به engineهای TensorRT، میزان بچ‌سایز بیشینه این مدل‌ها را برروی ۳۲ ثابت کردیم تا میزان رم مصرفی آنها از کارت‌های گرافیک محدود شود. با اعمال این مساله میزان utilization مدل‌ها نیز محدود میشد اما زمانی که ۶ رپلیکا از مدل برروی یک کارت گرافیک قرار می‌گرفت عملا سرعت پردازش از قرار دادن ۲ رپلیکا کمتر می‌شد و دلیل آن ظرفیت محدود پردازشی کارت گرافیک بود. به کمک ویژگی اختصاص منابع مجازی در ray، رپلیکاهای مدل رو در سطح ۴ کارت گرافیک پخش کردیم تا میزان utilization محدودکننده سرعت نباشد. با این کار از باتل‌نک شدن توان پردازش کارت‌های گرافیک جلوگیری کردیم و در نتیجه مرحله امبد کردن تصاویر باتل‌نک نمی‌شود.

گام ششم: بهبود سرعت جابه‌جایی داده بین دیپلویمنت‌ها

حال که تمام جاب در ray قرار دارد یکی از مواردی که می‌تواند باعث افت سرعت شود، جابه‌جایی داده حجیم بین دیپلویمنت‌های ray است. مثلا انتقال torch.tensor یا PIL.Image و یا موارد دیگری که به صورت python native نیستند موجب ایجاد سربار serialization زیادی می‌شوند. تبدیل این موارد به فرمت‌های پایه پایتون یا استفاده از np.array که توسط ray به خوبی پشتیبانی می‌شود باعث افزایش سرعت ۱۰ درصدی شد.

گام هفتم: دیپلوی کردن Qdrant برروی GPU

پایگاه‌داده برداری Qdrant در حالت پیش‌فرض و تا قبل از نسخه ۱.۱۳ فقط از cpu استفاده می‌کرد و فرآیند ساخت ایندکس HNSW و هم فرآیند جستجو در آن و سایر پردازش‌ها همه برروی cpu انجام می‌شدند. فرآیند ساخت ایندکس یک فرآیند سنگین می‌باشد که با اضافه شدن بردارهای جدید باید مجددا اجرا شود و از آن جایی که داده‌های ترب پیوسته در حالت تغییر هستند و در نتیجه بردارها باید بروز باشند این فرآیند همواره در حال اجرا است و بخش زیادی از تعداد هسته موجود را به خود اختصاص می‌دهد و عموما موجب کند شدن درخواست‌های جستجو نیز می‌شود.

در نسخه ۱.۱۳ قابلیت استفاده از GPU برای ساخت ایندکس اضافه شد که به کمک این قابلیت می‌شود سنگین‌ترین پردازش این پایگاه‌داده را بر عهده gpu گذاشت که منجر به افزایش سرعت ۱۰ برابری در ساخت ایندکس می‌شود و cpu نیز در این فرایند درگیر نخواهد شد و کاملا به درخواست‌های جستجو اختصاص پیدا می‌کند.

مورد دیگری که در افزایش سرعت جاب موثر بود و به Qdrant مرتبط می‌شود این است که در زمان ارسال درخواست upsert یا سایر درخواست‌ها به Qdrant می‌توان با قرار دادن wait=False صرفا بعد از دریافت شدن داده توسط پایگاه‌داده و قرار گرفتن در WAL پاسخ درخواست داده شود و به این صورت سنگینی پردازش ساخت ایندکس و قرار دادن داده جدید در پایگاه دانش یا سایر فرآیندهایی write موجب کند شدن فرآیند اصلی جاب نشوند.

جمع‌بندی

با اعمال گام‌های بالا در جاب جدید افزایش سرعت نزدیک به ۸ برابری را تجربه کردیم که به ما این امکان را می‌داد که در زمان بسیار کوتاه‌تری محصولات را مجددا آپدیت کنیم تا تغییرات جدید را بررسی و تست کنیم. همچنین هدف اصلی این تغییرات ایجاد قابلیت استفاده از مدل‌های بزرگتر بود. بعد از این افزایش سرعت توانستیم مدل اولیه یعنی Marqo-FashionCLIP را با سرعت ایندکس ۴۶۰ محصول در ثانیه و SigLIP2-so400m را با سرعت ایندکس ۱۶۰ محصول در ثانیه و به لطف Ray و پردازش موازی، قرار دادن هر دو بردار در ایندکس محصولات را با سرعت ۱۳۰ محصول در ثانیه داشته باشیم. این افزایش سرعت منجر شد در فرصت ۲ روزه بتوانیم مدل جدید را تست کنیم و با افزایش ۲۰ درصدی در شاخص‌های رضایت کاربران مواجه شویم.