Coverage for collectra/pipelines/base.py: 76.34%
558 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Core pipeline execution and workflow management for Collectra.
3This module contains the main Collectra workflow class.
5It provides the core functionality for defining, excuting and managing
6workflows with task dependencies.
8The module supports:
9 - Workflow definition and configuration management
10 - Task dependency graph construction and execution
11 - Workflow visualization and rendering
12 - Training and inference execution modes
14Classes:
15 Collectra: Main workflow management class
16"""
18from __future__ import annotations
20__all__ = ["Collectra"]
22import copy
23import datetime
24import json
25import shlex
26import shutil
27import stat
28import sys
29import time
30from pathlib import Path
31from typing import Annotated
33from cappa import Arg
34import graphviz
35import networkx as nx
36import yaml
37from rich.progress import track
39from collectra.utils import change_dir, load_class_from_string, remove_exif, write_yaml
40from collectra.cli import command, group
41from utils.get_types import get_param_types, get_return_type, unpack_types
43from ..commons.base import TaskContext
44from ..logger import get_logger
45from ..tasks.base import (
46 Task,
47 TaskNode,
48)
49from ..types.base import (
50 Artefact,
51 ArtefactNode,
52 NodeStatus,
53)
54from ..types.images import Image, ImageCrop
55from ..types.links import Link
56from ..types.texts import Text
57from .node_graph_manager import NodeGraphManager
59logger = get_logger(__name__)
62class Collectra:
63 """Main workflow management class for Collectra pipelines.
65 This class orchestrates the execution of complex data processing workflows
66 by managing task dependencies, data flow, and execution state.
67 """
69 # Constants
70 RESULTS_FILE = "results.yaml"
71 METADATA_KEY = "collectra_results_metadata"
72 PIPELINE_FILE = "pipeline.yaml"
73 PIPELINE_METADATA_KEY = "collectra_pipeline_metadata"
75 # Graph visualization colors
76 COLOR_PROCESSING = "orange"
77 COLOR_SUCCESS = "green"
78 COLOR_FAILURE = "red"
79 COLOR_FONT = "black"
81 def __init__(
82 self,
83 name: str,
84 ext: str,
85 version: str,
86 path: str | Path = "",
87 pipeline_metadata: dict | None = None,
88 **kwargs,
89 ):
90 """Initialize the Collectra workflow.
92 Args:
93 name (str): Name of the workflow.
94 ext (str): File extension that the workflow accepts.
95 version (str): Version of the workflow.
96 path (str, optional): Path to the workflow files. Defaults to name.
97 **kwargs: Additional configuration parameters.
98 """
99 self.name: str = name
100 self.ext: str = ext
101 self.version: str = version
102 self.pipeline_metadata = {
103 **(pipeline_metadata or {}),
104 "name": name,
105 "ext": ext,
106 "version": version,
107 }
108 self.node_manager: NodeGraphManager = NodeGraphManager()
109 self.path: Path = Path.cwd() / name if not path else Path(path)
110 self.data: dict = kwargs
112 def __call__(self, task_name: str, **kwargs):
113 """Allow instance to be called directly to run workflow."""
114 self.run(task_name, **kwargs)
116 @classmethod
117 def from_file(cls, pipeline_path: str | Path) -> "Collectra":
118 """Load a pipeline from a directory or a pipeline YAML file."""
119 path = Path(pipeline_path).expanduser().resolve()
120 config_path = path / cls.PIPELINE_FILE if path.is_dir() else path
121 with config_path.open() as stream:
122 metadata = yaml.safe_load(stream) or {}
123 initials = metadata.pop(cls.PIPELINE_METADATA_KEY)
124 return cls(
125 initials["name"],
126 initials["ext"],
127 initials["version"],
128 path=config_path.parent,
129 pipeline_metadata=initials,
130 **metadata,
131 )
133 @command(name="run")
134 def cli_run(
135 self,
136 inputs: list[str],
137 task: str = "",
138 output: Path | None = None,
139 single: bool = False,
140 verbose: bool = False,
141 usage: bool = False,
142 render: bool = False,
143 ) -> None:
144 """Run the complete pipeline or start at a named task."""
145 from collectra.logger import setup_logging
146 from collectra.utils import resolve_files, valid_raw_files
148 setup_logging(verbose=verbose)
149 files = resolve_files(inputs, [f".{self.ext}", *valid_raw_files()])
150 if not files:
151 logger.warning("No supported input files found for: %s", ", ".join(inputs))
152 return
153 self.run(
154 task,
155 files=files,
156 output=str(output) if output else None,
157 single=single,
158 usage=usage,
159 render=render,
160 )
162 @command(name="gui")
163 def launch_gui(self, *inputs: Path, debug: bool = False) -> None:
164 """Open the GUI with optional result folders, results.yaml files, or parent directories."""
165 from collectra.gui.backend import start
167 self._ensure_workflow_connected()
168 start(pipeline=self, debug=debug, inputs=inputs)
170 @command(name="partition")
171 def cli_partition(
172 self,
173 inputs: Annotated[
174 list[str],
175 Arg(help="Result folders, results.yaml files, or parent directories."),
176 ],
177 train: Annotated[
178 str,
179 Arg(help="Share of results labelled 'train' (e.g. 80%, 0.8, or 100)."),
180 ] = "",
181 validation: Annotated[
182 str,
183 Arg(help="Share of results labelled 'validation' (e.g. 20%, 0.2, or 25)."),
184 ] = "",
185 val: Annotated[
186 str, Arg(help="Share of results labelled 'val' (e.g. 20%, 0.2, or 25).")
187 ] = "",
188 test: Annotated[
189 str, Arg(help="Share of results labelled 'test' (e.g. 10%, 0.1, or 10).")
190 ] = "",
191 seed: Annotated[
192 int, Arg(help="Random seed used to shuffle results before splitting.")
193 ] = 42,
194 partition: Annotated[
195 list[str] | None,
196 Arg(
197 help="Custom partition as NAME=VALUE (e.g. holdout=5%); may be repeated."
198 ),
199 ] = None,
200 force: Annotated[
201 bool,
202 Arg(
203 help="Override existing partitions instead of refusing to change them."
204 ),
205 ] = False,
206 ) -> None:
207 """Assign result partitions using percentages, fractions, or file counts.
209 For example: partition results --train 80% --validation 20%.
210 Use repeated --partition NAME=VALUE options for custom labels.
211 All values must use the same form (percentages, fractions, or counts).
212 Fails if any result already has a partition unless --force is given.
213 Unassigned results retain their existing partition.
214 """
215 from collectra.partition import partition_results
217 partitions = {
218 label: value
219 for label, value in (
220 ("train", train),
221 ("validation", validation),
222 ("val", val),
223 ("test", test),
224 )
225 if value != ""
226 }
227 for specification in partition or []:
228 label, separator, value = specification.partition("=")
229 if not separator or not label.strip() or not value.strip():
230 raise ValueError("Custom partitions must use NAME=VALUE")
231 if label in partitions:
232 raise ValueError(f"Partition {label!r} was specified more than once")
233 partitions[label] = value
234 partition_results(inputs, self.ext, partitions, seed, force=force)
236 @command(name="install")
237 def cli_install(self, name: str, bin_dir: Path | None = None) -> None:
238 """Install this pipeline as a standalone command."""
239 destination = (bin_dir or Path.home() / ".local" / "bin").expanduser()
240 destination.mkdir(parents=True, exist_ok=True)
241 launcher = destination / name
242 pipeline_path = (self.path / self.PIPELINE_FILE).resolve()
243 code = (
244 "from pathlib import Path\n"
245 "from collectra.main import main\n"
246 f"main(pipeline_path=Path({str(pipeline_path)!r}), program_name={name!r})\n"
247 )
248 # Kernel shebang parsing cannot quote an interpreter path containing
249 # spaces. Let the shell invoke the exact interpreter as one argument.
250 launcher.write_text(
251 "#!/bin/sh\n"
252 f'exec {shlex.join([sys.executable, "-c", code])} "$@"\n',
253 encoding="utf-8",
254 )
255 launcher.chmod(
256 launcher.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
257 )
258 print(f"Installed {name!r} at {launcher}")
260 @command(name="install-app")
261 def cli_install_app(
262 self,
263 *inputs: Path,
264 name: str = "Collectra",
265 app_dir: Path | None = None,
266 icon: Path | None = None,
267 force: bool = False,
268 ) -> Path:
269 """Create a macOS Applications launcher for this pipeline's GUI.
271 Optional inputs are result folders or files opened when the app starts.
272 The default destination is ~/Applications; icon accepts an .icns or PNG file. Pipeline metadata supplies the default icon.
273 """
274 from .app_launcher import install_app
276 pipeline_icon = self.pipeline_metadata.get("icon") or self.pipeline_metadata.get("logo")
277 if icon is None and pipeline_icon:
278 icon = Path(pipeline_icon).expanduser()
279 if not icon.is_absolute():
280 icon = self.path / icon
281 bundle = install_app(
282 self.path / self.PIPELINE_FILE,
283 name=name,
284 app_dir=app_dir,
285 icon=icon,
286 inputs=inputs,
287 force=force,
288 )
289 print(f"Installed {name!r} at {bundle}")
290 return bundle
292 @group(name="task")
293 def cli_tasks(self) -> dict[str, Task]:
294 """Operate on one of the tasks in this pipeline."""
295 self._ensure_workflow_connected()
296 return {
297 node.name: node.get_task() for node in self.node_manager.get_task_nodes()
298 }
300 @group(name="artefact")
301 def cli_artefacts(self) -> dict[str, object]:
302 """Operate on one of the artefacts in this pipeline."""
303 self._ensure_workflow_connected()
304 artefacts: dict[str, object] = {}
305 for node in self.node_manager.get_artefact_nodes():
306 artefact_types = [
307 type_ for type_ in node.types if issubclass(type_, Artefact)
308 ]
309 if not artefact_types:
310 continue
311 # IO inference can add a base and a more-specific type to a node.
312 # Prefer the most specific class so its CLI capabilities win.
313 artefact_type = max(
314 artefact_types,
315 key=lambda type_: len(type_.mro()),
316 )
317 artefacts[node.name] = artefact_type.cli_commands(node, self.ext)
318 return artefacts
320 def task(self, task_name: str) -> dict:
321 """Get a task from the workflow by name.
323 Args:
324 task_name (str): Name of the task to retrieve.
326 Returns:
327 dict | None: Task configuration dictionary or None if not found.
328 """
329 data = self.data.get(task_name, None)
330 if not data:
331 raise ValueError(f"Task {task_name} not found in workflow")
332 data["name"] = task_name
333 return data
335 def _check_is_root(self, node_name: str) -> bool:
336 """Check if a node is a root node (no parents or only data parents)."""
337 node = self.node_manager.resolve_node(node_name)
338 parents = self.node_manager.get_parents(node)
339 for parent in parents:
340 grandparents = self.node_manager.get_parents(parent)
341 if len(grandparents) > 0:
342 return False
343 if isinstance(parent, TaskNode):
344 return False
345 return True
347 def run(self, task_name: str = "", **kwargs):
348 """Run the workflow from the specified task or from all root tasks.
350 Args:
351 task_name: Name of the task to start from. If empty, starts from
352 all root tasks in the workflow.
353 **kwargs: Run options including:
354 - single: Run only the starting task without propagating.
355 - files: List of file paths to process.
356 - usage: Whether to track usage metrics.
357 - render: Whether to render workflow visualization.
358 - output: Output directory for results.
359 """
360 try:
361 self._ensure_workflow_connected()
362 starting_nodes = self._get_starting_nodes(task_name)
363 run_options = self._extract_run_options(kwargs)
364 self._process_files(starting_nodes, run_options)
365 except Exception as e:
366 logger.error("Workflow run failed: %s", e, exc_info=True)
368 def _ensure_workflow_connected(self):
369 """Ensure the workflow graph is connected before execution.
371 If the workflow graph has no nodes, calls connect() to build
372 the graph from the workflow configuration.
373 """
374 if self.node_manager.is_empty():
375 self.connect()
377 def _get_starting_nodes(self, task_name: str) -> list[TaskNode]:
378 """Get the starting nodes for workflow execution.
380 Args:
381 task_name: Specific task to start from, or empty for all root tasks.
383 Returns:
384 List of starting task nodes.
385 """
386 if task_name:
387 return self._get_specific_starting_node(task_name)
388 return self._get_root_starting_nodes()
390 def _get_specific_starting_node(self, task_name: str) -> list[TaskNode]:
391 """Get a specific task node as the starting point."""
392 task_node = self.node_manager.resolve_node(task_name)
393 if not isinstance(task_node, TaskNode):
394 raise TypeError(f"Task {task_name} not found in workflow")
395 return [task_node]
397 def _get_root_starting_nodes(self) -> list[TaskNode]:
398 """Get all root task nodes as starting points."""
399 starting_nodes = []
400 task_nodes = self.node_manager.get_task_nodes()
401 for task_node in task_nodes:
402 if self._check_is_root(task_node.name):
403 starting_nodes.append(task_node)
404 return starting_nodes
406 def _extract_run_options(self, kwargs: dict) -> dict:
407 """Extract and return run options from kwargs."""
408 options = {
409 "single": kwargs.pop("single", False),
410 "files": kwargs.pop("files", []),
411 "usage": kwargs.pop("usage", False),
412 "render": kwargs.pop("render", False),
413 "output": kwargs.pop("output", None),
414 }
415 options["task_overrides"] = kwargs.copy()
416 return options
418 def _get_root_input(self) -> str:
419 """Get the name of the root input artefact node for the workflow."""
420 for artefact_node in self.node_manager.get_artefact_nodes():
421 if self._check_is_root(artefact_node.name):
422 return artefact_node.name
423 raise ValueError("No root input artefact node found in the workflow")
425 def _process_files(self, starting_nodes: list[TaskNode], options: dict):
426 """Process each file in the workflow.
428 Args:
429 starting_nodes: List of nodes to start execution from.
430 options: Execution options (single, files, usage, render, output).
431 """
432 for file_path in options["files"]:
433 try:
434 self._process_single_file(starting_nodes, file_path, options)
435 except Exception as e:
436 logger.error("Error processing file: %s", e, exc_info=True)
438 def _process_single_file(
439 self, starting_nodes: list[TaskNode], file_path: str, options: dict
440 ):
441 """Process a single file through the workflow.
443 Args:
444 starting_nodes: Nodes to start from.
445 file_path: Path to file to process.
446 options: Execution options.
447 """
448 key = (
449 "file"
450 if Path(file_path).suffix == f".{self.ext}"
451 else self._get_root_input()
452 )
453 key, value = self._create_collectra_file(key, file_path, options["output"])
455 # Create task context for this file
456 context = TaskContext(
457 usage_file=(
458 Path.cwd() / str(Path(value) / "usage.yaml")
459 if options["usage"]
460 else None
461 ),
462 render=options["render"],
463 file_path=Path(value),
464 single_run=options["single"],
465 )
467 # Set context for all tasks
468 self._set_task_contexts(context)
470 input_data = {key: value}
471 self._init_input_data(starting_nodes, **input_data)
472 self.render("workflow_run", file=str(value), render=options["render"])
474 self._run_nodes(
475 starting_nodes,
476 single=options["single"],
477 key=key,
478 value=value,
479 render=options["render"],
480 task_overrides=options["task_overrides"],
481 )
483 def _create_collectra_file(
484 self, key: str, value: str | Path, output_dir: str | Path | None = None
485 ) -> tuple[str, str | Path]:
486 """Create or prepare a collectra file structure for processing.
488 Handles three cases:
489 - Existing collectra folder with results.yaml
490 - Single image file
491 - Directory without results.yaml
492 """
493 self._validate_file_inputs(key, value)
494 savef = Path(value)
496 # Case A: Existing collectra folder with results.yaml
497 if self._is_existing_collectra_folder(savef):
498 return key, self._handle_existing_folder(savef, output_dir)
500 # Case B: Single image file
501 if self._is_image_file(savef):
502 savef = self._handle_image_file(savef, output_dir)
504 # Case C: Directory without results.yaml - create results file
505 if savef.is_dir():
506 self._create_results_file(savef, key, value)
508 return "file", savef
510 def _validate_file_inputs(self, key: str, value: str | Path):
511 """Validate input parameters for file creation."""
512 if not key or not value:
513 raise ValueError("Both key and value must be provided for input data.")
515 savef = Path(value)
516 if not savef.exists():
517 raise FileNotFoundError(f"Path {savef} does not exist.")
519 def _is_existing_collectra_folder(self, path: Path) -> bool:
520 """Check if path is an existing collectra folder with results."""
521 return (
522 path.is_dir()
523 and path.suffix == f".{self.ext}"
524 and (path / self.RESULTS_FILE).exists()
525 )
527 def _is_image_file(self, path: Path) -> bool:
528 """Check if path is an image file."""
529 return path.is_file() and path.suffix.lower() in Image.image_types()
531 def _handle_existing_folder(
532 self, savef: Path, output_dir: str | Path | None
533 ) -> Path:
534 """Handle existing collectra folder, optionally copying to output."""
535 if output_dir:
536 output_path = Path(output_dir) / savef.name
537 output_path.mkdir(parents=True, exist_ok=True)
538 shutil.copytree(savef, output_path, dirs_exist_ok=True)
539 savef = output_path
541 self._remove_exif_from_images(savef)
542 return savef
544 def _remove_exif_from_images(self, directory: Path):
545 """Remove EXIF data from all images in directory."""
546 for img_file in directory.glob("*"):
547 if img_file.suffix.lower() in Image.image_types():
548 remove_exif(img_file, img_file)
550 def _handle_image_file(
551 self, file_path: Path, output_dir: str | Path | None
552 ) -> Path:
553 """Handle single image file by creating collectra folder structure."""
554 folder_name = file_path.name.replace(file_path.suffix, f".{self.ext}")
555 savef = (
556 Path(output_dir) / folder_name
557 if output_dir
558 else file_path.parent / folder_name
559 )
560 savef.mkdir(parents=True, exist_ok=True)
562 remove_exif(file_path, savef / file_path.name)
563 return savef
565 def _create_results_file(self, directory: Path, key: str, value: str | Path):
566 """Create a new results.yaml file in the given directory."""
567 with change_dir(directory):
568 results = {
569 self.METADATA_KEY: {
570 "workflow": self.name,
571 "version": self.version,
572 "timestamp": datetime.datetime.now().isoformat(),
573 },
574 key: {
575 "type": "collectra.Image",
576 "id": f"{key}",
577 "data": value.name if isinstance(value, Path) else str(value),
578 },
579 }
580 self._write_results_file(results)
582 def _populate_active_paths(
583 self,
584 nodes: list[TaskNode | ArtefactNode] | list[TaskNode] | list[ArtefactNode],
585 **kwargs,
586 ):
587 """Recursively populate active paths in the workflow with input data.
589 Traverses the workflow graph from the given nodes, processing artefact nodes
590 and propagating values to their children.
592 Args:
593 nodes: List of nodes to start populating from.
594 **kwargs: Key-value pairs of input data to populate into matching nodes.
595 """
596 for node in nodes:
597 if isinstance(node, ArtefactNode):
598 value = kwargs.get(node.name, None)
599 node.process(node.name, value, **kwargs)
601 children = self.node_manager.get_children(node)
602 self._populate_active_paths(children, **kwargs)
604 def _init_input_data(
605 self,
606 starting_nodes: (
607 list[TaskNode | ArtefactNode] | list[TaskNode] | list[ArtefactNode]
608 ),
609 **kwargs,
610 ):
611 """Initialize input data for workflow execution.
613 Resets all nodes, reconnects the workflow, and populates parent data
614 nodes with the provided input values.
616 Args:
617 starting_nodes: List of nodes that will begin workflow execution.
618 **kwargs: Key-value pairs of input data to initialize.
619 """
620 self._reset_nodes()
621 self.connect()
622 for node in starting_nodes:
623 parents = self.node_manager.get_parents_artefact(node)
624 for parent in parents:
625 value = kwargs.get(parent.name, None)
626 parent.process(parent.name, value, **kwargs)
627 self._populate_active_paths(starting_nodes, **kwargs)
628 self._resolve_links()
630 def _resolve_links(self) -> None:
631 """Bind loaded links to artefacts in the active pipeline graph."""
632 artefacts = {
633 item.id: item
634 for node in self.node_manager.get_artefact_nodes()
635 for item in node.items.values()
636 }
637 for node in self.node_manager.get_artefact_nodes():
638 for item in node.items.values():
639 if type(item) is Link:
640 item.bind(artefacts)
642 def _set_task_contexts(self, context: TaskContext):
643 """Set context for all tasks in the workflow.
645 Args:
646 context: TaskContext to apply to all tasks.
647 """
648 for task_node in self.node_manager.get_task_nodes():
649 task = task_node.get_task()
650 if isinstance(task, Task):
651 task.context = context
653 def _reset_nodes(self):
654 """Reset all nodes in the workflow to their initial state.
656 Clears items and ensemble items from artefact nodes and sets all node
657 statuses to NOT_READY.
658 """
659 self.node_manager.reset_all_nodes()
661 def save_run(
662 self, key: str, value: str | Path, artefact_node: ArtefactNode | None = None
663 ):
664 """Save workflow execution results to file.
666 Args:
667 key: Key identifier for the data.
668 value: Path where results should be saved.
669 artefact_node: Specific artefact node to save, or None for all nodes.
670 """
671 if not self._should_save(key, value, artefact_node):
672 return
674 savef = Path(value)
675 if not savef.is_dir():
676 return
678 artefact_nodes = self._get_artefact_nodes_to_save(artefact_node)
679 self._save_results_to_directory(savef, key, value, artefact_nodes)
680 self._log_save_success(savef, artefact_node)
682 def _should_save(
683 self, key: str, value: str | Path, artefact_node: ArtefactNode | None
684 ) -> bool:
685 """Check if saving should proceed."""
686 if not key or not value:
687 logger.info("No save location specified, skipping save.")
688 return False
690 if artefact_node is not None and artefact_node.status != NodeStatus.READY:
691 logger.info("No new data for %s, skipping save.", artefact_node.name)
692 return False
694 return True
696 def _get_artefact_nodes_to_save(
697 self, artefact_node: ArtefactNode | None = None
698 ) -> list[ArtefactNode]:
699 """Get list of artefact nodes to save."""
700 if artefact_node is not None:
701 return [artefact_node]
703 return self.node_manager.get_artefact_nodes()
705 def _save_results_to_directory(
706 self,
707 directory: Path,
708 key: str,
709 value: str | Path,
710 artefact_nodes: list[ArtefactNode],
711 ):
712 """Save results to directory with proper formatting."""
713 with change_dir(directory):
714 results = self._load_or_create_results(key, value)
715 results = self._collect_node_data(results, artefact_nodes)
716 self._write_results_file(results)
718 def _load_or_create_results(self, key: str, value: str | Path) -> dict:
719 """Load existing results file or create new results dictionary."""
720 resultsf = Path(self.RESULTS_FILE)
722 if resultsf.exists():
723 with open(resultsf, "r") as f:
724 results = yaml.safe_load(f)
725 results[self.METADATA_KEY][
726 "timestamp"
727 ] = datetime.datetime.now().isoformat()
728 else:
729 results = {
730 self.METADATA_KEY: {
731 "workflow": self.name,
732 "version": self.version,
733 "timestamp": datetime.datetime.now().isoformat(),
734 },
735 key: {
736 "type": "collectra.Image",
737 "data": value.name if isinstance(value, Path) else str(value),
738 },
739 }
741 return results
743 def _collect_node_data(
744 self, results: dict, artefact_nodes: list[ArtefactNode]
745 ) -> dict:
746 """Collect data from all ready nodes and add to results."""
747 for artefact_node in artefact_nodes:
748 if artefact_node.status != NodeStatus.READY:
749 continue
751 name = artefact_node.name
752 results[name] = self._serialize_node_items(artefact_node)
754 return results
756 def _serialize_node_items(self, artefact_node: ArtefactNode) -> dict | list:
757 """Serialize all items in an artefact node."""
758 serialized_items = [item.serialize() for item in artefact_node.items.values()]
759 return serialized_items[0] if len(serialized_items) == 1 else serialized_items
761 def _write_results_file(self, results: dict):
762 """Write results dictionary to YAML file."""
763 write_yaml(results, self.RESULTS_FILE)
765 def _log_save_success(self, savef: Path, artefact_node: ArtefactNode | None):
766 """Log successful save operation."""
767 if artefact_node:
768 logger.info("Results saved to %s for %s", savef, artefact_node.name)
769 else:
770 logger.info("All results saved to %s", savef)
772 def _run_nodes(
773 self,
774 nodes: list[TaskNode | ArtefactNode] | list[TaskNode] | list[ArtefactNode],
775 *args,
776 **kwargs,
777 ):
778 """Execute a list of nodes (tasks or artefact nodes) in the workflow.
780 This method recursively processes nodes, handling both task execution
781 and artefact node updates. It maintains visualization state and manages
782 child node execution.
783 """
784 single_run = kwargs.get("single", False)
785 render = kwargs.get("render", False)
787 self._mark_nodes_as_processing(nodes)
788 self.render("workflow_run", file=kwargs.get("value", ""), render=render)
790 child_tasks = self._execute_nodes(nodes, args, kwargs, render)
792 if child_tasks and not single_run:
793 self._run_nodes(child_tasks, *args, **kwargs)
795 def _mark_nodes_as_processing(
796 self, nodes: list[TaskNode | ArtefactNode] | list[TaskNode] | list[ArtefactNode]
797 ):
798 """Mark nodes as currently being processed in the graph."""
799 for node in nodes:
800 self.node_manager.set_node_attr(
801 str(node.name), color=self.COLOR_PROCESSING, fontcolor=self.COLOR_FONT
802 )
804 def _execute_nodes(
805 self,
806 nodes: list[TaskNode | ArtefactNode] | list[TaskNode] | list[ArtefactNode],
807 args: tuple,
808 kwargs: dict,
809 render: bool,
810 ) -> list[TaskNode]:
811 """Execute all nodes and collect child tasks.
813 Returns:
814 List of child tasks that are ready to execute.
815 """
816 child_tasks_of_artefact_nodes = dict()
818 for node in nodes:
819 if isinstance(node, TaskNode):
820 self._execute_task_node(node, args, kwargs, render)
821 elif isinstance(node, ArtefactNode):
822 child_tasks = self._execute_artefact_node(node, args, kwargs, render)
823 child_tasks_of_artefact_nodes.update(child_tasks)
825 return self._check_tasks_ready(list(child_tasks_of_artefact_nodes.values()))
827 def _execute_task_node(
828 self, node: TaskNode, args: tuple, kwargs: dict, render: bool
829 ):
830 """Execute a single task node."""
831 if self._check_task_ready(node) != NodeStatus.READY:
832 logger.info(f"[yellow]Task {node.name} is not ready, skipping.[/yellow]")
833 return
835 children = self._get_node_children(node)
836 results = self._run_task(node, **kwargs)
838 self._update_node_status(node, bool(results))
839 self.render("workflow_run", file=kwargs.get("value", ""), render=render)
841 self._run_nodes(children, *results, **kwargs)
843 def _execute_artefact_node(
844 self, node: ArtefactNode, args: tuple, kwargs: dict, render: bool
845 ) -> dict:
846 """Execute a single artefact node and return child tasks.
848 Returns:
849 Dictionary of child task names to TaskNode objects.
850 """
851 # Process incoming data
852 for arg in args:
853 if (
854 isinstance(arg, Artefact)
855 and node.name == arg.get_name()
856 and node.check_type(type(arg))
857 ):
858 self._check_existing(node, arg)
860 # Mark as READY even if no items matched — parent task has completed
861 if node.status == NodeStatus.NOT_READY:
862 node.status = NodeStatus.READY
864 # Save results
865 self.save_run(
866 kwargs.get("key", ""), kwargs.get("value", ""), artefact_node=node
867 )
869 # Update visualization
870 self._update_node_status(node, True)
871 self.render("workflow_run", file=kwargs.get("value", ""), render=render)
873 # Collect child tasks
874 children = self._get_node_children(node)
875 return {child.name: child for child in children if isinstance(child, TaskNode)}
877 def _get_node_children(
878 self, node: TaskNode | ArtefactNode
879 ) -> list[TaskNode | ArtefactNode]:
880 """Get all children of a node."""
881 return self.node_manager.get_children(node)
883 def _update_node_status(self, node: TaskNode | ArtefactNode, success: bool):
884 """Update node visualization status based on execution result."""
885 color = self.COLOR_SUCCESS if success else self.COLOR_FAILURE
886 self.node_manager.set_node_attr(
887 str(node.name), color=color, fontcolor=self.COLOR_FONT
888 )
890 def _check_task_ready(self, task_node: TaskNode) -> NodeStatus:
891 """Check if a task node is ready to execute.
893 A task is ready when all its parent artefact nodes have READY status.
895 Args:
896 task_node: The task node to check.
898 Returns:
899 The updated status of the task node.
900 """
901 parents = self.node_manager.get_parents_artefact(task_node)
902 if all(parent.status == NodeStatus.READY for parent in parents):
903 task_node.status = NodeStatus.READY
904 return task_node.status
906 def _check_tasks_ready(self, task_nodes: list[TaskNode]) -> list[TaskNode]:
907 """Filter task nodes to return only those that are ready to execute.
909 Args:
910 task_nodes: List of task nodes to check.
912 Returns:
913 List of task nodes that have READY status.
914 """
915 ready_tasks: list[TaskNode] = list()
916 ready_tasks = [
917 task_node
918 for task_node in task_nodes
919 if self._check_task_ready(task_node) == NodeStatus.READY
920 ]
921 return ready_tasks
923 def _check_existing(self, node: ArtefactNode, new_item: Artefact) -> None:
924 """Check if an identical item already exists in the artefact node.
926 If an identical item is found, update the existing item with the new item's ID.
927 If no identical item is found, add the new item to the artefact node.
929 Args:
930 node: The artefact node to check.
931 new_item: The new data item to check or add.
932 """
933 items = node.items
934 for existing_id, existing_item in items.items():
935 if self._is_identical_item(existing_item, new_item):
936 new_item.id = existing_item.id
937 items[existing_id] = new_item
938 return
940 # No identical item found, add as new
941 node.add_item(new_item)
943 def _is_identical_item(self, original_item: Artefact, new_item: Artefact) -> bool:
944 """Check if two data items are identical.
946 Compares both parent relationships and serialized data attributes.
947 """
948 if set(new_item.parents) != set(original_item.parents):
949 return False
951 return self._compare_serialized_data(original_item, new_item)
953 def _compare_serialized_data(
954 self, original_item: Artefact, new_item: Artefact
955 ) -> bool:
956 """Compare serialized data of two items, ignoring IDs and parents."""
957 temp_original = self._prepare_for_comparison(copy.deepcopy(original_item))
958 temp_new = self._prepare_for_comparison(copy.deepcopy(new_item))
960 return json.dumps(temp_original, sort_keys=True) == json.dumps(
961 temp_new, sort_keys=True
962 )
964 def _prepare_for_comparison(self, item: Artefact) -> dict:
965 """Prepare data item for comparison by normalizing and removing metadata."""
966 serialized = item.serialize()
967 serialized.pop("parents", None)
968 serialized.pop("id", None)
970 # Normalize string data for comparison
971 if isinstance(serialized.get("data"), str):
972 serialized["data"] = serialized["data"].strip().lower()
974 return serialized
976 def _apply_task_run_overrides(self, task: Task, overrides: dict) -> None:
977 """Apply runtime CLI overrides to a task before execution."""
978 if not overrides:
979 return
980 merged_overrides = self._merge_task_params(task.name, dict(overrides))
981 for key, value in merged_overrides.items():
982 setattr(task, key, value)
984 def _run_task(self, task_node: TaskNode, **kwargs) -> list:
985 """Execute a task node and return its output results.
987 Prepares inputs from parent artefact nodes, runs the task, and collects
988 all output data items.
990 Args:
991 task_node: The task node to execute.
992 **kwargs: Additional execution parameters.
994 Returns:
995 List of Artefact objects produced by the task execution.
997 Raises:
998 TypeError: If the node does not contain a valid Task.
999 """
1000 task = task_node.get_task()
1001 if not isinstance(task, Task):
1002 raise TypeError(f"Node {task_node.name} is not a Task")
1003 logger.info("Attempting to run task: %s", task.name)
1004 self._apply_task_run_overrides(task, kwargs.get("task_overrides", {}))
1005 parents = self.node_manager.get_parents_artefact(task_node)
1006 entries = task.prepare_inputs(parents)
1007 results = list()
1008 if len(entries) == 0:
1009 logger.warning("No input data found for task %s, skipping.", task.name)
1010 return results
1011 with change_dir(self.path):
1012 for entry in track(
1013 entries,
1014 total=len(entries),
1015 description=f"Running {task.name}",
1016 ):
1017 if not isinstance(entry, list):
1018 entry = list(entry)
1019 entry_result = self._execute_entries(entry, task)
1020 results.extend(entry_result)
1021 return results
1023 def _execute_entries(self, entries: list[Artefact], task: Task) -> list[Artefact]:
1024 """Execute a task with the given input entries.
1026 Runs the task with the provided data entries, sets parent relationships
1027 on output items, and handles any execution errors.
1029 Args:
1030 entries: List of input Artefact objects for the task.
1031 task: The task to execute.
1033 Returns:
1034 List of Artefact objects produced by the task, with parent relationships set.
1035 """
1036 output: Artefact | list[Artefact] | None = task.run(*entries)
1037 try:
1038 if output is None:
1039 return list()
1040 if not isinstance(output, list):
1041 output = [output]
1042 for result in output:
1043 result.set_parents(entries)
1044 return output
1045 except Exception as e:
1046 logger.error(
1047 "Error executing entries for task %s: %s (entries: %s)",
1048 task.name,
1049 e,
1050 entries,
1051 exc_info=True,
1052 )
1053 return list()
1055 def _merge_task_params(self, task_name: str, kwargs: dict) -> dict:
1056 """Merge task-specific parameters with provided kwargs."""
1057 task_params = self.data.get(task_name, dict()).get("params", dict())
1058 return task_params | kwargs
1060 def ensemble(self, folder: list[Path], output: str | Path):
1061 """Create ensembled files from source folders.
1063 Delegates to EnsembleProcessor for the actual ensemble logic.
1065 Args:
1066 folder: List of source folders containing collectra results
1067 output: Output folder path for ensembled results
1068 """
1069 from collectra.pipelines.ensemble import EnsembleProcessor
1071 self._ensure_workflow_connected()
1072 processor = EnsembleProcessor(
1073 node_manager=self.node_manager,
1074 ext=self.ext,
1075 name=self.name,
1076 version=self.version,
1077 )
1078 processor.ensemble(folder, Path(output))
1080 @property
1081 def metadata(self) -> dict:
1082 """Get pipeline metadata dictionary.
1084 Returns:
1085 Dictionary containing pipeline name, extension, and version
1086 under the PIPELINE_METADATA_KEY.
1087 """
1088 return {self.PIPELINE_METADATA_KEY: self.pipeline_metadata.copy()}
1090 def gui_graph(self) -> dict:
1091 """Return the live workflow graph in the shape expected by the GUI."""
1092 self._ensure_workflow_connected()
1093 nodes = []
1094 for name in self.node_manager.get_node_names():
1095 node = self.node_manager.resolve_node(name)
1096 inputs = list(self.node_manager.flow.predecessors(name))
1097 outputs = list(self.node_manager.flow.successors(name))
1098 params = {}
1100 if isinstance(node, TaskNode):
1101 task = node.get_task()
1102 detail = task.get_class_path()
1103 params = getattr(task, "params", {})
1104 if "LLM" in detail:
1105 node_type = "llm"
1106 elif "ObjectDetection" in detail or "ImageClassifier" in detail:
1107 node_type = "detector"
1108 else:
1109 node_type = "processor"
1110 else:
1111 type_names = sorted(type_.get_class_path() for type_ in node.types)
1112 detail = ", ".join(type_names)
1113 node_type = "data" if node.items else "implicit"
1115 nodes.append(
1116 {
1117 "id": name,
1118 "label": name,
1119 "nodeType": node_type,
1120 "detail": detail,
1121 "inputs": inputs,
1122 "outputs": outputs,
1123 "params": params,
1124 }
1125 )
1127 edges = [
1128 {"source": source, "target": target}
1129 for source, target in self.node_manager.flow.edges
1130 ]
1131 return {"nodes": nodes, "edges": edges, "ext": self.ext}
1133 def save(self, task_name: str | None = None):
1134 """Save the workflow configuration to pipeline.yaml."""
1135 with change_dir(self.path):
1136 # Pull updated data from file if it exists to ensure we don't overwrite changes
1137 data_diff: dict | None = None
1138 if Path(self.PIPELINE_FILE).exists():
1139 with open(self.PIPELINE_FILE, "r") as f:
1140 file_data = yaml.safe_load(f) or dict()
1141 file_data.pop(self.PIPELINE_METADATA_KEY, None)
1142 # compare diffs between current data and file data, prioritizing file data for any keys that exist in both
1143 data_diff = dict()
1144 for key, data in file_data.items():
1145 if key != task_name and file_data[key] != self.data.get(key, None):
1146 data_diff[key] = data
1147 data_to_save = (
1148 self.data | data_diff if data_diff and len(data_diff) > 0 else self.data
1149 )
1150 config = self.metadata | data_to_save
1151 self._write_pipeline_config(config)
1153 def _write_pipeline_config(self, config: dict):
1154 """Write pipeline configuration to YAML file."""
1155 write_yaml(config, self.PIPELINE_FILE)
1157 def _get_io_list(self, io: list | str) -> list:
1158 """Convert an IO specification to a list format.
1160 Args:
1161 io: Either a single string or a list of IO identifiers.
1163 Returns:
1164 List containing the IO identifier(s).
1165 """
1166 return [io] if isinstance(io, str) else io
1168 def _check_task_io(self, io_type: str, task: Task, task_config: dict) -> list:
1169 """Extract IO node information from task configuration.
1171 Args:
1172 io_type: Either 'input' or 'output'.
1173 task: The task object.
1174 task_config: Configuration dictionary for the task.
1176 Returns:
1177 List of tuples containing (node_name, types).
1178 """
1179 if io_type not in task_config:
1180 return list()
1182 param_types = list(
1183 unpack_types(
1184 task.run, get_param_types if io_type == "input" else get_return_type
1185 ).items()
1186 )
1187 io = self._get_io_list(task_config.get(io_type, []))
1188 nodes: list = list()
1189 for item in io:
1190 for _, types in param_types:
1191 if not isinstance(types, tuple):
1192 types = (types,)
1193 nodes.append((item, types))
1194 return nodes
1196 def _check_data_assignment(self, collectra_data: dict, data: dict) -> dict:
1197 """Check and resolve data references in configuration.
1199 If a value references another key in collectra_data, replace it with
1200 the actual value.
1202 Args:
1203 collectra_data: The main workflow data dictionary.
1204 data: The dictionary to check for references.
1206 Returns:
1207 Updated data dictionary with resolved references.
1208 """
1209 for key, value in data.items():
1210 if isinstance(value, str) and value in collectra_data:
1211 data[key] = collectra_data[value]
1212 return data
1214 def connect(self):
1215 """Initialise the DAG representing the workflow. Preparing it for execution."""
1216 relations: dict[str, dict] = dict()
1217 with change_dir(self.path):
1218 self._build_nodes_and_relations(relations)
1219 self._create_graph_edges(relations)
1221 def _build_nodes_and_relations(self, relations: dict):
1222 """Build workflow nodes and track their relationships."""
1223 for name, value in self.data.items():
1224 data = copy.deepcopy(value)
1225 if not isinstance(data, dict) or "type" not in data:
1226 continue
1228 cls_ = load_class_from_string(data.pop("type"))
1229 data = self._check_data_assignment(self.data, data)
1230 obj = self._create_or_reuse_object(name, cls_, data)
1232 if isinstance(obj, Task):
1233 self._process_task_node(obj, data, relations)
1234 else:
1235 self.node_manager.add_artefact_node(name, obj=obj)
1237 def _create_or_reuse_object(self, name: str, cls_: type, data: dict):
1238 """Create new object or reuse existing task from graph."""
1239 existing_task = self.node_manager.try_resolve_existing_task(name, cls_)
1240 if existing_task is not None:
1241 return existing_task
1242 return cls_(name, **data)
1244 def _process_task_node(self, task: Task, data: dict, relations: dict):
1245 """Process a task node by adding it to the graph and tracking IO."""
1246 task_key = task.name
1247 task.pipeline = self
1248 relations[task_key] = {
1249 "input": self._get_io_list(data.get("input", [])),
1250 "output": self._get_io_list(data.get("output", [])),
1251 }
1252 self.node_manager.add_task_node(task_key, task)
1254 # Add artefact nodes for task inputs and outputs
1255 ios = self._extract_task_io_nodes(task, data)
1256 for io_key, io_types in ios:
1257 self.node_manager.add_artefact_node(io_key, types=list(io_types))
1259 def _extract_task_io_nodes(
1260 self, task: Task, task_config: dict
1261 ) -> list[tuple[str, tuple]]:
1262 """Extract input and output node information from task configuration."""
1263 ios = []
1264 ios.extend(self._check_task_io("input", task, task_config))
1265 ios.extend(self._check_task_io("output", task, task_config))
1266 return ios
1268 def _create_graph_edges(self, relations: dict):
1269 """Create edges in the workflow graph based on task relationships."""
1270 for task_name, io in relations.items():
1271 inputs = io.get("input", [])
1272 outputs = io.get("output", [])
1274 for input_node in inputs:
1275 if self.node_manager.has_node(input_node):
1276 self.node_manager.add_edge(input_node, task_name)
1278 for output_node in self._get_io_list(outputs):
1279 if self.node_manager.has_node(output_node):
1280 self.node_manager.add_edge(task_name, output_node)
1282 def render(self, dest: str | Path = "", file="", render=False) -> None:
1283 """Render the workflow graph as an SVG visualization.
1285 Args:
1286 dest: Destination path for the rendered file.
1287 file: Optional file being processed (shown in title).
1288 render: Whether to actually render (allows conditional rendering).
1289 """
1290 if not render:
1291 return
1293 self._ensure_workflow_connected()
1294 visual_graph = self._prepare_graph_for_visualization()
1295 dot_str = self._generate_dot_string(visual_graph, file)
1297 dest_path = Path(dest) if dest else self.path / "workflow"
1298 graphviz.Source(dot_str).render(filename=dest_path, format="svg", cleanup=True)
1299 time.sleep(0.1) # Brief pause for file system
1301 def _prepare_graph_for_visualization(self) -> nx.DiGraph:
1302 """Prepare a copy of the graph for visualization.
1304 Removes internal node objects that shouldn't be visualized.
1305 """
1306 visual_graph = self.node_manager.copy_graph()
1307 for node_id in visual_graph.nodes:
1308 node_data = visual_graph.nodes[node_id]
1309 if "node" in node_data:
1310 del node_data["node"]
1311 return visual_graph
1313 def _generate_dot_string(self, graph: nx.DiGraph, file: str = "") -> str:
1314 """Generate DOT format string from graph.
1316 Args:
1317 graph: The networkx graph to convert.
1318 file: Optional file name to include in graph title.
1320 Returns:
1321 DOT format string representation of the graph.
1322 """
1323 dot_str = nx.nx_pydot.to_pydot(graph).to_string()
1325 if file:
1326 title = f"{self.name} Workflow - Processing: {file}"
1327 dot_str = dot_str.replace(
1328 "{", f'{{\nlabel="{title}";\nlabelloc="t";\nfontsize=20;\n\n', 1
1329 )
1331 return dot_str