Merge pull request #665 from NVIDIA/xutongr/time

Skip stale resource messages in multi-worker processing

Author: xutongNV

deployments/charts/service/migrations/005_v6_3_0_schema.json

@@ -0,0 +1,24 @@
+{
+  "operations": [
+    {
+      "add_column": {
+        "table": "resources",
+        "column": {
+          "name": "last_updated",
+          "type": "TIMESTAMP",
+          "nullable": true
+        }
+      }
+    },
+    {
+      "add_column": {
+        "table": "resources",
+        "column": {
+          "name": "last_usage_updated",
+          "type": "TIMESTAMP",
+          "nullable": true
+        }
+      }
+    }
+  ]
+}

src/service/agent/helpers.py

@@ -206,25 +206,30 @@ def queue_update_group_job(message: backend_messages.UpdatePodBody):
     update_task.send_job_to_queue()
 
 
-def update_resource(backend: str, message: backend_messages.UpdateNodeBody):
+def update_resource(backend: str, message: backend_messages.UpdateNodeBody,
+                    message_timestamp: datetime.datetime | None = None):
     context = objects.WorkflowServiceContext.get()
     postgres = context.database
     # If delete flag is set, delegate to delete_resource and ignore all other fields
     if message.delete:
         delete_resource(backend, message)
         return
 
+    timestamp = message_timestamp if message_timestamp else common.current_time()
+
     commit_cmd = '''
         INSERT INTO resources
         (name, backend, available, allocatable_fields, label_fields, usage_fields,
-         non_workflow_usage_fields, taints, conditions)
-        VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb[], %s::text[])
+         non_workflow_usage_fields, taints, conditions, last_updated)
+        VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb[], %s::text[], %s)
         ON CONFLICT (name, backend) DO UPDATE SET
         available = %s,
         allocatable_fields = %s,
         label_fields = %s,
         taints = %s::jsonb[],
-        conditions = %s::text[]
+        conditions = %s::text[],
+        last_updated = %s
+        WHERE resources.last_updated IS NULL OR resources.last_updated < %s
     '''
 
     resource = {
@@ -249,48 +254,62 @@ def update_resource(backend: str, message: backend_messages.UpdateNodeBody):
         resource['non_workflow_usage_fields'],
         resource['taints'],
         resource['conditions'],
+        timestamp,
         resource['available'],
         resource['allocatable_fields'],
         resource['label_fields'],
         resource['taints'],
-        resource['conditions']
+        resource['conditions'],
+        timestamp,
+        timestamp,
     )
 
-    postgres.execute_commit_command(commit_cmd, columns)
-    pool_config = connectors.fetch_verbose_pool_config(postgres, backend)
-    resource_entry = workflow.ResourcesEntry(hostname=message.hostname,
-                                             label_fields=message.label_fields,
-                                             backend=backend, taints=message.taints,
-                                             # Dummy placeholder values below
-                                             exposed_fields={},
-                                             usage_fields={},
-                                             non_workflow_usage_fields={},
-                                             allocatable_fields={},
-                                             pool_platform_labels={},
-                                             resource_type=connectors.BackendResourceType.SHARED,
-                                             conditions=message.conditions)
-    config_helpers.update_node_pool_platform(resource_entry, backend, pool_config)
-
-
-def update_resource_usage(backend: str, message: backend_messages.UpdateNodeUsageBody):
+    rowcount = postgres.execute_commit_command(commit_cmd, columns)
+    if rowcount > 0:
+        pool_config = connectors.fetch_verbose_pool_config(postgres, backend)
+        resource_type = connectors.BackendResourceType.SHARED
+        resource_entry = workflow.ResourcesEntry(
+            hostname=message.hostname,
+            label_fields=message.label_fields,
+            backend=backend, taints=message.taints,
+            # Dummy placeholder values below
+            exposed_fields={},
+            usage_fields={},
+            non_workflow_usage_fields={},
+            allocatable_fields={},
+            pool_platform_labels={},
+            resource_type=resource_type,
+            conditions=message.conditions)
+        config_helpers.update_node_pool_platform(resource_entry, backend, pool_config)
+
+
+def update_resource_usage(backend: str, message: backend_messages.UpdateNodeUsageBody,
+                          message_timestamp: datetime.datetime | None = None):
     context = objects.WorkflowServiceContext.get()
     postgres = context.database
+    timestamp = message_timestamp if message_timestamp else common.current_time()
+
     commit_cmd = '''
         INSERT INTO resources
-        (name, backend, usage_fields, non_workflow_usage_fields)
-        VALUES (%s, %s, %s, %s)
+        (name, backend, usage_fields, non_workflow_usage_fields, last_usage_updated)
+        VALUES (%s, %s, %s, %s, %s)
         ON CONFLICT (name, backend) DO UPDATE SET
         usage_fields = %s,
-        non_workflow_usage_fields = %s
+        non_workflow_usage_fields = %s,
+        last_usage_updated = %s
+        WHERE resources.last_usage_updated IS NULL OR resources.last_usage_updated < %s
     '''
 
     columns = (
         message.hostname,
         backend,
         postgres.encode_hstore(message.usage_fields),
         postgres.encode_hstore(message.non_workflow_usage_fields),
+        timestamp,
         postgres.encode_hstore(message.usage_fields),
-        postgres.encode_hstore(message.non_workflow_usage_fields)
+        postgres.encode_hstore(message.non_workflow_usage_fields),
+        timestamp,
+        timestamp,
     )
 
     postgres.execute_commit_command(commit_cmd, columns)
