feat: prioritize gpus (#1768)

* feat: prioritize GPUs * fix: migrate db * fix: add priority * fix: db * fix: more * fix: migration --------- Co-authored-by: vansangpfiev <[email protected]>
janhq · Dec 10, 2024 · 0fa83b2 · 0fa83b2
1 parent 6300732
commit 0fa83b2
Show file tree

Hide file tree

Showing 8 changed files with 331 additions and 65 deletions.
diff --git a/engine/controllers/hardware.cc b/engine/controllers/hardware.cc
@@ -38,7 +38,7 @@ void Hardware::Activate(
       ahc.gpus.push_back(g.asInt());
     }
   }
-  std::sort(ahc.gpus.begin(), ahc.gpus.end());
+
   if (!hw_svc_->IsValidConfig(ahc)) {
     Json::Value ret;
     ret["message"] = "Invalid GPU index provided.";

diff --git a/engine/database/hardware.cc b/engine/database/hardware.cc
@@ -5,22 +5,23 @@
 
 namespace cortex::db {
 
-Hardwares::Hardwares() : db_(cortex::db::Database::GetInstance().db()) {}
+Hardware::Hardware() : db_(cortex::db::Database::GetInstance().db()) {}
 
-Hardwares::Hardwares(SQLite::Database& db) : db_(db) {}
+Hardware::Hardware(SQLite::Database& db) : db_(db) {}
 
-Hardwares::~Hardwares() {}
+
+Hardware::~Hardware() {}
 
 cpp::result<std::vector<HardwareEntry>, std::string>
-Hardwares::LoadHardwareList() const {
+Hardware::LoadHardwareList() const {
   try {
     db_.exec("BEGIN TRANSACTION;");
     cortex::utils::ScopeExit se([this] { db_.exec("COMMIT;"); });
     std::vector<HardwareEntry> entries;
     SQLite::Statement query(
         db_,
         "SELECT uuid, type, "
-        "hardware_id, software_id, activated FROM hardware");
+        "hardware_id, software_id, activated, priority FROM hardware");
 
     while (query.executeStep()) {
       HardwareEntry entry;
@@ -29,6 +30,7 @@ Hardwares::LoadHardwareList() const {
       entry.hardware_id = query.getColumn(2).getInt();
       entry.software_id = query.getColumn(3).getInt();
       entry.activated = query.getColumn(4).getInt();
+      entry.priority = query.getColumn(5).getInt();
       entries.push_back(entry);
     }
     return entries;
@@ -37,19 +39,20 @@ Hardwares::LoadHardwareList() const {
     return cpp::fail(e.what());
   }
 }
-cpp::result<bool, std::string> Hardwares::AddHardwareEntry(
+cpp::result<bool, std::string> Hardware::AddHardwareEntry(
     const HardwareEntry& new_entry) {
   try {
     SQLite::Statement insert(
         db_,
         "INSERT INTO hardware (uuid, type, "
-        "hardware_id, software_id, activated) VALUES (?, ?, "
-        "?, ?, ?)");
+        "hardware_id, software_id, activated, priority) VALUES (?, ?, "
+        "?, ?, ?, ?)");
     insert.bind(1, new_entry.uuid);
     insert.bind(2, new_entry.type);
     insert.bind(3, new_entry.hardware_id);
     insert.bind(4, new_entry.software_id);
     insert.bind(5, new_entry.activated);
+    insert.bind(6, new_entry.priority);
     insert.exec();
     CTL_INF("Inserted: " << new_entry.ToJsonString());
     return true;
@@ -58,17 +61,19 @@ cpp::result<bool, std::string> Hardwares::AddHardwareEntry(
     return cpp::fail(e.what());
   }
 }
-cpp::result<bool, std::string> Hardwares::UpdateHardwareEntry(
+cpp::result<bool, std::string> Hardware::UpdateHardwareEntry(
     const std::string& id, const HardwareEntry& updated_entry) {
   try {
-    SQLite::Statement upd(db_,
-                          "UPDATE hardware "
-                          "SET hardware_id = ?, software_id = ?, activated = ? "
-                          "WHERE uuid = ?");
+    SQLite::Statement upd(
+        db_,
+        "UPDATE hardware "
+        "SET hardware_id = ?, software_id = ?, activated = ?, priority = ? "
+        "WHERE uuid = ?");
     upd.bind(1, updated_entry.hardware_id);
     upd.bind(2, updated_entry.software_id);
     upd.bind(3, updated_entry.activated);
-    upd.bind(4, id);
+    upd.bind(4, updated_entry.priority);
+    upd.bind(5, id);
     if (upd.exec() == 1) {
       CTL_INF("Updated: " << updated_entry.ToJsonString());
       return true;
@@ -79,7 +84,7 @@ cpp::result<bool, std::string> Hardwares::UpdateHardwareEntry(
   }
 }
 
-cpp::result<bool, std::string> Hardwares::DeleteHardwareEntry(
+cpp::result<bool, std::string> Hardware::DeleteHardwareEntry(
     const std::string& id) {
   try {
     SQLite::Statement del(db_, "DELETE from hardware WHERE uuid = ?");

diff --git a/engine/database/hardware.h b/engine/database/hardware.h
@@ -4,8 +4,8 @@
 #include <trantor/utils/Logger.h>
 #include <string>
 #include <vector>
-#include "utils/result.hpp"
 #include "utils/json_helper.h"
+#include "utils/result.hpp"
 
 namespace cortex::db {
 struct HardwareEntry {
@@ -14,33 +14,34 @@ struct HardwareEntry {
   int hardware_id;
   int software_id;
   bool activated;
+  int priority;
   std::string ToJsonString() const {
     Json::Value root;
     root["uuid"] = uuid;
     root["type"] = type;
     root["hardware_id"] = hardware_id;
     root["software_id"] = software_id;
     root["activated"] = activated;
+    root["priority"] = priority;
     return json_helper::DumpJsonString(root);
   }
 };
 
-class Hardwares {
+class Hardware {
 
  private:
   SQLite::Database& db_;
 
-
  public:
-  Hardwares();
-  Hardwares(SQLite::Database& db);
-  ~Hardwares();
+  Hardware();
+  Hardware(SQLite::Database& db);
+  ~Hardware();
 
   cpp::result<std::vector<HardwareEntry>, std::string> LoadHardwareList() const;
-  cpp::result<bool, std::string> AddHardwareEntry(const HardwareEntry& new_entry);
+  cpp::result<bool, std::string> AddHardwareEntry(
+      const HardwareEntry& new_entry);
   cpp::result<bool, std::string> UpdateHardwareEntry(
       const std::string& id, const HardwareEntry& updated_entry);
-  cpp::result<bool, std::string> DeleteHardwareEntry(
-      const std::string& id);
+  cpp::result<bool, std::string> DeleteHardwareEntry(const std::string& id);
 };
 }  // namespace cortex::db
diff --git a/engine/migrations/db_helper.h b/engine/migrations/db_helper.h
@@ -4,23 +4,28 @@
 namespace cortex::mgr {
 #include <iostream>
 #include <stdexcept>
-#include <vector>
 #include <string>
+#include <vector>
 
-inline bool ColumnExists(SQLite::Database& db, const std::string& table_name, const std::string& column_name) {
-    try {
-        SQLite::Statement query(db, "SELECT " + column_name + " FROM " + table_name + " LIMIT 0");
-        return true;
-    } catch (std::exception&) {
-        return false;
-    }
+inline bool ColumnExists(SQLite::Database& db, const std::string& table_name,
+                         const std::string& column_name) {
+  try {
+    SQLite::Statement query(
+        db, "SELECT " + column_name + " FROM " + table_name + " LIMIT 0");
+    return true;
+  } catch (std::exception&) {
+    return false;
+  }
 }
 
-inline void AddColumnIfNotExists(SQLite::Database& db, const std::string& table_name, 
-                              const std::string& column_name, const std::string& column_type) {
-    if (!ColumnExists(db, table_name, column_name)) {
-        std::string sql = "ALTER TABLE " + table_name + " ADD COLUMN " + column_name + " " + column_type;
-        db.exec(sql);
-    }
+inline void AddColumnIfNotExists(SQLite::Database& db,
+                                 const std::string& table_name,
+                                 const std::string& column_name,
+                                 const std::string& column_type) {
+  if (!ColumnExists(db, table_name, column_name)) {
+    std::string sql = "ALTER TABLE " + table_name + " ADD COLUMN " +
+                      column_name + " " + column_type;
+    db.exec(sql);
+  }
 }
-}
+}  // namespace cortex::mgr
diff --git a/engine/migrations/migration_manager.cc b/engine/migrations/migration_manager.cc
@@ -7,7 +7,7 @@
 #include "utils/widechar_conv.h"
 #include "v0/migration.h"
 #include "v1/migration.h"
-
+#include "v2/migration.h"
 namespace cortex::migr {
 
 namespace {
@@ -141,9 +141,11 @@ cpp::result<bool, std::string> MigrationManager::DoUpFolderStructure(
   switch (version) {
     case 0:
       return v0::MigrateFolderStructureUp();
-      break;
     case 1:
       return v1::MigrateFolderStructureUp();
+    case 2:
+      return v2::MigrateFolderStructureUp();
+
       break;
 
     default:
@@ -155,9 +157,10 @@ cpp::result<bool, std::string> MigrationManager::DoDownFolderStructure(
   switch (version) {
     case 0:
       return v0::MigrateFolderStructureDown();
-      break;
     case 1:
       return v1::MigrateFolderStructureDown();
+    case 2:
+      return v2::MigrateFolderStructureDown();
       break;
 
     default:
@@ -191,9 +194,10 @@ cpp::result<bool, std::string> MigrationManager::DoUpDB(int version) {
   switch (version) {
     case 0:
       return v0::MigrateDBUp(db_);
-      break;
     case 1:
       return v1::MigrateDBUp(db_);
+    case 2:
+      return v2::MigrateDBUp(db_);
       break;
 
     default:
@@ -205,9 +209,10 @@ cpp::result<bool, std::string> MigrationManager::DoDownDB(int version) {
   switch (version) {
     case 0:
       return v0::MigrateDBDown(db_);
-      break;
     case 1:
       return v1::MigrateDBDown(db_);
+    case 2:
+      return v2::MigrateDBDown(db_);
       break;
 
     default:

diff --git a/engine/migrations/schema_version.h b/engine/migrations/schema_version.h
@@ -1,4 +1,5 @@
 #pragma once
 
 //Track the current schema version
-#define SCHEMA_VERSION 1
+#define SCHEMA_VERSION 2
+