From 8daca59ca61dac8fd7f68f3135d69092c70e599f Mon Sep 17 00:00:00 2001 From: popkirby Date: Sat, 8 Jul 2023 07:27:26 +0900 Subject: [PATCH] perf(backend): use mutex for nsfw model loading (#11109) Co-authored-by: tamaina --- CHANGELOG.md | 1 + packages/backend/package.json | 1 + packages/backend/src/core/AiService.ts | 10 +++++++++- pnpm-lock.yaml | 9 +++++++++ 4 files changed, 20 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 288ab0796..434420522 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,7 @@ ### Server - JSON.parse の回数を削減することで、ストリーミングのパフォーマンスを向上しました +- nsfwjs のモデルロードを排他することで、重複ロードによってメモリ使用量が増加しないように - 連合の配送ジョブのパフォーマンスを向上(ロック機構の見直し、Redisキャッシュの活用) ## 13.13.2 diff --git a/packages/backend/package.json b/packages/backend/package.json index 288a8ea7b..87f7f9ddf 100644 --- a/packages/backend/package.json +++ b/packages/backend/package.json @@ -75,6 +75,7 @@ "accepts": "1.3.8", "ajv": "8.12.0", "archiver": "5.3.1", + "async-mutex": "^0.4.0", "autwh": "0.1.0", "bcryptjs": "2.4.3", "blurhash": "2.0.5", diff --git a/packages/backend/src/core/AiService.ts b/packages/backend/src/core/AiService.ts index 02501b832..c0596446d 100644 --- a/packages/backend/src/core/AiService.ts +++ b/packages/backend/src/core/AiService.ts @@ -4,6 +4,7 @@ import { dirname } from 'node:path'; import { Inject, Injectable } from '@nestjs/common'; import * as nsfw from 'nsfwjs'; import si from 'systeminformation'; +import { Mutex } from 'async-mutex'; import type { Config } from '@/config.js'; import { DI } from '@/di-symbols.js'; import { bindThis } from '@/decorators.js'; @@ -17,6 +18,7 @@ let isSupportedCpu: undefined | boolean = undefined; @Injectable() export class AiService { private model: nsfw.NSFWJS; + private modelLoadMutex: Mutex = new Mutex(); constructor( @Inject(DI.config) @@ -39,7 +41,13 @@ export class AiService { const tf = await import('@tensorflow/tfjs-node'); - if (this.model == null) this.model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 }); + if (this.model == null) { + await this.modelLoadMutex.runExclusive(async () => { + if (this.model == null) { + this.model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 }); + } + }); + } const buffer = await fs.promises.readFile(path); const image = await tf.node.decodeImage(buffer, 3) as any; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 69d25e807..0d99cede2 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -140,6 +140,9 @@ importers: archiver: specifier: 5.3.1 version: 5.3.1 + async-mutex: + specifier: ^0.4.0 + version: 0.4.0 autwh: specifier: 0.1.0 version: 0.1.0 @@ -9468,6 +9471,12 @@ packages: resolution: {integrity: sha512-csOlWGAcRFJaI6m+F2WKdnMKr4HhdhFVBk0H/QbJFMCr+uO2kwohwXQPxw/9OCxp05r5ghVBFSyioixx3gfkNQ==} dev: true + /async-mutex@0.4.0: + resolution: {integrity: sha512-eJFZ1YhRR8UN8eBLoNzcDPcy/jqjsg6I1AP+KvWQX80BqOSW1oJPJXDylPUEeMr2ZQvHgnQ//Lp6f3RQ1zI7HA==} + dependencies: + tslib: 2.6.0 + dev: false + /async-settle@1.0.0: resolution: {integrity: sha512-VPXfB4Vk49z1LHHodrEQ6Xf7W4gg1w0dAPROHngx7qgDjqmIQ+fXmwgGXTW/ITLai0YLSvWepJOP9EVpMnEAcw==} engines: {node: '>= 0.10'}