@@ -331,9 +350,11 @@ def delete_resource(backend: str,
 
 
 def clean_resources(backend: str,
-                    message: backend_messages.NodeInventoryBody):
+                    message: backend_messages.NodeInventoryBody,
+                    message_timestamp: datetime.datetime | None = None):
     context = objects.WorkflowServiceContext.get()
     postgres = context.database
+    timestamp = message_timestamp if message_timestamp else common.current_time()
 
     # Track all resources from resources table
     cmd = 'SELECT name FROM resources where backend = %s'
@@ -343,8 +364,13 @@ def clean_resources(backend: str,
     # Find nodes that exist in the database but not in the message
     stale_nodes = db_node_names - set(message.hostnames)
     if stale_nodes:
-        commit_cmd = 'DELETE FROM resources WHERE name IN %s and backend = %s'
-        postgres.execute_commit_command(commit_cmd, (tuple(stale_nodes), backend))
+        # Only delete nodes whose last_updated is older than this inventory message,
+        # preventing a stale inventory from removing nodes added after it was created
+        commit_cmd = '''
+            DELETE FROM resources WHERE name IN %s AND backend = %s
+            AND (last_updated IS NULL OR last_updated < %s)
+        '''
+        postgres.execute_commit_command(commit_cmd, (tuple(stale_nodes), backend, timestamp))
 
 
 def clean_tasks(backend: str,
@@ -622,11 +648,14 @@ async def get_messages():
                 elif message_body.monitor_pod:
                     create_monitor_job(message_body.monitor_pod)
                 elif message_body.update_node:
-                    update_resource(name, message_body.update_node)
+                    update_resource(name, message_body.update_node,
+                                    message_timestamp=message.timestamp)
                 elif message_body.update_node_usage:
-                    update_resource_usage(name, message_body.update_node_usage)
+                    update_resource_usage(name, message_body.update_node_usage,
+                                          message_timestamp=message.timestamp)
                 elif message_body.node_inventory:
-                    clean_resources(name, message_body.node_inventory)
+                    clean_resources(name, message_body.node_inventory,
+                                    message_timestamp=message.timestamp)
                 elif message_body.task_list:
                     clean_tasks(name, message_body.task_list)
                 elif message_body.heartbeat:

src/service/agent/message_worker.py

@@ -227,12 +227,15 @@ def process_message(self, message_id: str, message_json: str, backend_name: str)
                 helpers.send_pod_conditions(message_body.update_pod,
                                             self.workflow_config.max_event_log_lines)
             elif message_body.update_node:
-                helpers.update_resource(backend_name, message_body.update_node)
+                helpers.update_resource(backend_name, message_body.update_node,
+                                        message_timestamp=message.timestamp)
             elif message_body.update_node_usage:
                 helpers.update_resource_usage(
-                    backend_name, message_body.update_node_usage)
+                    backend_name, message_body.update_node_usage,
+                    message_timestamp=message.timestamp)
             elif message_body.node_inventory:
-                helpers.clean_resources(backend_name, message_body.node_inventory)
+                helpers.clean_resources(backend_name, message_body.node_inventory,
+                                        message_timestamp=message.timestamp)
             elif message_body.pod_event:
                 helpers.send_pod_event(
                     message_body.pod_event, self.workflow_config.max_event_log_lines)

src/utils/connectors/postgres.py

@@ -475,14 +475,17 @@ def execute_fetch_command(self, command: str,
                     cur.close()
 
     @retry
-    def execute_commit_command(self, command: str, args: Tuple):
+    def execute_commit_command(self, command: str, args: Tuple) -> int:
         """
         Connects and executes a command that updates the database.
 
         Args:
             command (str): The command to execute.
             args (Tuple): Any args for the command.
 
+        Returns:
+            int: The number of rows affected by the command.
+
         Raises:
             OSMODatabaseError: Error while executing the database command.
         """
@@ -491,8 +494,10 @@ def execute_commit_command(self, command: str, args: Tuple):
             try:
                 cur = conn.cursor()
                 cur.execute(command, args)
+                rowcount = cur.rowcount
                 cur.close()
                 conn.commit()
+                return rowcount
             except (psycopg2.DatabaseError, psycopg2.InterfaceError) as error:
                 try:
                     if cur is not None:
@@ -994,6 +999,8 @@ def _init_tables(self):
                 usage_fields HSTORE,
                 non_workflow_usage_fields HSTORE,
                 conditions TEXT[],
+                last_updated TIMESTAMP,
+                last_usage_updated TIMESTAMP,
                 PRIMARY KEY (name, backend)
             );
         '''