[Draft] Update all stages to use new IngestControlMessage class instead of Morpheus ControlMessage

api/src/nv_ingest_api/primitives/control_message_task.py

@@ -1,10 +1,12 @@
+from uuid import UUID
+
 from pydantic import BaseModel, Field, ConfigDict
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 
 class ControlMessageTask(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    name: str
-    id: str
+    type: str
+    id: Union[str, UUID]
     properties: Dict[str, Any] = Field(default_factory=dict)

api/src/nv_ingest_api/primitives/ingest_control_message.py

@@ -12,6 +12,46 @@
 logger = logging.getLogger(__name__)
 
 
+def remove_task_by_type(ctrl_msg, task: str):
+    """
+    Remove a task from the control message by matching its type.
+
+    This function iterates over the tasks in the control message, and if it finds a task
+    whose type matches the provided task string, it removes that task (using its unique id)
+    and returns the task's properties.
+
+    Parameters
+    ----------
+    ctrl_msg : IngestControlMessage
+        The control message from which to remove the task.
+    task : str
+        The task type to remove.
+
+    Returns
+    -------
+    dict
+        The properties of the removed task.
+
+    Raises
+    ------
+    ValueError
+        If no task with the given type is found.
+    """
+    task_obj = None
+    for t in ctrl_msg.get_tasks():
+        if t.type == task:
+            task_obj = t
+            break
+
+    if task_obj is None:
+        err_msg = f"process_control_message: Task '{task}' not found in control message."
+        logger.error(err_msg)
+        raise ValueError(err_msg)
+
+    removed_task = ctrl_msg.remove_task(task_obj.id)
+    return removed_task.properties
+
+
 class IngestControlMessage:
     """
     A control message class for ingesting tasks and managing associated metadata,
@@ -53,14 +93,18 @@ def has_task(self, task_id: str) -> bool:
         """
         return task_id in self._tasks
 
-    def remove_task(self, task_id: str) -> None:
+    def remove_task(self, task_id: str) -> ControlMessageTask:
         """
         Remove a task from the control message. Logs a warning if the task does not exist.
         """
         if task_id in self._tasks:
+            _task = self._tasks[task_id]
+
             del self._tasks[task_id]
+
+            return _task
         else:
-            logger.warning(f"Attempted to remove non-existent task with id: {task_id}")
+            raise RuntimeError(f"Attempted to remove non-existent task with id: {task_id}")
 
     def config(self, config: Dict[str, Any] = None) -> Dict[str, Any]:
         """
@@ -154,7 +198,8 @@ def filter_timestamp(self, regex_filter: str) -> Dict[str, datetime]:
         Retrieve timestamps whose keys match the regex filter.
         """
         pattern = re.compile(regex_filter)
-        return {key: ts for key, ts in self._timestamps.items() if pattern.search(key)}
+        timestamps_snapshot = self._timestamps.copy()
+        return {key: ts for key, ts in timestamps_snapshot.items() if pattern.search(key)}
 
     def get_timestamp(self, key: str, fail_if_nonexist: bool = False) -> datetime:
         """
@@ -188,12 +233,14 @@ def set_timestamp(self, key: str, timestamp: Any) -> None:
         """
         if isinstance(timestamp, datetime):
             self._timestamps[key] = timestamp
+
         elif isinstance(timestamp, str):
             try:
                 dt = datetime.fromisoformat(timestamp)
                 self._timestamps[key] = dt
             except ValueError as e:
                 raise ValueError(f"Invalid timestamp format: {timestamp}") from e
+
         else:
             raise ValueError("timestamp must be a datetime object or ISO format string")
 

client/src/nv_ingest_client/cli/util/click.py

@@ -22,17 +22,13 @@
 from nv_ingest_client.primitives.tasks import StoreTask
 from nv_ingest_client.primitives.tasks import VdbUploadTask
 from nv_ingest_client.primitives.tasks.caption import CaptionTaskSchema
-from nv_ingest_client.primitives.tasks.chart_extraction import ChartExtractionSchema
-from nv_ingest_client.primitives.tasks.chart_extraction import ChartExtractionTask
 from nv_ingest_client.primitives.tasks.dedup import DedupTaskSchema
 from nv_ingest_client.primitives.tasks.embed import EmbedTaskSchema
 from nv_ingest_client.primitives.tasks.extract import ExtractTaskSchema
 from nv_ingest_client.primitives.tasks.filter import FilterTaskSchema
 from nv_ingest_client.primitives.tasks.split import SplitTaskSchema
 from nv_ingest_client.primitives.tasks.store import StoreEmbedTaskSchema
 from nv_ingest_client.primitives.tasks.store import StoreTaskSchema
-from nv_ingest_client.primitives.tasks.table_extraction import TableExtractionSchema
-from nv_ingest_client.primitives.tasks.table_extraction import TableExtractionTask
 from nv_ingest_client.primitives.tasks.vdb_upload import VdbUploadTaskSchema
 from nv_ingest_client.util.util import generate_matching_files
 
@@ -115,15 +111,6 @@ def click_validate_task(ctx, param, value):
                 task_options = check_schema(ExtractTaskSchema, options, task_id, json_options)
                 new_task_id = f"{task_id}_{task_options.document_type}"
                 new_task = [(new_task_id, ExtractTask(**task_options.model_dump()))]
-
-                if task_options.extract_tables is True:
-                    subtask_options = check_schema(TableExtractionSchema, {}, "table_data_extract", "{}")
-                    new_task.append(("table_data_extract", TableExtractionTask(**subtask_options.model_dump())))
-
-                if task_options.extract_charts is True:
-                    subtask_options = check_schema(ChartExtractionSchema, {}, "chart_data_extract", "{}")
-                    new_task.append(("chart_data_extract", ChartExtractionTask(**subtask_options.model_dump())))
-
             elif task_id == "store":
                 task_options = check_schema(StoreTaskSchema, options, task_id, json_options)
                 new_task_id = f"{task_id}"

docker-compose.yaml

@@ -76,7 +76,7 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              device_ids: ["1"]
+              device_ids: ["0"]
               capabilities: [gpu]
     runtime: nvidia
     profiles:
@@ -165,7 +165,7 @@ services:
     cap_add:
       - sys_nice
     environment:
-      - CUDA_VISIBLE_DEVICES=0
+      - CUDA_VISIBLE_DEVICES=-1
       - EMBEDDING_NIM_MODEL_NAME=${EMBEDDING_NIM_MODEL_NAME:-nvidia/llama-3.2-nv-embedqa-1b-v2}
       - INGEST_LOG_LEVEL=DEFAULT
       # Message client for development

src/microservice_entrypoint.py

@@ -4,7 +4,7 @@
 
 import json
 
-from morpheus.config import Config
+from morpheus.config import Config, ExecutionMode
 from morpheus.config import CppConfig
 from morpheus.config import PipelineModes
 from morpheus.utils.logger import configure_logging
@@ -88,9 +88,10 @@ def cli(
     logging.basicConfig(level=log_level, format="%(asctime)s - %(levelname)s - %(message)s")
     configure_logging(log_level=log_level)
 
-    CppConfig.set_should_use_cpp(use_cpp)
+    CppConfig.set_should_use_cpp(False)
 
     morpheus_pipeline_config = Config()
+    morpheus_pipeline_config.execution_mode = ExecutionMode.CPU
     morpheus_pipeline_config.debug = True if log_level == "DEBUG" else False
     morpheus_pipeline_config.log_level = log_level
     morpheus_pipeline_config.pipeline_batch_size = pipeline_batch_size

src/nv_ingest/modules/extractors/__init__.py

@@ -2,10 +2,8 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from .docx_extractor import DocxExtractorLoaderFactory
 from .pdf_extractor import PDFExtractorLoaderFactory
 
 __all__ = [
     "PDFExtractorLoaderFactory",
-    "DocxExtractorLoaderFactory",
 ]