Track cluster creation time and use it in the DRA create token

Author: heerener

hpc_provisioner/src/hpc_provisioner/aws_queries.py

@@ -17,7 +17,9 @@
 from hpc_provisioner.dynamodb_actions import (
     SubnetAlreadyRegisteredException,
     dynamodb_client,
+    dynamodb_resource,
     free_subnet,
+    get_cluster_by_name,
     get_registered_subnets,
     get_subnet,
     register_subnet,
@@ -418,8 +420,7 @@ def create_dra(
     filesystem_id: str,
     mountpoint: str,
     bucket: str,
-    vlab_id: str,
-    project_id: str,
+    cluster: Cluster,
     writable: bool = False,
 ) -> dict:
     logger.debug(
@@ -439,6 +440,13 @@ def create_dra(
         s3_config["AutoExportPolicy"] = {"Events": ["NEW", "CHANGED", "DELETED"]}
 
     logger.debug(f"s3 config: {s3_config}")
+    dynamo_cluster = get_cluster_by_name(
+        dynamodb_resource=dynamodb_resource(), cluster_name=cluster.name
+    )
+    if not dynamo_cluster:
+        raise RuntimeError(f"Clould not retrieve cluster {cluster.name} from dynamodb")
+    if dynamo_cluster["creation_time"] == 0:
+        raise ValueError(f"Creation time for {cluster.name} is 0; something is wrong")
 
     dra = fsx_client.create_data_repository_association(
         FileSystemId=filesystem_id,
@@ -447,12 +455,12 @@ def create_dra(
         BatchImportMetaDataOnCreate=True,
         ImportedFileChunkSize=1024,
         S3=s3_config,
-        ClientRequestToken=f"{vlab_id}-{project_id}-{mountpoint.split('/')[-1]}",
+        ClientRequestToken=f"{dynamo_cluster['creation_time']}-{cluster.vlab_id[:21]}-{cluster.project_id[:21]}-{mountpoint.split('/')[-1]}",
         Tags=[
             {"Key": "Name", "Value": f"{filesystem_id}-{mountpoint}"},
             {"Key": BILLING_TAG_KEY, "Value": BILLING_TAG_VALUE},
-            {"Key": VLAB_TAG_KEY, "Value": vlab_id},
-            {"Key": PROJECT_TAG_KEY, "Value": project_id},
+            {"Key": VLAB_TAG_KEY, "Value": cluster.vlab_id},
+            {"Key": PROJECT_TAG_KEY, "Value": cluster.project_id},
         ],
     )
 

hpc_provisioner/src/hpc_provisioner/cluster.py

@@ -20,11 +20,13 @@ class Cluster:
     tier: str
     vlab_id: str
     provisioning_launched: bool
+    creation_time: int
 
     def __init__(
         self,
         project_id: str,
         vlab_id: str,
+        creation_time: int,
         tier: str = "debug",
         benchmark: bool = False,
         dev: bool = False,
@@ -45,6 +47,7 @@ def __init__(
         else:
             self.admin_ssh_key_name = self.name
         self.provisioning_launched = provisioning_launched
+        self.creation_time = creation_time
 
     @property
     def name(self):

hpc_provisioner/src/hpc_provisioner/handlers.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import logging.config
+import time
 from importlib.metadata import version
 
 import boto3
@@ -142,7 +143,7 @@ def pcluster_create_request_handler(event, _context=None):
       * precreate ssh keys
     """
 
-    cluster = _get_vlab_query_params(event)
+    cluster = _get_vlab_query_params(event, set_creation_time=True)
     dynamo = dynamodb_resource()
     try:
         register_cluster(dynamo, cluster)
@@ -267,7 +268,13 @@ def pcluster_delete_handler(event, _context=None):
     return response_json(pc_output)
 
 
-def _get_vlab_query_params(incoming_event) -> Cluster:
+def _get_vlab_query_params(incoming_event, set_creation_time=False) -> Cluster:
+    """
+    Retrieve the query parameters from the incoming event and create a Cluster object based on them.
+    If set_creation_time is set, this implies that this is the first time creating this Cluster object
+    (ie. we're being called from pcluster_create_request_handler) - this helps us keep track of
+    when the cluster was requested.
+    """
     logger.debug(f"Getting query params from event {incoming_event}")
     event = copy.deepcopy(incoming_event)
 
@@ -292,13 +299,17 @@ def _get_vlab_query_params(incoming_event) -> Cluster:
                 )
                 if param in event.get("queryStringParameters", {}):
                     params[param] = event["queryStringParameters"][param]
+    # we set creation_time to 0 if this is not the initial request
+    # this shouldn't even be necessary as we only use it in a context where it was retrieved from dynamo
+    # but it will allow us to check for a "suspicious" value if we ever need it somewhere else
     cluster = Cluster(
         project_id=params["project_id"],
         vlab_id=params["vlab_id"],
         tier=params["tier"],
         benchmark=params.get("benchmark", "").lower() == "true",
         dev=params.get("dev", "").lower() == "true",
         include_lustre=params.get("include_lustre", "").lower() == "true",
+        creation_time=int(time.time()) if set_creation_time else 0,
     )
 
     logger.debug(f"Params: {params}")

hpc_provisioner/src/hpc_provisioner/pcluster_manager.py

@@ -193,8 +193,7 @@ def fsx_precreate(cluster: Cluster, filesystems: list) -> bool:
                 filesystem_id=fs["FileSystemId"],
                 mountpoint=dra["mountpoint"],
                 bucket=get_fs_bucket(dra["name"], cluster),
-                vlab_id=cluster.vlab_id,
-                project_id=cluster.project_id,
+                cluste=cluster,
                 writable=dra["writable"],
             )
             return True