Implement GPU mode

fixes #67 Signed-off-by: Marcel Klehr <mklehr@gmx.net>
nextcloud · Dec 4, 2022 · a0d74e6 · a0d74e6
1 parent d2849a3
commit a0d74e6
Show file tree

Hide file tree

Showing 10 changed files with 396 additions and 129 deletions.
diff --git a/Makefile b/Makefile
@@ -39,6 +39,8 @@ remove-binaries:
 	# make it download appropriate tf binaries
 	rm -rf node_modules/@tensorflow/tfjs-node/deps/lib/*
 	rm -rf node_modules/@tensorflow/tfjs-node/lib/*
+	rm -rf node_modules/@tensorflow/tfjs-node-gpu/lib/*
+	rm -rf node_modules/@tensorflow/tfjs-node-gpu/deps/lib/*
 
 remove-devdeps:
 	rm -rf node_modules

diff --git a/appinfo/routes.php b/appinfo/routes.php
@@ -33,6 +33,7 @@
 		['name' => 'admin#nodejs', 'url' => '/admin/nodejs', 'verb' => 'GET'],
 		['name' => 'admin#libtensorflow', 'url' => '/admin/libtensorflow', 'verb' => 'GET'],
 		['name' => 'admin#wasmtensorflow', 'url' => '/admin/wasmtensorflow', 'verb' => 'GET'],
+		['name' => 'admin#gputensorflow', 'url' => '/admin/gputensorflow', 'verb' => 'GET'],
 		['name' => 'admin#cron', 'url' => '/admin/cron', 'verb' => 'GET'],
 		['name' => 'admin#get_setting', 'url' => '/admin/settings/{setting}', 'verb' => 'GET'],
 		['name' => 'admin#set_setting', 'url' => '/admin/settings/{setting}', 'verb' => 'PUT'],

diff --git a/lib/Controller/AdminController.php b/lib/Controller/AdminController.php
@@ -167,6 +167,20 @@ public function wasmtensorflow(): JSONResponse {
 		return new JSONResponse(['wasmtensorflow' => true]);
 	}
 
+	public function gputensorflow(): JSONResponse {
+		try {
+			exec($this->settingsService->getSetting('node_binary') . ' ' . __DIR__ . '/../../src/test_gputensorflow.js' . ' 2>&1', $output, $returnCode);
+		} catch (\Throwable $e) {
+			return new JSONResponse(['gputensorflow' => false]);
+		}
+
+		if ($returnCode !== 0) {
+			return new JSONResponse(['gputensorflow' => false]);
+		}
+
+		return new JSONResponse(['gputensorflow' => true]);
+	}
+
 	public function cron(): JSONResponse {
 		$cron = $this->config->getAppValue('core', 'backgroundjobs_mode', '');
 		return new JSONResponse(['cron' => $cron]);

diff --git a/lib/Migration/InstallDeps.php b/lib/Migration/InstallDeps.php
@@ -50,14 +50,18 @@ class InstallDeps implements IRepairStep {
 	private string $tfjsPath;
 	private IClientService $clientService;
 	private LoggerInterface $logger;
+	private string $runTfjsGpuInstall;
+	private string $tfjsGPUPath;
 
 	public function __construct(IConfig $config, IClientService $clientService, LoggerInterface $logger) {
 		$this->config = $config;
 		$this->binaryDir = dirname(__DIR__, 2) . '/bin/';
 		$this->preGypBinaryDir = dirname(__DIR__, 2) . '/node_modules/@mapbox/node-pre-gyp/bin/';
 		$this->ffmpegDir = dirname(__DIR__, 2) . '/node_modules/ffmpeg-static/';
 		$this->tfjsInstallScript = dirname(__DIR__, 2) . '/node_modules/@tensorflow/tfjs-node/scripts/install.js';
+		$this->runTfjsGpuInstall = dirname(__DIR__, 2) . '/node_modules/@tensorflow/tfjs-node-gpu/scripts/install.js';
 		$this->tfjsPath = dirname(__DIR__, 2) . '/node_modules/@tensorflow/tfjs-node/';
+		$this->tfjsGPUPath = dirname(__DIR__, 2) . '/node_modules/@tensorflow/tfjs-node-gpu/';
 		$this->clientService = $clientService;
 		$this->logger = $logger;
 	}
@@ -116,6 +120,7 @@ public function run(IOutput $output): void {
 		$this->setBinariesPermissions();
 
 		$this->runTfjsInstall($binaryPath);
+		$this->runTfjsGpuInstall($binaryPath);
 	}
 
 	protected function testBinary(string $binaryPath): ?string {
@@ -147,10 +152,28 @@ protected function runTfjsInstall(string $nodeBinary) : void {
 		}
 		chdir($oriCwd);
 		if ($returnCode !== 0) {
+			$this->logger->error('Failed to install Tensorflow.js: '.trim(implode("\n", $output)));
 			throw new \Exception('Failed to install Tensorflow.js: '.trim(implode("\n", $output)));
 		}
 	}
 
+	protected function runTfjsGpuInstall(string $nodeBinary) : void {
+		$oriCwd = getcwd();
+		chdir($this->tfjsGPUPath);
+		$cmd = 'PATH='.escapeshellcmd($this->preGypBinaryDir).':'.escapeshellcmd($this->binaryDir).':$PATH ' . escapeshellcmd($nodeBinary) . ' ' . escapeshellarg($this->tfjsGpuInstallScript) . ' gpu ' . escapeshellarg('download');
+		try {
+			exec($cmd . ' 2>&1', $output, $returnCode); // Appending  2>&1 to avoid leaking sterr
+		} catch (\Throwable $e) {
+			$this->logger->error('Failed to install Tensorflow.js for GPU: '.$e->getMessage(), ['exception' => $e]);
+			throw new \Exception('Failed to install Tensorflow.js for GPU: '.$e->getMessage());
+		}
+		chdir($oriCwd);
+		if ($returnCode !== 0) {
+			$this->logger->error('Failed to install Tensorflow.js for GPU: '.trim(implode("\n", $output)));
+			throw new \Exception('Failed to install Tensorflow.js for GPU: '.trim(implode("\n", $output)));
+		}
+	}
+
 	protected function downloadNodeBinary(string $server, string $version, string $arch, string $flavor = '') : string {
 		$name = 'node-'.$version.'-linux-'.$arch;
 		if ($flavor !== '') {