diff options
Diffstat (limited to 'lib/private/TextToImage')
| -rw-r--r-- | lib/private/TextToImage/Db/Task.php | 101 | ||||
| -rw-r--r-- | lib/private/TextToImage/Db/TaskMapper.php | 111 | ||||
| -rw-r--r-- | lib/private/TextToImage/Manager.php | 325 | ||||
| -rw-r--r-- | lib/private/TextToImage/RemoveOldTasksBackgroundJob.php | 63 | ||||
| -rw-r--r-- | lib/private/TextToImage/TaskBackgroundJob.php | 47 |
5 files changed, 647 insertions, 0 deletions
diff --git a/lib/private/TextToImage/Db/Task.php b/lib/private/TextToImage/Db/Task.php new file mode 100644 index 00000000000..cc92c865220 --- /dev/null +++ b/lib/private/TextToImage/Db/Task.php @@ -0,0 +1,101 @@ +<?php + +declare(strict_types=1); + +/** + * SPDX-FileCopyrightText: 2023 Nextcloud GmbH and Nextcloud contributors + * SPDX-License-Identifier: AGPL-3.0-or-later + */ + +namespace OC\TextToImage\Db; + +use DateTime; +use OCP\AppFramework\Db\Entity; +use OCP\AppFramework\Utility\ITimeFactory; +use OCP\TextToImage\Task as OCPTask; + +/** + * @method setLastUpdated(DateTime $lastUpdated) + * @method DateTime getLastUpdated() + * @method setInput(string $type) + * @method string getInput() + * @method setResultPath(string $resultPath) + * @method string getResultPath() + * @method setStatus(int $type) + * @method int getStatus() + * @method setUserId(?string $userId) + * @method string|null getUserId() + * @method setAppId(string $type) + * @method string getAppId() + * @method setIdentifier(string $identifier) + * @method string|null getIdentifier() + * @method setNumberOfImages(int $numberOfImages) + * @method int getNumberOfImages() + * @method setCompletionExpectedAt(DateTime $at) + * @method DateTime getCompletionExpectedAt() + */ +class Task extends Entity { + protected $lastUpdated; + protected $type; + protected $input; + protected $status; + protected $userId; + protected $appId; + protected $identifier; + protected $numberOfImages; + protected $completionExpectedAt; + + /** + * @var string[] + */ + public static array $columns = ['id', 'last_updated', 'input', 'status', 'user_id', 'app_id', 'identifier', 'number_of_images', 'completion_expected_at']; + + /** + * @var string[] + */ + public static array $fields = ['id', 'lastUpdated', 'input', 'status', 'userId', 'appId', 'identifier', 'numberOfImages', 'completionExpectedAt']; + + + public function __construct() { + // add types in constructor + $this->addType('id', 'integer'); + $this->addType('lastUpdated', 'datetime'); + $this->addType('input', 'string'); + $this->addType('status', 'integer'); + $this->addType('userId', 'string'); + $this->addType('appId', 'string'); + $this->addType('identifier', 'string'); + $this->addType('numberOfImages', 'integer'); + $this->addType('completionExpectedAt', 'datetime'); + } + + public function toRow(): array { + return array_combine(self::$columns, array_map(function ($field) { + return $this->{'get' . ucfirst($field)}(); + }, self::$fields)); + } + + public static function fromPublicTask(OCPTask $task): Task { + /** @var Task $dbTask */ + $dbTask = Task::fromParams([ + 'id' => $task->getId(), + 'lastUpdated' => \OCP\Server::get(ITimeFactory::class)->getDateTime(), + 'status' => $task->getStatus(), + 'numberOfImages' => $task->getNumberOfImages(), + 'input' => $task->getInput(), + 'userId' => $task->getUserId(), + 'appId' => $task->getAppId(), + 'identifier' => $task->getIdentifier(), + 'completionExpectedAt' => $task->getCompletionExpectedAt(), + ]); + return $dbTask; + } + + public function toPublicTask(): OCPTask { + $task = new OCPTask($this->getInput(), $this->getAppId(), $this->getNumberOfImages(), $this->getuserId(), $this->getIdentifier()); + $task->setId($this->getId()); + $task->setStatus($this->getStatus()); + $task->setCompletionExpectedAt($this->getCompletionExpectedAt()); + return $task; + } +} diff --git a/lib/private/TextToImage/Db/TaskMapper.php b/lib/private/TextToImage/Db/TaskMapper.php new file mode 100644 index 00000000000..37f492e14cf --- /dev/null +++ b/lib/private/TextToImage/Db/TaskMapper.php @@ -0,0 +1,111 @@ +<?php + +declare(strict_types=1); + +/** + * SPDX-FileCopyrightText: 2023 Nextcloud GmbH and Nextcloud contributors + * SPDX-License-Identifier: AGPL-3.0-or-later + */ + +namespace OC\TextToImage\Db; + +use OCP\AppFramework\Db\DoesNotExistException; +use OCP\AppFramework\Db\Entity; +use OCP\AppFramework\Db\MultipleObjectsReturnedException; +use OCP\AppFramework\Db\QBMapper; +use OCP\AppFramework\Utility\ITimeFactory; +use OCP\DB\Exception; +use OCP\DB\QueryBuilder\IQueryBuilder; +use OCP\IDBConnection; + +/** + * @extends QBMapper<Task> + */ +class TaskMapper extends QBMapper { + public function __construct( + IDBConnection $db, + private ITimeFactory $timeFactory, + ) { + parent::__construct($db, 'text2image_tasks', Task::class); + } + + /** + * @param int $id + * @return Task + * @throws Exception + * @throws DoesNotExistException + * @throws MultipleObjectsReturnedException + */ + public function find(int $id): Task { + $qb = $this->db->getQueryBuilder(); + $qb->select(Task::$columns) + ->from($this->tableName) + ->where($qb->expr()->eq('id', $qb->createPositionalParameter($id))); + return $this->findEntity($qb); + } + + /** + * @param int $id + * @param string|null $userId + * @return Task + * @throws DoesNotExistException + * @throws Exception + * @throws MultipleObjectsReturnedException + */ + public function findByIdAndUser(int $id, ?string $userId): Task { + $qb = $this->db->getQueryBuilder(); + $qb->select(Task::$columns) + ->from($this->tableName) + ->where($qb->expr()->eq('id', $qb->createPositionalParameter($id))); + if ($userId === null) { + $qb->andWhere($qb->expr()->isNull('user_id')); + } else { + $qb->andWhere($qb->expr()->eq('user_id', $qb->createPositionalParameter($userId))); + } + return $this->findEntity($qb); + } + + /** + * @param string $userId + * @param string $appId + * @param string|null $identifier + * @return array + * @throws Exception + */ + public function findUserTasksByApp(?string $userId, string $appId, ?string $identifier = null): array { + $qb = $this->db->getQueryBuilder(); + $qb->select(Task::$columns) + ->from($this->tableName) + ->where($qb->expr()->eq('user_id', $qb->createPositionalParameter($userId))) + ->andWhere($qb->expr()->eq('app_id', $qb->createPositionalParameter($appId))); + if ($identifier !== null) { + $qb->andWhere($qb->expr()->eq('identifier', $qb->createPositionalParameter($identifier))); + } + return $this->findEntities($qb); + } + + /** + * @param int $timeout + * @return Task[] the deleted tasks + * @throws Exception + */ + public function deleteOlderThan(int $timeout): array { + $datetime = $this->timeFactory->getDateTime(); + $datetime->sub(new \DateInterval('PT' . $timeout . 'S')); + $qb = $this->db->getQueryBuilder(); + $qb->select('*') + ->from($this->tableName) + ->where($qb->expr()->lt('last_updated', $qb->createPositionalParameter($datetime, IQueryBuilder::PARAM_DATETIME_MUTABLE))); + $deletedTasks = $this->findEntities($qb); + $qb = $this->db->getQueryBuilder(); + $qb->delete($this->tableName) + ->where($qb->expr()->lt('last_updated', $qb->createPositionalParameter($datetime, IQueryBuilder::PARAM_DATETIME_MUTABLE))); + $qb->executeStatement(); + return $deletedTasks; + } + + public function update(Entity $entity): Entity { + $entity->setLastUpdated($this->timeFactory->getDateTime()); + return parent::update($entity); + } +} diff --git a/lib/private/TextToImage/Manager.php b/lib/private/TextToImage/Manager.php new file mode 100644 index 00000000000..eec6cc3d241 --- /dev/null +++ b/lib/private/TextToImage/Manager.php @@ -0,0 +1,325 @@ +<?php + +declare(strict_types=1); + +/** + * SPDX-FileCopyrightText: 2023 Nextcloud GmbH and Nextcloud contributors + * SPDX-License-Identifier: AGPL-3.0-or-later + */ + +namespace OC\TextToImage; + +use OC\AppFramework\Bootstrap\Coordinator; +use OC\TextToImage\Db\Task as DbTask; +use OC\TextToImage\Db\TaskMapper; +use OCP\AppFramework\Db\DoesNotExistException; +use OCP\AppFramework\Db\MultipleObjectsReturnedException; +use OCP\BackgroundJob\IJobList; +use OCP\DB\Exception; +use OCP\Files\AppData\IAppDataFactory; +use OCP\Files\IAppData; +use OCP\Files\NotFoundException; +use OCP\Files\NotPermittedException; +use OCP\IConfig; +use OCP\IServerContainer; +use OCP\PreConditionNotMetException; +use OCP\TextToImage\Exception\TaskFailureException; +use OCP\TextToImage\Exception\TaskNotFoundException; +use OCP\TextToImage\IManager; +use OCP\TextToImage\IProvider; +use OCP\TextToImage\IProviderWithUserId; +use OCP\TextToImage\Task; +use Psr\Log\LoggerInterface; +use RuntimeException; +use Throwable; + +class Manager implements IManager { + /** @var ?list<IProvider> */ + private ?array $providers = null; + private IAppData $appData; + + public function __construct( + private IServerContainer $serverContainer, + private Coordinator $coordinator, + private LoggerInterface $logger, + private IJobList $jobList, + private TaskMapper $taskMapper, + private IConfig $config, + IAppDataFactory $appDataFactory, + ) { + $this->appData = $appDataFactory->get('core'); + } + + /** + * @inheritDoc + */ + public function getProviders(): array { + $context = $this->coordinator->getRegistrationContext(); + if ($context === null) { + return []; + } + + if ($this->providers !== null) { + return $this->providers; + } + + $this->providers = []; + + foreach ($context->getTextToImageProviders() as $providerServiceRegistration) { + $class = $providerServiceRegistration->getService(); + try { + /** @var IProvider $provider */ + $provider = $this->serverContainer->get($class); + $this->providers[] = $provider; + } catch (Throwable $e) { + $this->logger->error('Failed to load Text to image provider ' . $class, [ + 'exception' => $e, + ]); + } + } + + return $this->providers; + } + + /** + * @inheritDoc + */ + public function hasProviders(): bool { + $context = $this->coordinator->getRegistrationContext(); + if ($context === null) { + return false; + } + return count($context->getTextToImageProviders()) > 0; + } + + /** + * @inheritDoc + */ + public function runTask(Task $task): void { + $this->logger->debug('Running TextToImage Task'); + if (!$this->hasProviders()) { + throw new PreConditionNotMetException('No text to image provider is installed that can handle this task'); + } + $providers = $this->getPreferredProviders(); + + foreach ($providers as $provider) { + $this->logger->debug('Trying to run Text2Image provider ' . $provider::class); + try { + $task->setStatus(Task::STATUS_RUNNING); + $completionExpectedAt = new \DateTime('now'); + $completionExpectedAt->add(new \DateInterval('PT' . $provider->getExpectedRuntime() . 'S')); + $task->setCompletionExpectedAt($completionExpectedAt); + if ($task->getId() === null) { + $this->logger->debug('Inserting Text2Image task into DB'); + $taskEntity = $this->taskMapper->insert(DbTask::fromPublicTask($task)); + $task->setId($taskEntity->getId()); + } else { + $this->logger->debug('Updating Text2Image task in DB'); + $this->taskMapper->update(DbTask::fromPublicTask($task)); + } + try { + $folder = $this->appData->getFolder('text2image'); + } catch (NotFoundException) { + $this->logger->debug('Creating folder in appdata for Text2Image results'); + $folder = $this->appData->newFolder('text2image'); + } + try { + $folder = $folder->getFolder((string)$task->getId()); + } catch (NotFoundException) { + $this->logger->debug('Creating new folder in appdata Text2Image results folder'); + $folder = $folder->newFolder((string)$task->getId()); + } + $this->logger->debug('Creating result files for Text2Image task'); + $resources = []; + $files = []; + for ($i = 0; $i < $task->getNumberOfImages(); $i++) { + $file = $folder->newFile((string)$i); + $files[] = $file; + $resource = $file->write(); + if ($resource !== false && $resource !== true && is_resource($resource)) { + $resources[] = $resource; + } else { + throw new RuntimeException('Text2Image generation using provider "' . $provider->getName() . '" failed: Couldn\'t open file to write.'); + } + } + $this->logger->debug('Calling Text2Image provider\'s generate method'); + if ($provider instanceof IProviderWithUserId) { + $provider->setUserId($task->getUserId()); + } + $provider->generate($task->getInput(), $resources); + for ($i = 0; $i < $task->getNumberOfImages(); $i++) { + if (is_resource($resources[$i])) { + // If $resource hasn't been closed yet, we'll do that here + fclose($resources[$i]); + } + } + $task->setStatus(Task::STATUS_SUCCESSFUL); + $this->logger->debug('Updating Text2Image task in DB'); + $this->taskMapper->update(DbTask::fromPublicTask($task)); + return; + } catch (\RuntimeException|\Throwable $e) { + for ($i = 0; $i < $task->getNumberOfImages(); $i++) { + if (isset($files, $files[$i])) { + try { + $files[$i]->delete(); + } catch (NotPermittedException $e) { + $this->logger->warning('Failed to clean up Text2Image result file after error', ['exception' => $e]); + } + } + } + + $this->logger->info('Text2Image generation using provider "' . $provider->getName() . '" failed', ['exception' => $e]); + $task->setStatus(Task::STATUS_FAILED); + try { + $this->taskMapper->update(DbTask::fromPublicTask($task)); + } catch (Exception $e) { + $this->logger->warning('Failed to update database after Text2Image error', ['exception' => $e]); + } + throw new TaskFailureException('Text2Image generation using provider "' . $provider->getName() . '" failed: ' . $e->getMessage(), 0, $e); + } + } + + $task->setStatus(Task::STATUS_FAILED); + try { + $this->taskMapper->update(DbTask::fromPublicTask($task)); + } catch (Exception $e) { + $this->logger->warning('Failed to update database after Text2Image error', ['exception' => $e]); + } + throw new TaskFailureException('Could not run task'); + } + + /** + * @inheritDoc + */ + public function scheduleTask(Task $task): void { + if (!$this->hasProviders()) { + throw new PreConditionNotMetException('No text to image provider is installed that can handle this task'); + } + $this->logger->debug('Scheduling Text2Image Task'); + $task->setStatus(Task::STATUS_SCHEDULED); + $completionExpectedAt = new \DateTime('now'); + $completionExpectedAt->add(new \DateInterval('PT' . $this->getPreferredProviders()[0]->getExpectedRuntime() . 'S')); + $task->setCompletionExpectedAt($completionExpectedAt); + $taskEntity = DbTask::fromPublicTask($task); + $this->taskMapper->insert($taskEntity); + $task->setId($taskEntity->getId()); + $this->jobList->add(TaskBackgroundJob::class, [ + 'taskId' => $task->getId() + ]); + } + + /** + * @inheritDoc + */ + public function runOrScheduleTask(Task $task) : void { + if (!$this->hasProviders()) { + throw new PreConditionNotMetException('No text to image provider is installed that can handle this task'); + } + $providers = $this->getPreferredProviders(); + $maxExecutionTime = (int)ini_get('max_execution_time'); + // Offload the task to a background job if the expected runtime of the likely provider is longer than 80% of our max execution time + if ($providers[0]->getExpectedRuntime() > $maxExecutionTime * 0.8) { + $this->scheduleTask($task); + return; + } + $this->runTask($task); + } + + /** + * @inheritDoc + */ + public function deleteTask(Task $task): void { + $taskEntity = DbTask::fromPublicTask($task); + $this->taskMapper->delete($taskEntity); + $this->jobList->remove(TaskBackgroundJob::class, [ + 'taskId' => $task->getId() + ]); + } + + /** + * Get a task from its id + * + * @param int $id The id of the task + * @return Task + * @throws RuntimeException If the query failed + * @throws TaskNotFoundException If the task could not be found + */ + public function getTask(int $id): Task { + try { + $taskEntity = $this->taskMapper->find($id); + return $taskEntity->toPublicTask(); + } catch (DoesNotExistException $e) { + throw new TaskNotFoundException('Could not find task with the provided id'); + } catch (MultipleObjectsReturnedException $e) { + throw new RuntimeException('Could not uniquely identify task with given id', 0, $e); + } catch (Exception $e) { + throw new RuntimeException('Failure while trying to find task by id: ' . $e->getMessage(), 0, $e); + } + } + + /** + * Get a task from its user id and task id + * If userId is null, this can only get a task that was scheduled anonymously + * + * @param int $id The id of the task + * @param string|null $userId The user id that scheduled the task + * @return Task + * @throws RuntimeException If the query failed + * @throws TaskNotFoundException If the task could not be found + */ + public function getUserTask(int $id, ?string $userId): Task { + try { + $taskEntity = $this->taskMapper->findByIdAndUser($id, $userId); + return $taskEntity->toPublicTask(); + } catch (DoesNotExistException $e) { + throw new TaskNotFoundException('Could not find task with the provided id and user id'); + } catch (MultipleObjectsReturnedException $e) { + throw new RuntimeException('Could not uniquely identify task with given id and user id', 0, $e); + } catch (Exception $e) { + throw new RuntimeException('Failure while trying to find task by id and user id: ' . $e->getMessage(), 0, $e); + } + } + + /** + * Get a list of tasks scheduled by a specific user for a specific app + * and optionally with a specific identifier. + * This cannot be used to get anonymously scheduled tasks + * + * @param string $userId + * @param string $appId + * @param string|null $identifier + * @return Task[] + * @throws RuntimeException + */ + public function getUserTasksByApp(?string $userId, string $appId, ?string $identifier = null): array { + try { + $taskEntities = $this->taskMapper->findUserTasksByApp($userId, $appId, $identifier); + return array_map(static function (DbTask $taskEntity) { + return $taskEntity->toPublicTask(); + }, $taskEntities); + } catch (Exception $e) { + throw new RuntimeException('Failure while trying to find tasks by appId and identifier: ' . $e->getMessage(), 0, $e); + } + } + + /** + * @return list<IProvider> + */ + private function getPreferredProviders() { + $providers = $this->getProviders(); + $json = $this->config->getAppValue('core', 'ai.text2image_provider', ''); + if ($json !== '') { + try { + $id = json_decode($json, true, 512, JSON_THROW_ON_ERROR); + $provider = current(array_filter($providers, fn ($provider) => $provider->getId() === $id)); + if ($provider !== false && $provider !== null) { + $providers = [$provider]; + } + } catch (\JsonException $e) { + $this->logger->warning('Failed to decode Text2Image setting `ai.text2image_provider`', ['exception' => $e]); + } + } + + return $providers; + } +} diff --git a/lib/private/TextToImage/RemoveOldTasksBackgroundJob.php b/lib/private/TextToImage/RemoveOldTasksBackgroundJob.php new file mode 100644 index 00000000000..225984a7956 --- /dev/null +++ b/lib/private/TextToImage/RemoveOldTasksBackgroundJob.php @@ -0,0 +1,63 @@ +<?php + +declare(strict_types=1); + +/** + * SPDX-FileCopyrightText: 2023 Nextcloud GmbH and Nextcloud contributors + * SPDX-License-Identifier: AGPL-3.0-or-later + */ + + +namespace OC\TextToImage; + +use OC\TextToImage\Db\TaskMapper; +use OCP\AppFramework\Utility\ITimeFactory; +use OCP\BackgroundJob\TimedJob; +use OCP\DB\Exception; +use OCP\Files\AppData\IAppDataFactory; +use OCP\Files\IAppData; +use OCP\Files\NotFoundException; +use OCP\Files\NotPermittedException; +use Psr\Log\LoggerInterface; + +class RemoveOldTasksBackgroundJob extends TimedJob { + public const MAX_TASK_AGE_SECONDS = 60 * 60 * 24 * 7; // 1 week + + private IAppData $appData; + + public function __construct( + ITimeFactory $timeFactory, + private TaskMapper $taskMapper, + private LoggerInterface $logger, + IAppDataFactory $appDataFactory, + ) { + parent::__construct($timeFactory); + $this->appData = $appDataFactory->get('core'); + $this->setInterval(60 * 60 * 24); + $this->setTimeSensitivity(self::TIME_INSENSITIVE); + } + + /** + * @param mixed $argument + * @inheritDoc + */ + protected function run($argument) { + try { + $deletedTasks = $this->taskMapper->deleteOlderThan(self::MAX_TASK_AGE_SECONDS); + $folder = $this->appData->getFolder('text2image'); + foreach ($deletedTasks as $deletedTask) { + try { + $folder->getFolder((string)$deletedTask->getId())->delete(); + } catch (NotFoundException) { + // noop + } catch (NotPermittedException $e) { + $this->logger->warning('Failed to delete stale text to image task files', ['exception' => $e]); + } + } + } catch (Exception $e) { + $this->logger->warning('Failed to delete stale text to image tasks', ['exception' => $e]); + } catch (NotFoundException) { + // noop + } + } +} diff --git a/lib/private/TextToImage/TaskBackgroundJob.php b/lib/private/TextToImage/TaskBackgroundJob.php new file mode 100644 index 00000000000..16990005530 --- /dev/null +++ b/lib/private/TextToImage/TaskBackgroundJob.php @@ -0,0 +1,47 @@ +<?php + +declare(strict_types=1); + +/** + * SPDX-FileCopyrightText: 2023 Nextcloud GmbH and Nextcloud contributors + * SPDX-License-Identifier: AGPL-3.0-or-later + */ + + +namespace OC\TextToImage; + +use OCP\AppFramework\Utility\ITimeFactory; +use OCP\BackgroundJob\QueuedJob; +use OCP\EventDispatcher\IEventDispatcher; +use OCP\TextToImage\Events\TaskFailedEvent; +use OCP\TextToImage\Events\TaskSuccessfulEvent; +use OCP\TextToImage\IManager; + +class TaskBackgroundJob extends QueuedJob { + public function __construct( + ITimeFactory $timeFactory, + private IManager $text2imageManager, + private IEventDispatcher $eventDispatcher, + ) { + parent::__construct($timeFactory); + // We want to avoid overloading the machine with these jobs + // so we only allow running one job at a time + $this->setAllowParallelRuns(false); + } + + /** + * @param array{taskId: int} $argument + * @inheritDoc + */ + protected function run($argument) { + $taskId = $argument['taskId']; + $task = $this->text2imageManager->getTask($taskId); + try { + $this->text2imageManager->runTask($task); + $event = new TaskSuccessfulEvent($task); + } catch (\Throwable $e) { + $event = new TaskFailedEvent($task, $e->getMessage()); + } + $this->eventDispatcher->dispatchTyped($event); + } +} |