refactor: smart search queue (#17977)

This commit is contained in:
Jason Rasmussen
2025-04-29 17:44:28 -04:00
committed by GitHub
parent 038a82c4f1
commit 2e8a286540
5 changed files with 51 additions and 36 deletions

View File

@@ -1,11 +1,10 @@
import { SystemConfig } from 'src/config';
import { ImmichWorker, JobName, JobStatus } from 'src/enum';
import { WithoutProperty } from 'src/repositories/asset.repository';
import { SmartInfoService } from 'src/services/smart-info.service';
import { getCLIPModelInfo } from 'src/utils/misc';
import { assetStub } from 'test/fixtures/asset.stub';
import { systemConfigStub } from 'test/fixtures/system-config.stub';
import { newTestService, ServiceMocks } from 'test/utils';
import { makeStream, newTestService, ServiceMocks } from 'test/utils';
describe(SmartInfoService.name, () => {
let sut: SmartInfoService;
@@ -152,38 +151,31 @@ describe(SmartInfoService.name, () => {
await sut.handleQueueEncodeClip({});
expect(mocks.asset.getAll).not.toHaveBeenCalled();
expect(mocks.asset.getWithout).not.toHaveBeenCalled();
expect(mocks.search.setDimensionSize).not.toHaveBeenCalled();
});
it('should queue the assets without clip embeddings', async () => {
mocks.asset.getWithout.mockResolvedValue({
items: [assetStub.image],
hasNextPage: false,
});
mocks.assetJob.streamForEncodeClip.mockReturnValue(makeStream([assetStub.image]));
await sut.handleQueueEncodeClip({ force: false });
expect(mocks.job.queueAll).toHaveBeenCalledWith([
{ name: JobName.SMART_SEARCH, data: { id: assetStub.image.id } },
]);
expect(mocks.asset.getWithout).toHaveBeenCalledWith({ skip: 0, take: 1000 }, WithoutProperty.SMART_SEARCH);
expect(mocks.assetJob.streamForEncodeClip).toHaveBeenCalledWith(false);
expect(mocks.search.setDimensionSize).not.toHaveBeenCalled();
});
it('should queue all the assets', async () => {
mocks.asset.getAll.mockResolvedValue({
items: [assetStub.image],
hasNextPage: false,
});
mocks.assetJob.streamForEncodeClip.mockReturnValue(makeStream([assetStub.image]));
await sut.handleQueueEncodeClip({ force: true });
expect(mocks.job.queueAll).toHaveBeenCalledWith([
{ name: JobName.SMART_SEARCH, data: { id: assetStub.image.id } },
]);
expect(mocks.asset.getAll).toHaveBeenCalled();
expect(mocks.assetJob.streamForEncodeClip).toHaveBeenCalledWith(true);
expect(mocks.search.setDimensionSize).toHaveBeenCalledExactlyOnceWith(512);
});
});

View File

@@ -3,12 +3,10 @@ import { SystemConfig } from 'src/config';
import { JOBS_ASSET_PAGINATION_SIZE } from 'src/constants';
import { OnEvent, OnJob } from 'src/decorators';
import { DatabaseLock, ImmichWorker, JobName, JobStatus, QueueName } from 'src/enum';
import { WithoutProperty } from 'src/repositories/asset.repository';
import { ArgOf } from 'src/repositories/event.repository';
import { BaseService } from 'src/services/base.service';
import { JobOf } from 'src/types';
import { JobItem, JobOf } from 'src/types';
import { getCLIPModelInfo, isSmartSearchEnabled } from 'src/utils/misc';
import { usePagination } from 'src/utils/pagination';
@Injectable()
export class SmartInfoService extends BaseService {
@@ -79,18 +77,18 @@ export class SmartInfoService extends BaseService {
await this.searchRepository.setDimensionSize(dimSize);
}
const assetPagination = usePagination(JOBS_ASSET_PAGINATION_SIZE, (pagination) => {
return force
? this.assetRepository.getAll(pagination, { isVisible: true })
: this.assetRepository.getWithout(pagination, WithoutProperty.SMART_SEARCH);
});
for await (const assets of assetPagination) {
await this.jobRepository.queueAll(
assets.map((asset) => ({ name: JobName.SMART_SEARCH, data: { id: asset.id } })),
);
let queue: JobItem[] = [];
const assets = this.assetJobRepository.streamForEncodeClip(force);
for await (const asset of assets) {
queue.push({ name: JobName.SMART_SEARCH, data: { id: asset.id } });
if (queue.length >= JOBS_ASSET_PAGINATION_SIZE) {
await this.jobRepository.queueAll(queue);
queue = [];
}
}
await this.jobRepository.queueAll(queue);
return JobStatus.SUCCESS;
}