refactor(server): jobs (#2023)

* refactor: job to domain

* chore: regenerate open api

* chore: tests

* fix: missing breaks

* fix: get asset with missing exif data

---------

Co-authored-by: Alex Tran <alex.tran1502@gmail.com>
This commit is contained in:
Jason Rasmussen
2023-03-20 11:55:28 -04:00
committed by GitHub
parent db6b14361d
commit 386eef046d
68 changed files with 1355 additions and 907 deletions

View File

@@ -1,4 +1,12 @@
import { BadRequestException } from '@nestjs/common';
export * from './upload_location.constant';
export const MACHINE_LEARNING_URL = process.env.IMMICH_MACHINE_LEARNING_URL || 'http://immich-machine-learning:3003';
export const MACHINE_LEARNING_ENABLED = MACHINE_LEARNING_URL !== 'false';
export function assertMachineLearningEnabled() {
if (!MACHINE_LEARNING_ENABLED) {
throw new BadRequestException('Machine learning is not enabled.');
}
}

View File

@@ -2,12 +2,22 @@ import { AssetEntity, AssetType } from '@app/infra/db/entities';
export interface AssetSearchOptions {
isVisible?: boolean;
type?: AssetType;
}
export enum WithoutProperty {
THUMBNAIL = 'thumbnail',
ENCODED_VIDEO = 'encoded-video',
EXIF = 'exif',
CLIP_ENCODING = 'clip-embedding',
OBJECT_TAGS = 'object-tags',
}
export const IAssetRepository = 'IAssetRepository';
export interface IAssetRepository {
getByIds(ids: string[]): Promise<AssetEntity[]>;
getWithout(property: WithoutProperty): Promise<AssetEntity[]>;
deleteAll(ownerId: string): Promise<void>;
getAll(options?: AssetSearchOptions): Promise<AssetEntity[]>;
save(asset: Partial<AssetEntity>): Promise<AssetEntity>;

View File

@@ -3,6 +3,7 @@ import { APIKeyService } from './api-key';
import { AssetService } from './asset';
import { AuthService } from './auth';
import { DeviceInfoService } from './device-info';
import { JobService } from './job';
import { MediaService } from './media';
import { OAuthService } from './oauth';
import { SearchService } from './search';
@@ -18,6 +19,7 @@ const providers: Provider[] = [
APIKeyService,
AuthService,
DeviceInfoService,
JobService,
MediaService,
OAuthService,
SmartInfoService,

View File

@@ -18,3 +18,4 @@ export * from './system-config';
export * from './tag';
export * from './user';
export * from './user-token';
export * from './util';

View File

@@ -0,0 +1,2 @@
export * from './job-command.dto';
export * from './job-id.dto';

View File

@@ -0,0 +1,14 @@
import { ApiProperty } from '@nestjs/swagger';
import { IsBoolean, IsEnum, IsNotEmpty, IsOptional } from 'class-validator';
import { JobCommand } from '../job.constants';
export class JobCommandDto {
@IsNotEmpty()
@IsEnum(JobCommand)
@ApiProperty({ type: 'string', enum: JobCommand, enumName: 'JobCommand' })
command!: JobCommand;
@IsOptional()
@IsBoolean()
force!: boolean;
}

View File

@@ -0,0 +1,10 @@
import { ApiProperty } from '@nestjs/swagger';
import { IsEnum, IsNotEmpty } from 'class-validator';
import { QueueName } from '../job.constants';
export class JobIdDto {
@IsNotEmpty()
@IsEnum(QueueName)
@ApiProperty({ type: String, enum: QueueName, enumName: 'JobName' })
jobId!: QueueName;
}

View File

@@ -1,3 +1,6 @@
export * from './dto';
export * from './job.constants';
export * from './job.interface';
export * from './job.repository';
export * from './job.service';
export * from './response-dto';

View File

@@ -2,32 +2,63 @@ export enum QueueName {
THUMBNAIL_GENERATION = 'thumbnail-generation-queue',
METADATA_EXTRACTION = 'metadata-extraction-queue',
VIDEO_CONVERSION = 'video-conversion-queue',
MACHINE_LEARNING = 'machine-learning-queue',
BACKGROUND_TASK = 'background-task',
OBJECT_TAGGING = 'object-tagging-queue',
CLIP_ENCODING = 'clip-encoding-queue',
BACKGROUND_TASK = 'background-task-queue',
STORAGE_TEMPLATE_MIGRATION = 'storage-template-migration-queue',
SEARCH = 'search-queue',
}
export enum JobCommand {
START = 'start',
PAUSE = 'pause',
EMPTY = 'empty',
}
export enum JobName {
// upload
ASSET_UPLOADED = 'asset-uploaded',
VIDEO_CONVERSION = 'mp4-conversion',
// conversion
QUEUE_VIDEO_CONVERSION = 'queue-video-conversion',
VIDEO_CONVERSION = 'video-conversion',
// thumbnails
QUEUE_GENERATE_THUMBNAILS = 'queue-generate-thumbnails',
GENERATE_JPEG_THUMBNAIL = 'generate-jpeg-thumbnail',
GENERATE_WEBP_THUMBNAIL = 'generate-webp-thumbnail',
// metadata
QUEUE_METADATA_EXTRACTION = 'queue-metadata-extraction',
EXIF_EXTRACTION = 'exif-extraction',
EXTRACT_VIDEO_METADATA = 'extract-video-metadata',
REVERSE_GEOCODING = 'reverse-geocoding',
// user deletion
USER_DELETION = 'user-deletion',
USER_DELETE_CHECK = 'user-delete-check',
// storage template
STORAGE_TEMPLATE_MIGRATION = 'storage-template-migration',
SYSTEM_CONFIG_CHANGE = 'system-config-change',
OBJECT_DETECTION = 'detect-object',
IMAGE_TAGGING = 'tag-image',
// object tagging
QUEUE_OBJECT_TAGGING = 'queue-object-tagging',
DETECT_OBJECTS = 'detect-objects',
CLASSIFY_IMAGE = 'classify-image',
// cleanup
DELETE_FILES = 'delete-files',
// search
SEARCH_INDEX_ASSETS = 'search-index-assets',
SEARCH_INDEX_ASSET = 'search-index-asset',
SEARCH_INDEX_ALBUMS = 'search-index-albums',
SEARCH_INDEX_ALBUM = 'search-index-album',
SEARCH_REMOVE_ALBUM = 'search-remove-album',
SEARCH_REMOVE_ASSET = 'search-remove-asset',
// clip
QUEUE_ENCODE_CLIP = 'queue-clip-encode',
ENCODE_CLIP = 'clip-encode',
}

View File

@@ -1,31 +1,35 @@
import { AlbumEntity, AssetEntity, UserEntity } from '@app/infra/db/entities';
export interface IAlbumJob {
export interface IBaseJob {
force?: boolean;
}
export interface IAlbumJob extends IBaseJob {
album: AlbumEntity;
}
export interface IAssetJob {
export interface IAssetJob extends IBaseJob {
asset: AssetEntity;
}
export interface IBulkEntityJob {
export interface IBulkEntityJob extends IBaseJob {
ids: string[];
}
export interface IAssetUploadedJob {
export interface IAssetUploadedJob extends IBaseJob {
asset: AssetEntity;
fileName: string;
}
export interface IDeleteFilesJob {
export interface IDeleteFilesJob extends IBaseJob {
files: Array<string | null | undefined>;
}
export interface IUserDeletionJob {
export interface IUserDeletionJob extends IBaseJob {
user: UserEntity;
}
export interface IReverseGeocodingJob {
export interface IReverseGeocodingJob extends IBaseJob {
assetId: string;
latitude: number;
longitude: number;

View File

@@ -2,6 +2,7 @@ import { JobName, QueueName } from './job.constants';
import {
IAssetJob,
IAssetUploadedJob,
IBaseJob,
IBulkEntityJob,
IDeleteFilesJob,
IReverseGeocodingJob,
@@ -17,21 +18,45 @@ export interface JobCounts {
}
export type JobItem =
// Asset Upload
| { name: JobName.ASSET_UPLOADED; data: IAssetUploadedJob }
// Transcoding
| { name: JobName.QUEUE_VIDEO_CONVERSION; data: IBaseJob }
| { name: JobName.VIDEO_CONVERSION; data: IAssetJob }
// Thumbnails
| { name: JobName.QUEUE_GENERATE_THUMBNAILS; data: IBaseJob }
| { name: JobName.GENERATE_JPEG_THUMBNAIL; data: IAssetJob }
| { name: JobName.GENERATE_WEBP_THUMBNAIL; data: IAssetJob }
| { name: JobName.EXIF_EXTRACTION; data: IAssetUploadedJob }
| { name: JobName.REVERSE_GEOCODING; data: IReverseGeocodingJob }
// User Deletion
| { name: JobName.USER_DELETE_CHECK }
| { name: JobName.USER_DELETION; data: IUserDeletionJob }
// Storage Template
| { name: JobName.STORAGE_TEMPLATE_MIGRATION }
| { name: JobName.SYSTEM_CONFIG_CHANGE }
// Metadata Extraction
| { name: JobName.QUEUE_METADATA_EXTRACTION; data: IBaseJob }
| { name: JobName.EXIF_EXTRACTION; data: IAssetUploadedJob }
| { name: JobName.EXTRACT_VIDEO_METADATA; data: IAssetUploadedJob }
| { name: JobName.OBJECT_DETECTION; data: IAssetJob }
| { name: JobName.IMAGE_TAGGING; data: IAssetJob }
| { name: JobName.REVERSE_GEOCODING; data: IReverseGeocodingJob }
// Object Tagging
| { name: JobName.QUEUE_OBJECT_TAGGING; data: IBaseJob }
| { name: JobName.DETECT_OBJECTS; data: IAssetJob }
| { name: JobName.CLASSIFY_IMAGE; data: IAssetJob }
// Clip Embedding
| { name: JobName.QUEUE_ENCODE_CLIP; data: IBaseJob }
| { name: JobName.ENCODE_CLIP; data: IAssetJob }
// Filesystem
| { name: JobName.DELETE_FILES; data: IDeleteFilesJob }
// Search
| { name: JobName.SEARCH_INDEX_ASSETS }
| { name: JobName.SEARCH_INDEX_ASSET; data: IBulkEntityJob }
| { name: JobName.SEARCH_INDEX_ALBUMS }
@@ -43,6 +68,7 @@ export const IJobRepository = 'IJobRepository';
export interface IJobRepository {
queue(item: JobItem): Promise<void>;
pause(name: QueueName): Promise<void>;
empty(name: QueueName): Promise<void>;
isActive(name: QueueName): Promise<boolean>;
getJobCounts(name: QueueName): Promise<JobCounts>;

View File

@@ -0,0 +1,170 @@
import { BadRequestException } from '@nestjs/common';
import { newJobRepositoryMock } from '../../test';
import { IJobRepository, JobCommand, JobName, JobService, QueueName } from '../job';
describe(JobService.name, () => {
let sut: JobService;
let jobMock: jest.Mocked<IJobRepository>;
beforeEach(async () => {
jobMock = newJobRepositoryMock();
sut = new JobService(jobMock);
});
it('should work', () => {
expect(sut).toBeDefined();
});
describe('getAllJobStatus', () => {
it('should get all job statuses', async () => {
jobMock.getJobCounts.mockResolvedValue({
active: 1,
completed: 1,
failed: 1,
delayed: 1,
waiting: 1,
});
await expect(sut.getAllJobsStatus()).resolves.toEqual({
'background-task-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
'clip-encoding-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
'metadata-extraction-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
'object-tagging-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
'search-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
'storage-template-migration-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
'thumbnail-generation-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
'video-conversion-queue': {
active: 1,
completed: 1,
delayed: 1,
failed: 1,
waiting: 1,
},
});
});
});
describe('handleCommand', () => {
it('should handle a pause command', async () => {
await sut.handleCommand(QueueName.METADATA_EXTRACTION, { command: JobCommand.PAUSE, force: false });
expect(jobMock.pause).toHaveBeenCalledWith(QueueName.METADATA_EXTRACTION);
});
it('should handle an empty command', async () => {
await sut.handleCommand(QueueName.METADATA_EXTRACTION, { command: JobCommand.EMPTY, force: false });
expect(jobMock.empty).toHaveBeenCalledWith(QueueName.METADATA_EXTRACTION);
});
it('should not start a job that is already running', async () => {
jobMock.isActive.mockResolvedValue(true);
await expect(
sut.handleCommand(QueueName.VIDEO_CONVERSION, { command: JobCommand.START, force: false }),
).rejects.toBeInstanceOf(BadRequestException);
expect(jobMock.queue).not.toHaveBeenCalled();
});
it('should handle a start video conversion command', async () => {
jobMock.isActive.mockResolvedValue(false);
await sut.handleCommand(QueueName.VIDEO_CONVERSION, { command: JobCommand.START, force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.QUEUE_VIDEO_CONVERSION, data: { force: false } });
});
it('should handle a start storage template migration command', async () => {
jobMock.isActive.mockResolvedValue(false);
await sut.handleCommand(QueueName.STORAGE_TEMPLATE_MIGRATION, { command: JobCommand.START, force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.STORAGE_TEMPLATE_MIGRATION });
});
it('should handle a start object tagging command', async () => {
jobMock.isActive.mockResolvedValue(false);
await sut.handleCommand(QueueName.OBJECT_TAGGING, { command: JobCommand.START, force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.QUEUE_OBJECT_TAGGING, data: { force: false } });
});
it('should handle a start clip encoding command', async () => {
jobMock.isActive.mockResolvedValue(false);
await sut.handleCommand(QueueName.CLIP_ENCODING, { command: JobCommand.START, force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.QUEUE_ENCODE_CLIP, data: { force: false } });
});
it('should handle a start metadata extraction command', async () => {
jobMock.isActive.mockResolvedValue(false);
await sut.handleCommand(QueueName.METADATA_EXTRACTION, { command: JobCommand.START, force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.QUEUE_METADATA_EXTRACTION, data: { force: false } });
});
it('should handle a start thumbnail generation command', async () => {
jobMock.isActive.mockResolvedValue(false);
await sut.handleCommand(QueueName.THUMBNAIL_GENERATION, { command: JobCommand.START, force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.QUEUE_GENERATE_THUMBNAILS, data: { force: false } });
});
it('should throw a bad request when an invalid queue is used', async () => {
jobMock.isActive.mockResolvedValue(false);
await expect(
sut.handleCommand(QueueName.BACKGROUND_TASK, { command: JobCommand.START, force: false }),
).rejects.toBeInstanceOf(BadRequestException);
expect(jobMock.queue).not.toHaveBeenCalled();
});
});
});

View File

@@ -0,0 +1,68 @@
import { assertMachineLearningEnabled } from '@app/common';
import { BadRequestException, Inject, Injectable, Logger } from '@nestjs/common';
import { JobCommandDto } from './dto';
import { JobCommand, JobName, QueueName } from './job.constants';
import { IJobRepository } from './job.repository';
import { AllJobStatusResponseDto } from './response-dto';
@Injectable()
export class JobService {
private logger = new Logger(JobService.name);
constructor(@Inject(IJobRepository) private jobRepository: IJobRepository) {}
handleCommand(queueName: QueueName, dto: JobCommandDto): Promise<void> {
this.logger.debug(`Handling command: queue=${queueName},force=${dto.force}`);
switch (dto.command) {
case JobCommand.START:
return this.start(queueName, dto);
case JobCommand.PAUSE:
return this.jobRepository.pause(queueName);
case JobCommand.EMPTY:
return this.jobRepository.empty(queueName);
}
}
async getAllJobsStatus(): Promise<AllJobStatusResponseDto> {
const response = new AllJobStatusResponseDto();
for (const queueName of Object.values(QueueName)) {
response[queueName] = await this.jobRepository.getJobCounts(queueName);
}
return response;
}
private async start(name: QueueName, { force }: JobCommandDto): Promise<void> {
const isActive = await this.jobRepository.isActive(name);
if (isActive) {
throw new BadRequestException(`Job is already running`);
}
switch (name) {
case QueueName.VIDEO_CONVERSION:
return this.jobRepository.queue({ name: JobName.QUEUE_VIDEO_CONVERSION, data: { force } });
case QueueName.STORAGE_TEMPLATE_MIGRATION:
return this.jobRepository.queue({ name: JobName.STORAGE_TEMPLATE_MIGRATION });
case QueueName.OBJECT_TAGGING:
assertMachineLearningEnabled();
return this.jobRepository.queue({ name: JobName.QUEUE_OBJECT_TAGGING, data: { force } });
case QueueName.CLIP_ENCODING:
assertMachineLearningEnabled();
return this.jobRepository.queue({ name: JobName.QUEUE_ENCODE_CLIP, data: { force } });
case QueueName.METADATA_EXTRACTION:
return this.jobRepository.queue({ name: JobName.QUEUE_METADATA_EXTRACTION, data: { force } });
case QueueName.THUMBNAIL_GENERATION:
return this.jobRepository.queue({ name: JobName.QUEUE_GENERATE_THUMBNAILS, data: { force } });
default:
throw new BadRequestException(`Invalid job name: ${name}`);
}
}
}

View File

@@ -0,0 +1,41 @@
import { ApiProperty } from '@nestjs/swagger';
import { QueueName } from '../job.constants';
export class JobCountsDto {
@ApiProperty({ type: 'integer' })
active!: number;
@ApiProperty({ type: 'integer' })
completed!: number;
@ApiProperty({ type: 'integer' })
failed!: number;
@ApiProperty({ type: 'integer' })
delayed!: number;
@ApiProperty({ type: 'integer' })
waiting!: number;
}
export class AllJobStatusResponseDto implements Record<QueueName, JobCountsDto> {
@ApiProperty({ type: JobCountsDto })
[QueueName.THUMBNAIL_GENERATION]!: JobCountsDto;
@ApiProperty({ type: JobCountsDto })
[QueueName.METADATA_EXTRACTION]!: JobCountsDto;
@ApiProperty({ type: JobCountsDto })
[QueueName.VIDEO_CONVERSION]!: JobCountsDto;
@ApiProperty({ type: JobCountsDto })
[QueueName.OBJECT_TAGGING]!: JobCountsDto;
@ApiProperty({ type: JobCountsDto })
[QueueName.CLIP_ENCODING]!: JobCountsDto;
@ApiProperty({ type: JobCountsDto })
[QueueName.STORAGE_TEMPLATE_MIGRATION]!: JobCountsDto;
@ApiProperty({ type: JobCountsDto })
[QueueName.BACKGROUND_TASK]!: JobCountsDto;
@ApiProperty({ type: JobCountsDto })
[QueueName.SEARCH]!: JobCountsDto;
}

View File

@@ -0,0 +1 @@
export * from './all-job-status-response.dto';

View File

@@ -3,9 +3,9 @@ import { AssetType } from '@app/infra/db/entities';
import { Inject, Injectable, Logger } from '@nestjs/common';
import { join } from 'path';
import sanitize from 'sanitize-filename';
import { IAssetRepository, mapAsset } from '../asset';
import { IAssetRepository, mapAsset, WithoutProperty } from '../asset';
import { CommunicationEvent, ICommunicationRepository } from '../communication';
import { IAssetJob, IJobRepository, JobName } from '../job';
import { IAssetJob, IBaseJob, IJobRepository, JobName } from '../job';
import { IStorageRepository } from '../storage';
import { IMediaRepository } from './media.repository';
@@ -21,6 +21,22 @@ export class MediaService {
@Inject(IStorageRepository) private storageRepository: IStorageRepository,
) {}
async handleQueueGenerateThumbnails(job: IBaseJob): Promise<void> {
try {
const { force } = job;
const assets = force
? await this.assetRepository.getAll()
: await this.assetRepository.getWithout(WithoutProperty.THUMBNAIL);
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.GENERATE_JPEG_THUMBNAIL, data: { asset } });
}
} catch (error: any) {
this.logger.error('Failed to queue generate thumbnail jobs', error.stack);
}
}
async handleGenerateJpegThumbnail(data: IAssetJob): Promise<void> {
const { asset } = data;
@@ -52,8 +68,8 @@ export class MediaService {
asset.resizePath = jpegThumbnailPath;
await this.jobRepository.queue({ name: JobName.GENERATE_WEBP_THUMBNAIL, data: { asset } });
await this.jobRepository.queue({ name: JobName.IMAGE_TAGGING, data: { asset } });
await this.jobRepository.queue({ name: JobName.OBJECT_DETECTION, data: { asset } });
await this.jobRepository.queue({ name: JobName.CLASSIFY_IMAGE, data: { asset } });
await this.jobRepository.queue({ name: JobName.DETECT_OBJECTS, data: { asset } });
await this.jobRepository.queue({ name: JobName.ENCODE_CLIP, data: { asset } });
this.communicationRepository.send(CommunicationEvent.UPLOAD_SUCCESS, asset.ownerId, mapAsset(asset));
@@ -71,8 +87,8 @@ export class MediaService {
asset.resizePath = jpegThumbnailPath;
await this.jobRepository.queue({ name: JobName.GENERATE_WEBP_THUMBNAIL, data: { asset } });
await this.jobRepository.queue({ name: JobName.IMAGE_TAGGING, data: { asset } });
await this.jobRepository.queue({ name: JobName.OBJECT_DETECTION, data: { asset } });
await this.jobRepository.queue({ name: JobName.CLASSIFY_IMAGE, data: { asset } });
await this.jobRepository.queue({ name: JobName.DETECT_OBJECTS, data: { asset } });
await this.jobRepository.queue({ name: JobName.ENCODE_CLIP, data: { asset } });
this.communicationRepository.send(CommunicationEvent.UPLOAD_SUCCESS, asset.ownerId, mapAsset(asset));

View File

@@ -5,7 +5,7 @@ export interface MachineLearningInput {
}
export interface IMachineLearningRepository {
tagImage(input: MachineLearningInput): Promise<string[]>;
classifyImage(input: MachineLearningInput): Promise<string[]>;
detectObjects(input: MachineLearningInput): Promise<string[]>;
encodeImage(input: MachineLearningInput): Promise<number[]>;
encodeText(input: string): Promise<number[]>;

View File

@@ -1,6 +1,13 @@
import { AssetEntity } from '@app/infra/db/entities';
import { newJobRepositoryMock, newMachineLearningRepositoryMock, newSmartInfoRepositoryMock } from '../../test';
import { IJobRepository } from '../job';
import {
assetEntityStub,
newAssetRepositoryMock,
newJobRepositoryMock,
newMachineLearningRepositoryMock,
newSmartInfoRepositoryMock,
} from '../../test';
import { IAssetRepository, WithoutProperty } from '../asset';
import { IJobRepository, JobName } from '../job';
import { IMachineLearningRepository } from './machine-learning.interface';
import { ISmartInfoRepository } from './smart-info.repository';
import { SmartInfoService } from './smart-info.service';
@@ -12,35 +19,63 @@ const asset = {
describe(SmartInfoService.name, () => {
let sut: SmartInfoService;
let assetMock: jest.Mocked<IAssetRepository>;
let jobMock: jest.Mocked<IJobRepository>;
let smartMock: jest.Mocked<ISmartInfoRepository>;
let machineMock: jest.Mocked<IMachineLearningRepository>;
beforeEach(async () => {
assetMock = newAssetRepositoryMock();
smartMock = newSmartInfoRepositoryMock();
jobMock = newJobRepositoryMock();
machineMock = newMachineLearningRepositoryMock();
sut = new SmartInfoService(jobMock, smartMock, machineMock);
sut = new SmartInfoService(assetMock, jobMock, smartMock, machineMock);
});
it('should work', () => {
expect(sut).toBeDefined();
});
describe('handleQueueObjectTagging', () => {
it('should queue the assets without tags', async () => {
assetMock.getWithout.mockResolvedValue([assetEntityStub.image]);
await sut.handleQueueObjectTagging({ force: false });
expect(jobMock.queue.mock.calls).toEqual([
[{ name: JobName.CLASSIFY_IMAGE, data: { asset: assetEntityStub.image } }],
[{ name: JobName.DETECT_OBJECTS, data: { asset: assetEntityStub.image } }],
]);
expect(assetMock.getWithout).toHaveBeenCalledWith(WithoutProperty.OBJECT_TAGS);
});
it('should queue all the assets', async () => {
assetMock.getAll.mockResolvedValue([assetEntityStub.image]);
await sut.handleQueueObjectTagging({ force: true });
expect(jobMock.queue.mock.calls).toEqual([
[{ name: JobName.CLASSIFY_IMAGE, data: { asset: assetEntityStub.image } }],
[{ name: JobName.DETECT_OBJECTS, data: { asset: assetEntityStub.image } }],
]);
expect(assetMock.getAll).toHaveBeenCalled();
});
});
describe('handleTagImage', () => {
it('should skip assets without a resize path', async () => {
await sut.handleTagImage({ asset: { resizePath: '' } as AssetEntity });
await sut.handleClassifyImage({ asset: { resizePath: '' } as AssetEntity });
expect(smartMock.upsert).not.toHaveBeenCalled();
expect(machineMock.tagImage).not.toHaveBeenCalled();
expect(machineMock.classifyImage).not.toHaveBeenCalled();
});
it('should save the returned tags', async () => {
machineMock.tagImage.mockResolvedValue(['tag1', 'tag2', 'tag3']);
machineMock.classifyImage.mockResolvedValue(['tag1', 'tag2', 'tag3']);
await sut.handleTagImage({ asset });
await sut.handleClassifyImage({ asset });
expect(machineMock.tagImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
expect(machineMock.classifyImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
expect(smartMock.upsert).toHaveBeenCalledWith({
assetId: 'asset-1',
tags: ['tag1', 'tag2', 'tag3'],
@@ -48,19 +83,19 @@ describe(SmartInfoService.name, () => {
});
it('should handle an error with the machine learning pipeline', async () => {
machineMock.tagImage.mockRejectedValue(new Error('Unable to read thumbnail'));
machineMock.classifyImage.mockRejectedValue(new Error('Unable to read thumbnail'));
await sut.handleTagImage({ asset });
await sut.handleClassifyImage({ asset });
expect(smartMock.upsert).not.toHaveBeenCalled();
});
it('should no update the smart info if no tags were returned', async () => {
machineMock.tagImage.mockResolvedValue([]);
machineMock.classifyImage.mockResolvedValue([]);
await sut.handleTagImage({ asset });
await sut.handleClassifyImage({ asset });
expect(machineMock.tagImage).toHaveBeenCalled();
expect(machineMock.classifyImage).toHaveBeenCalled();
expect(smartMock.upsert).not.toHaveBeenCalled();
});
});
@@ -102,4 +137,53 @@ describe(SmartInfoService.name, () => {
expect(smartMock.upsert).not.toHaveBeenCalled();
});
});
describe('handleQueueEncodeClip', () => {
it('should queue the assets without clip embeddings', async () => {
assetMock.getWithout.mockResolvedValue([assetEntityStub.image]);
await sut.handleQueueEncodeClip({ force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { asset: assetEntityStub.image } });
expect(assetMock.getWithout).toHaveBeenCalledWith(WithoutProperty.CLIP_ENCODING);
});
it('should queue all the assets', async () => {
assetMock.getAll.mockResolvedValue([assetEntityStub.image]);
await sut.handleQueueEncodeClip({ force: true });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { asset: assetEntityStub.image } });
expect(assetMock.getAll).toHaveBeenCalled();
});
});
describe('handleEncodeClip', () => {
it('should skip assets without a resize path', async () => {
await sut.handleEncodeClip({ asset: { resizePath: '' } as AssetEntity });
expect(smartMock.upsert).not.toHaveBeenCalled();
expect(machineMock.encodeImage).not.toHaveBeenCalled();
});
it('should save the returned objects', async () => {
machineMock.encodeImage.mockResolvedValue([0.01, 0.02, 0.03]);
await sut.handleEncodeClip({ asset });
expect(machineMock.encodeImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
expect(smartMock.upsert).toHaveBeenCalledWith({
assetId: 'asset-1',
clipEmbedding: [0.01, 0.02, 0.03],
});
});
it('should handle an error with the machine learning pipeline', async () => {
machineMock.encodeImage.mockRejectedValue(new Error('Unable to read thumbnail'));
await sut.handleEncodeClip({ asset });
expect(smartMock.upsert).not.toHaveBeenCalled();
});
});
});

View File

@@ -1,6 +1,7 @@
import { MACHINE_LEARNING_ENABLED } from '@app/common';
import { Inject, Injectable, Logger } from '@nestjs/common';
import { IAssetJob, IJobRepository, JobName } from '../job';
import { IAssetRepository, WithoutProperty } from '../asset';
import { IAssetJob, IBaseJob, IJobRepository, JobName } from '../job';
import { IMachineLearningRepository } from './machine-learning.interface';
import { ISmartInfoRepository } from './smart-info.repository';
@@ -9,26 +10,24 @@ export class SmartInfoService {
private logger = new Logger(SmartInfoService.name);
constructor(
@Inject(IAssetRepository) private assetRepository: IAssetRepository,
@Inject(IJobRepository) private jobRepository: IJobRepository,
@Inject(ISmartInfoRepository) private repository: ISmartInfoRepository,
@Inject(IMachineLearningRepository) private machineLearning: IMachineLearningRepository,
) {}
async handleTagImage(data: IAssetJob) {
const { asset } = data;
if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
return;
}
async handleQueueObjectTagging({ force }: IBaseJob) {
try {
const tags = await this.machineLearning.tagImage({ thumbnailPath: asset.resizePath });
if (tags.length > 0) {
await this.repository.upsert({ assetId: asset.id, tags });
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
const assets = force
? await this.assetRepository.getAll()
: await this.assetRepository.getWithout(WithoutProperty.OBJECT_TAGS);
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.CLASSIFY_IMAGE, data: { asset } });
await this.jobRepository.queue({ name: JobName.DETECT_OBJECTS, data: { asset } });
}
} catch (error: any) {
this.logger.error(`Unable to run image tagging pipeline: ${asset.id}`, error?.stack);
this.logger.error(`Unable to queue object tagging`, error?.stack);
}
}
@@ -50,6 +49,38 @@ export class SmartInfoService {
}
}
async handleClassifyImage(data: IAssetJob) {
const { asset } = data;
if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
return;
}
try {
const tags = await this.machineLearning.classifyImage({ thumbnailPath: asset.resizePath });
if (tags.length > 0) {
await this.repository.upsert({ assetId: asset.id, tags });
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
}
} catch (error: any) {
this.logger.error(`Unable to run image tagging pipeline: ${asset.id}`, error?.stack);
}
}
async handleQueueEncodeClip({ force }: IBaseJob) {
try {
const assets = force
? await this.assetRepository.getAll()
: await this.assetRepository.getWithout(WithoutProperty.CLIP_ENCODING);
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.ENCODE_CLIP, data: { asset } });
}
} catch (error: any) {
this.logger.error(`Unable to queue clip encoding`, error?.stack);
}
}
async handleEncodeClip(data: IAssetJob) {
const { asset } = data;

View File

@@ -0,0 +1,5 @@
import { basename, extname } from 'node:path';
export function getFileNameWithoutExtension(path: string): string {
return basename(path, extname(path));
}

View File

@@ -3,6 +3,7 @@ import { IAssetRepository } from '../src';
export const newAssetRepositoryMock = (): jest.Mocked<IAssetRepository> => {
return {
getByIds: jest.fn(),
getWithout: jest.fn(),
getAll: jest.fn(),
deleteAll: jest.fn(),
save: jest.fn(),

View File

@@ -3,6 +3,7 @@ import { IJobRepository } from '../src';
export const newJobRepositoryMock = (): jest.Mocked<IJobRepository> => {
return {
empty: jest.fn(),
pause: jest.fn(),
queue: jest.fn().mockImplementation(() => Promise.resolve()),
isActive: jest.fn(),
getJobCounts: jest.fn(),

View File

@@ -2,7 +2,7 @@ import { IMachineLearningRepository } from '../src';
export const newMachineLearningRepositoryMock = (): jest.Mocked<IMachineLearningRepository> => {
return {
tagImage: jest.fn(),
classifyImage: jest.fn(),
detectObjects: jest.fn(),
encodeImage: jest.fn(),
encodeText: jest.fn(),

View File

@@ -1,7 +1,7 @@
import { AssetSearchOptions, IAssetRepository } from '@app/domain';
import { AssetSearchOptions, IAssetRepository, WithoutProperty } from '@app/domain';
import { Injectable } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { In, Not, Repository } from 'typeorm';
import { FindOptionsRelations, FindOptionsWhere, In, IsNull, Not, Repository } from 'typeorm';
import { AssetEntity, AssetType } from '../entities';
@Injectable()
@@ -65,4 +65,73 @@ export class AssetRepository implements IAssetRepository {
},
});
}
getWithout(property: WithoutProperty): Promise<AssetEntity[]> {
let relations: FindOptionsRelations<AssetEntity> = {};
let where: FindOptionsWhere<AssetEntity> | FindOptionsWhere<AssetEntity>[] = {};
switch (property) {
case WithoutProperty.THUMBNAIL:
where = [
{ resizePath: IsNull(), isVisible: true },
{ resizePath: '', isVisible: true },
{ webpPath: IsNull(), isVisible: true },
{ webpPath: '', isVisible: true },
];
break;
case WithoutProperty.ENCODED_VIDEO:
where = [
{ type: AssetType.VIDEO, encodedVideoPath: IsNull() },
{ type: AssetType.VIDEO, encodedVideoPath: '' },
];
break;
case WithoutProperty.EXIF:
relations = {
exifInfo: true,
};
where = {
isVisible: true,
resizePath: Not(IsNull()),
exifInfo: {
assetId: IsNull(),
},
};
break;
case WithoutProperty.CLIP_ENCODING:
relations = {
smartInfo: true,
};
where = {
isVisible: true,
smartInfo: {
clipEmbedding: IsNull(),
},
};
break;
case WithoutProperty.OBJECT_TAGS:
relations = {
smartInfo: true,
};
where = {
resizePath: IsNull(),
isVisible: true,
smartInfo: {
tags: IsNull(),
},
};
break;
default:
throw new Error(`Invalid getWithout property: ${property}`);
}
return this.repository.find({
relations,
where,
});
}
}

View File

@@ -1,18 +1,38 @@
import { IAssetJob, IJobRepository, IMetadataExtractionJob, JobCounts, JobItem, JobName, QueueName } from '@app/domain';
import {
IAssetJob,
IBaseJob,
IJobRepository,
IMetadataExtractionJob,
JobCounts,
JobItem,
JobName,
QueueName,
} from '@app/domain';
import { InjectQueue } from '@nestjs/bull';
import { BadRequestException, Logger } from '@nestjs/common';
import { Logger } from '@nestjs/common';
import { Queue } from 'bull';
export class JobRepository implements IJobRepository {
private logger = new Logger(JobRepository.name);
private queueMap: Record<QueueName, Queue> = {
[QueueName.STORAGE_TEMPLATE_MIGRATION]: this.storageTemplateMigration,
[QueueName.THUMBNAIL_GENERATION]: this.generateThumbnail,
[QueueName.METADATA_EXTRACTION]: this.metadataExtraction,
[QueueName.OBJECT_TAGGING]: this.objectTagging,
[QueueName.CLIP_ENCODING]: this.clipEmbedding,
[QueueName.VIDEO_CONVERSION]: this.videoTranscode,
[QueueName.BACKGROUND_TASK]: this.backgroundTask,
[QueueName.SEARCH]: this.searchIndex,
};
constructor(
@InjectQueue(QueueName.BACKGROUND_TASK) private backgroundTask: Queue,
@InjectQueue(QueueName.MACHINE_LEARNING) private machineLearning: Queue<IAssetJob>,
@InjectQueue(QueueName.METADATA_EXTRACTION) private metadataExtraction: Queue<IMetadataExtractionJob>,
@InjectQueue(QueueName.OBJECT_TAGGING) private objectTagging: Queue<IAssetJob | IBaseJob>,
@InjectQueue(QueueName.CLIP_ENCODING) private clipEmbedding: Queue<IAssetJob | IBaseJob>,
@InjectQueue(QueueName.METADATA_EXTRACTION) private metadataExtraction: Queue<IMetadataExtractionJob | IBaseJob>,
@InjectQueue(QueueName.STORAGE_TEMPLATE_MIGRATION) private storageTemplateMigration: Queue,
@InjectQueue(QueueName.THUMBNAIL_GENERATION) private thumbnail: Queue,
@InjectQueue(QueueName.VIDEO_CONVERSION) private videoTranscode: Queue<IAssetJob>,
@InjectQueue(QueueName.THUMBNAIL_GENERATION) private generateThumbnail: Queue,
@InjectQueue(QueueName.VIDEO_CONVERSION) private videoTranscode: Queue<IAssetJob | IBaseJob>,
@InjectQueue(QueueName.SEARCH) private searchIndex: Queue,
) {}
@@ -21,12 +41,16 @@ export class JobRepository implements IJobRepository {
return !!counts.active;
}
pause(name: QueueName) {
return this.queueMap[name].pause();
}
empty(name: QueueName) {
return this.getQueue(name).empty();
return this.queueMap[name].empty();
}
getJobCounts(name: QueueName): Promise<JobCounts> {
return this.getQueue(name).getJobCounts();
return this.queueMap[name].getJobCounts();
}
async queue(item: JobItem): Promise<void> {
@@ -39,21 +63,28 @@ export class JobRepository implements IJobRepository {
await this.backgroundTask.add(item.name, item.data);
break;
case JobName.OBJECT_DETECTION:
case JobName.IMAGE_TAGGING:
case JobName.ENCODE_CLIP:
await this.machineLearning.add(item.name, item.data);
case JobName.QUEUE_OBJECT_TAGGING:
case JobName.DETECT_OBJECTS:
case JobName.CLASSIFY_IMAGE:
await this.objectTagging.add(item.name, item.data);
break;
case JobName.QUEUE_ENCODE_CLIP:
case JobName.ENCODE_CLIP:
await this.clipEmbedding.add(item.name, item.data);
break;
case JobName.QUEUE_METADATA_EXTRACTION:
case JobName.EXIF_EXTRACTION:
case JobName.EXTRACT_VIDEO_METADATA:
case JobName.REVERSE_GEOCODING:
await this.metadataExtraction.add(item.name, item.data);
break;
case JobName.QUEUE_GENERATE_THUMBNAILS:
case JobName.GENERATE_JPEG_THUMBNAIL:
case JobName.GENERATE_WEBP_THUMBNAIL:
await this.thumbnail.add(item.name, item.data);
await this.generateThumbnail.add(item.name, item.data);
break;
case JobName.USER_DELETION:
@@ -68,6 +99,7 @@ export class JobRepository implements IJobRepository {
await this.backgroundTask.add(item.name, {});
break;
case JobName.QUEUE_VIDEO_CONVERSION:
case JobName.VIDEO_CONVERSION:
await this.videoTranscode.add(item.name, item.data);
break;
@@ -85,25 +117,7 @@ export class JobRepository implements IJobRepository {
break;
default:
// TODO inject remaining queues and map job to queue
this.logger.error('Invalid job', item);
}
}
private getQueue(name: QueueName) {
switch (name) {
case QueueName.STORAGE_TEMPLATE_MIGRATION:
return this.storageTemplateMigration;
case QueueName.THUMBNAIL_GENERATION:
return this.thumbnail;
case QueueName.METADATA_EXTRACTION:
return this.metadataExtraction;
case QueueName.VIDEO_CONVERSION:
return this.videoTranscode;
case QueueName.MACHINE_LEARNING:
return this.machineLearning;
default:
throw new BadRequestException('Invalid job name');
}
}
}

View File

@@ -7,7 +7,7 @@ const client = axios.create({ baseURL: MACHINE_LEARNING_URL });
@Injectable()
export class MachineLearningRepository implements IMachineLearningRepository {
tagImage(input: MachineLearningInput): Promise<string[]> {
classifyImage(input: MachineLearningInput): Promise<string[]> {
return client.post<string[]>('/image-classifier/tag-image', input).then((res) => res.data);
}