Initial commit to speed up submission

Author: ethany-nv

src/service/core/workflow/objects.py

@@ -18,6 +18,7 @@
 
 import collections
 import datetime
+import json
 import math
 from typing import Any, Dict, List, NamedTuple, Optional, Protocol, Set
 import yaml
@@ -1040,6 +1041,7 @@ def send_submit_workflow_to_queue(self,
 
         # Write workflow and group objects to the database
         workflow_obj.insert_to_db()
+        task_entries: list[tuple] = []
         for group_obj in workflow_obj.groups:
             group_obj.workflow_id_internal = workflow_obj.workflow_id
             group_obj.spec = \
@@ -1048,13 +1050,23 @@ def send_submit_workflow_to_queue(self,
             group_obj.insert_to_db()
             for task_obj, task_obj_spec in zip(group_obj.tasks, group_obj.spec.tasks):
                 task_obj.workflow_id_internal = workflow_obj.workflow_id
-                task_obj.insert_to_db(
-                    gpu_count=task_obj_spec.resources.gpu or 0,
-                    cpu_count=task_obj_spec.resources.cpu or 0,
-                    disk_count=common.convert_resource_value_str(
+                workflow_uuid = task_obj.workflow_uuid if task_obj.workflow_uuid else ''
+                task_entries.append((
+                    task_obj.workflow_id_internal, task_obj.name, task_obj.group_name,
+                    task_obj.task_db_key, task_obj.retry_id, task_obj.task_uuid,
+                    task.TaskGroupStatus.WAITING.name,
+                    kb_objects.construct_pod_name(workflow_uuid, task_obj.task_uuid),
+                    None,
+                    task_obj_spec.resources.gpu or 0,
+                    task_obj_spec.resources.cpu or 0,
+                    common.convert_resource_value_str(
                         task_obj_spec.resources.storage or '0', 'GiB'),
-                    memory_count=common.convert_resource_value_str(
-                        task_obj_spec.resources.memory or '0', 'GiB'))
+                    common.convert_resource_value_str(
+                        task_obj_spec.resources.memory or '0', 'GiB'),
+                    json.dumps(task_obj.exit_actions, default=common.pydantic_encoder),
+                    task_obj.lead,
+                ))
+        task.Task.batch_insert_to_db(postgres, task_entries)
 
         logs = f'{service_url}/api/workflow/{workflow_obj.workflow_id}/logs'
         context = WorkflowServiceContext.get()

src/utils/job/jobs.py

@@ -222,6 +222,7 @@ def execute(self, context: JobExecutionContext,
 
         self.workflow_id = workflow_obj.workflow_id
 
+        task_entries: list[tuple] = []
         for group_obj in workflow_obj.groups:
             group_obj.workflow_id_internal = workflow_obj.workflow_id
             group_obj.spec = \
@@ -230,19 +231,23 @@ def execute(self, context: JobExecutionContext,
             group_obj.insert_to_db()
             for task_obj, task_obj_spec in zip(group_obj.tasks, group_obj.spec.tasks):
                 task_obj.workflow_id_internal = workflow_obj.workflow_id
-                task_obj.insert_to_db(
-                    gpu_count=task_obj_spec.resources.gpu or 0,
-                    cpu_count=task_obj_spec.resources.cpu or 0,
-                    disk_count=common.convert_resource_value_str(
+                workflow_uuid = task_obj.workflow_uuid if task_obj.workflow_uuid else ''
+                task_entries.append((
+                    task_obj.workflow_id_internal, task_obj.name, task_obj.group_name,
+                    task_obj.task_db_key, task_obj.retry_id, task_obj.task_uuid,
+                    task.TaskGroupStatus.WAITING.name,
+                    kb_objects.construct_pod_name(workflow_uuid, task_obj.task_uuid),
+                    None,
+                    task_obj_spec.resources.gpu or 0,
+                    task_obj_spec.resources.cpu or 0,
+                    common.convert_resource_value_str(
                         task_obj_spec.resources.storage or '0', 'GiB'),
-                    memory_count=common.convert_resource_value_str(
-                        task_obj_spec.resources.memory or '0', 'GiB'))
-
-                current_timestamp = datetime.datetime.now()
-                time_elapsed = last_timestamp - current_timestamp
-                if time_elapsed > progress_iter_freq:
-                    progress_writer.report_progress()
-                    last_timestamp = current_timestamp
+                    common.convert_resource_value_str(
+                        task_obj_spec.resources.memory or '0', 'GiB'),
+                    json.dumps(task_obj.exit_actions, default=common.pydantic_encoder),
+                    task_obj.lead,
+                ))
+        task.Task.batch_insert_to_db(postgres, task_entries)
         progress_writer.report_progress()
 
         # Fetch workflow_obj to get latest info

src/utils/job/task.py

@@ -1041,6 +1041,42 @@ class Task(pydantic.BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
+    @staticmethod
+    def batch_insert_to_db(
+        database: connectors.PostgresConnector,
+        task_entries: List[Tuple],
+        batch_size: int = 100,
+    ):
+        """Batch-insert multiple tasks in a single query.
+
+        Args:
+            database: The Postgres connector instance.
+            task_entries: List of tuples, each containing the full set of
+                column values for a single task row (same order as insert_to_db).
+            batch_size: Maximum number of rows per INSERT statement.
+        """
+        if not task_entries:
+            return
+
+        for i in range(0, len(task_entries), batch_size):
+            chunk = task_entries[i:i + batch_size]
+            values_clause = ','.join(
+                ['(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)']
+                * len(chunk)
+            )
+            flat_args: List[Any] = []
+            for entry in chunk:
+                flat_args.extend(entry)
+
+            insert_cmd = f'''
+                INSERT INTO tasks
+                (workflow_id, name, group_name, task_db_key, retry_id, task_uuid,
+                 status, pod_name, failure_message, gpu_count, cpu_count,
+                 disk_count, memory_count, exit_actions, lead)
+                VALUES {values_clause} ON CONFLICT DO NOTHING;
+            '''
+            database.execute_commit_command(insert_cmd, tuple(flat_args))
+
     def insert_to_db(self, gpu_count: float, cpu_count: float, disk_count: float,
                      memory_count: float, status: TaskGroupStatus = TaskGroupStatus.WAITING,
                      failure_message: str | None = None):

src/utils/job/workflow.py

@@ -530,6 +530,10 @@ def validate_credentials(self, user: str):
         workflow_config = database.get_workflow_configs()
         dataset_config = database.get_dataset_configs()
         image_hash_map: Dict[str, str] = {}
+        default_user_bucket = connectors.UserProfile.fetch_from_db(database, user).bucket
+        default_service_bucket = dataset_config.default_bucket
+        user_creds = database.get_all_data_creds(user)
+        generic_cred_cache: Dict[str, Any] = {}
         for group in self.groups:
             for group_task in group.tasks:
                 response = self.validate_registry(
@@ -562,13 +566,18 @@ def validate_credentials(self, user: str):
                 self.validate_data(
                     user, dataset_config, group_task, seen_data_input,
                     seen_data_output, workflow_config.credential_config.disable_data_validation,
-                    seen_bucket_input, seen_bucket_output)
-                self.validate_generic_cred(user, database, group_task)
+                    seen_bucket_input, seen_bucket_output,
+                    default_user_bucket, default_service_bucket, user_creds)
+                self.validate_generic_cred(user, database, group_task,
+                                           generic_cred_cache)
 
     def validate_generic_cred(self, user: str, database: connectors.PostgresConnector,
-                              group_task: task.TaskSpec):
+                              group_task: task.TaskSpec,
+                              generic_cred_cache: Dict[str, Any]):
         for cred_name, cred_map in group_task.credentials.items():
-            payload = database.get_generic_cred(user, cred_name)
+            if cred_name not in generic_cred_cache:
+                generic_cred_cache[cred_name] = database.get_generic_cred(user, cred_name)
+            payload = generic_cred_cache[cred_name]
             if isinstance(cred_map, str):
                 continue
             elif isinstance(cred_map, Dict):
@@ -622,12 +631,10 @@ def validate_registry(self, user: str,
     def validate_data(self, user: str, dataset_config: connectors.DatasetConfig,
                       group_task: task.TaskSpec, seen_uri_input: Set[str],
                       seen_uri_output: Set[str], disabled_data: List[str],
-                      seen_bucket_input: Set[str], seen_bucket_output: Set[str]):
-
-        postgres = connectors.PostgresConnector.get_instance()
-        default_user_bucket = connectors.UserProfile.fetch_from_db(postgres, user).bucket
-        default_service_bucket = postgres.get_dataset_configs().default_bucket
-        user_creds = postgres.get_all_data_creds(user)
+                      seen_bucket_input: Set[str], seen_bucket_output: Set[str],
+                      default_user_bucket: Optional[str],
+                      default_service_bucket: str,
+                      user_creds: Dict[str, Any]):
 
         def _validate_input_output(data_spec: Union[task.InputType, task.OutputType, task.TaskKPI],
                                    is_input: bool):