feat: Set pilot stamp in VM env

DIRACGrid · Oct 11, 2022 · 77ae2f0 · 77ae2f0
1 parent d374644
commit 77ae2f0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/DIRAC/Resources/Computing/CloudComputingElement.py b/src/DIRAC/Resources/Computing/CloudComputingElement.py
@@ -286,7 +286,7 @@ def _getSSHKeyID(self):
         """
         return self.ceParameters.get("Instance_SSHKey", None)
 
-    def _getMetadata(self, executableFile):
+    def _getMetadata(self, executableFile, pilotStamp=""):
         """Builds metadata from configuration system, cloudinit template
          and dirac pilot job wrapper
 
@@ -309,6 +309,8 @@ def _getMetadata(self, executableFile):
                 filedef["content"] = self.proxy
             elif filedef["content"] == "EXECUTABLE_STR":
                 filedef["content"] = exe_str
+            elif "STAMP_STR" in filedef["content"]:
+                filedef["content"] = filedef["content"].replace("STAMP_STR", pilotStamp)
         ext_packages = self.ceParameters.get("Context_ExtPackages", None)
         if ext_packages:
             packages = [x.strip() for x in ext_packages.split(",")]
@@ -427,7 +429,6 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1):
         instParams["image"] = self._getImage()
         instParams["size"] = self._getFlavor()
         instParams["ex_keyname"] = self._getSSHKeyID()
-        instParams["ex_userdata"] = self._getMetadata(executableFile)
         instParams["ex_config_drive"] = True
 
         driver = self._getDriver()
@@ -437,6 +438,7 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1):
             instRandom = str(uuid.uuid4()).upper()[:8]
             instName = VM_NAME_PREFIX + instRandom
             instParams["name"] = instName
+            instParams["ex_userdata"] = self._getMetadata(executableFile, instRandom)
             try:
                 node = driver.create_node(**instParams)
             except Exception as err:
@@ -445,7 +447,9 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1):
             instIDs.append(VM_ID_PREFIX + node.id)
         if not instIDs:
             return S_ERROR("Failed to submit any instances.")
-        return S_OK(instIDs)
+        result = S_OK(instIDs)
+        result["PilotStampDict"] = stampDict
+        return result
 
     def killJob(self, jobIDList):
         """Stops VM instances

diff --git a/src/DIRAC/Resources/Computing/cloudinit.template b/src/DIRAC/Resources/Computing/cloudinit.template
@@ -81,6 +81,7 @@ write_files:
      cd /mnt/dirac
      export X509_USER_PROXY=/mnt/proxy.pem
      export PILOT_UUID="cloud://$(cat /var/lib/cloud/data/instance-id)"
+     export DIRAC_PILOT_STAMP="STAMP_STR"
      bash /mnt/run_pilot.sh &> /mnt/dirac/startup.log
 
 yum_repos: