From 6128b58cbd2d94a34896d10855a3ebc0b37742a7 Mon Sep 17 00:00:00 2001
From: Joachim Meyer <jmeyer@cs.uni-saarland.de>
Date: Fri, 16 Dec 2022 13:42:04 +0100
Subject: [PATCH] Fix `AssignedGPUs` parsing.

---
 scripts/condor_status_to_gpu_map.py |  6 +++++-
 src/htcondor_cc_sync_plugin.cpp     | 16 ++++++++--------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/scripts/condor_status_to_gpu_map.py b/scripts/condor_status_to_gpu_map.py
index 2f7e0c8..4d47121 100644
--- a/scripts/condor_status_to_gpu_map.py
+++ b/scripts/condor_status_to_gpu_map.py
@@ -17,10 +17,14 @@ if __name__ == "__main__":
 
     slot_gpu_map = {}
     for slot in condor_status:
-        machine = slot["Machine"]
+        machine = slot["Machine"].split('.')[0]
         gpu_map = {}
         if machine in slot_gpu_map:
             gpu_map = slot_gpu_map[machine]
+        else:
+            slot_gpu_map[machine] = {}
+        if not "AssignedGPUs" in slot:
+            continue
         gpus = slot["AssignedGPUs"].split(',')
         for gpu_id in gpus:
             gpu = slot["GPUs_" + gpu_id.strip().replace("-", "_")]
diff --git a/src/htcondor_cc_sync_plugin.cpp b/src/htcondor_cc_sync_plugin.cpp
index c3e3758..4061e53 100644
--- a/src/htcondor_cc_sync_plugin.cpp
+++ b/src/htcondor_cc_sync_plugin.cpp
@@ -267,12 +267,19 @@ std::string getRemoteHost(std::string_view slot) {
   return {++atIt, subIt};
 }
 
+std::string removeQuotes(const std::string &value) {
+  if (value[0] == '"' && value.back() == '"') {
+    return value.substr(1, value.size() - 2);
+  }
+  return value;
+}
+
 std::vector<std::string>
 CCSyncPlugin::getAccelerators(const CondorJob &job,
                               const std::string &hostname) {
   if (auto hostIt = gpuMap.find(hostname); hostIt != gpuMap.end()) {
     if (auto gpuIt = job.find("AssignedGPUs"); gpuIt != job.end()) {
-      std::istringstream gpuStream{gpuIt->second};
+      std::istringstream gpuStream{removeQuotes(gpuIt->second)};
       std::vector<std::string> gpus;
       std::string gpu;
       while (std::getline(gpuStream, gpu, ',')) {
@@ -313,13 +320,6 @@ bool jobHasRequiredClassAds(const CondorJob &job) {
          jobHasClassAd(job, "EnteredCurrentStatus");
 }
 
-std::string removeQuotes(const std::string &value) {
-  if (value[0] == '"' && value.back() == '"') {
-    return value.substr(1, value.size() - 2);
-  }
-  return value;
-}
-
 void CCSyncPlugin::sendPostRequest(const std::string &route,
                                    const std::string &body) const noexcept {
   dprintf(D_VERBOSE, "POST body: %s\n", body.c_str());