refactor(pipeline): replace networkx with rustworkx (#518)

Author: disrupted

kpops/pipeline/__init__.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, TypeAlias
 
-import networkx as nx
+import rustworkx as rx
 import yaml
 from pydantic import (
     SerializeAsAny,
@@ -37,7 +37,8 @@ class Pipeline:
     """Pipeline representation."""
 
     _component_index: dict[str, PipelineComponent] = field(default_factory=dict)
-    _graph: nx.DiGraph[str] = field(default_factory=nx.DiGraph)
+    _graph: rx.PyDiGraph[str, None] = field(default_factory=rx.PyDiGraph)
+    _node_index: dict[str, int] = field(default_factory=dict)
 
     @property
     def step_names(self) -> list[str]:
@@ -118,23 +119,23 @@ async def run_graph_layers(
             for layer_components in pending_layers:
                 await run_layer_parallel(layer_components)
 
-        graph: nx.DiGraph[str] = self._graph.copy()
+        graph = self._graph.copy()
 
         # We add an extra node to the graph, connecting all the leaf nodes to it
         # in that way we make this node the root of the graph, avoiding backtracking
-        root_node = "root_node_bfs"
-        graph.add_node(root_node)
+        root_node = graph.add_node("root_node_bfs")
 
-        for node in graph:
-            predecessors = list(graph.predecessors(node))
-            if not predecessors:
-                graph.add_edge(root_node, node)
+        for node in graph.node_indices():
+            if node != root_node and not list(graph.predecessors(node)):
+                graph.add_edge(root_node, node, None)
 
-        layers_graph: list[list[str]] = list(nx.bfs_layers(graph, root_node))
+        layers_graph = list(rx.bfs_layers(graph, [root_node]))
 
         sorted_layers: list[list[PipelineComponent]] = []
         for layer in layers_graph[1:]:
-            if parallel_components := self.__get_parallel_components_from(layer):
+            if parallel_components := self.__get_parallel_components_from(
+                [graph[idx] for idx in layer]
+            ):
                 sorted_layers.append(parallel_components)
 
         if reverse:
@@ -158,22 +159,27 @@ def __iter__(self) -> Iterator[PipelineComponent]:
     def __len__(self) -> int:
         return len(self.components)
 
+    def __get_or_add_node(self, node_id: str) -> int:
+        if node_id not in self._node_index:
+            self._node_index[node_id] = self._graph.add_node(node_id)
+        return self._node_index[node_id]
+
     def __add_to_graph(self, component: PipelineComponent):
-        self._graph.add_node(component.id)
+        node = self.__get_or_add_node(component.id)
 
         for input_topic in component.inputs:
-            self.__add_input(input_topic.id, component.id)
+            self.__add_input(input_topic.id, node)
 
         for output_topic in component.outputs:
-            self.__add_output(output_topic.id, component.id)
+            self.__add_output(output_topic.id, node)
 
-    def __add_output(self, topic_id: str, source: str) -> None:
-        self._graph.add_node(topic_id)
-        self._graph.add_edge(source, topic_id)
+    def __add_output(self, topic_id: str, source: int) -> None:
+        topic = self.__get_or_add_node(topic_id)
+        self._graph.add_edge(source, topic, None)
 
-    def __add_input(self, topic_id: str, target: str) -> None:
-        self._graph.add_node(topic_id)
-        self._graph.add_edge(topic_id, target)
+    def __add_input(self, topic_id: str, target: int) -> None:
+        topic = self.__get_or_add_node(topic_id)
+        self._graph.add_edge(topic, target, None)
 
     def __get_parallel_components_from(
         self, layer: list[str]
@@ -187,7 +193,7 @@ def gen_parallel_components() -> Iterator[PipelineComponent]:
         return list(gen_parallel_components())
 
     def __validate_graph(self) -> None:
-        if not nx.is_directed_acyclic_graph(self._graph):
+        if not rx.is_directed_acyclic_graph(self._graph):
             msg = "Pipeline is not a valid DAG."
             raise ValueError(msg)
 

pyproject.toml

@@ -2,36 +2,36 @@
 name = "kpops"
 version = "10.9.0"
 description = "KPOps is a tool to deploy Kafka pipelines to Kubernetes"
+readme = "README.md"
+requires-python = ">=3.11"
 authors = [
   { name = "bakdata", email = "opensource@bakdata.com" },
 ]
-readme = "README.md"
-keywords = ["kafka", "kubernetes", "stream-processing", "pipelines"]
+keywords = ["kafka", "kubernetes", "pipelines", "stream-processing"]
 classifiers = [
   "Environment :: Console",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: MIT License",
   "Operating System :: OS Independent",
-  "Topic :: Software Development :: Libraries :: Python Modules",
-  "Topic :: Software Development :: Libraries",
   "Topic :: Software Development",
-  "Intended Audience :: Developers",
+  "Topic :: Software Development :: Libraries",
+  "Topic :: Software Development :: Libraries :: Python Modules",
   "Typing :: Typed",
-  "License :: OSI Approved :: MIT License",
 ]
-requires-python = ">=3.11"
 dependencies = [
   "anyio>=4.3.0",
   "cachetools>=5.2.0",
   "croniter>=3.0.3",
   "dictdiffer>=0.9.0",
   "httpx>=0.24.1",
   "lightkube>=0.15.3",
-  "networkx>=3.1",
   "pydantic>=2.10.6",
   "pydantic-settings>=2.0.3",
   "pyhumps>=3.7.3",
   "python-schema-registry-client>=2.4.1",
   "pyyaml>=6.0",
   "rich>=12.4.4",
+  "rustworkx>=0.18.0",
   "typer>=0.12.5",
 ]
 
@@ -40,9 +40,9 @@ kpops = "kpops.cli.main:app"
 
 [dependency-groups]
 dev = [
+  "basedpyright>=1.31.5",
   "lefthook>=1.10.4",
   "polyfactory>=2.13.0",
-  "basedpyright>=1.31.5",
   "pytablewriter[from]>=1.0.0",
   "pytest>=8.3.2",
   "pytest-asyncio>=0.23.8",
@@ -54,31 +54,20 @@ dev = [
   "ruff>=0.9.9",
 ]
 docs = [
+  "mdx-truly-sane-lists>=1.3",
   "mike>=1.1.2",
   "mkdocs>=1.4.2",
   "mkdocs-exclude-search>=0.6.5",
   "mkdocs-glightbox>=0.3.1",
   "mkdocs-macros-plugin>=1.0.5",
   "mkdocs-material>=9.0.0",
   "mkdocstrings[python]>=0.25.1",
-  "mdx-truly-sane-lists>=1.3",
 ]
 
-[tool.uv]
-package = true
-trusted-publishing = "always"
-
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
-[tool.hatch.build.targets.wheel]
-packages = ["kpops"]
-
-[tool.pytest.ini_options]
-asyncio_mode = "auto"
-asyncio_default_fixture_loop_scope = "function"
-
 [tool.basedpyright]
 include = ["kpops", "tests", "hooks"]
 allowedUntypedLibraries = ["pytablewriter"]
@@ -114,6 +103,13 @@ strictParameterNoneValue = false
 # TODO: apply only to tests https://github.com/DetachHead/basedpyright/issues/668
 reportPrivateUsage = false
 
+[tool.hatch.build.targets.wheel]
+packages = ["kpops"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+
 [tool.ruff]
 output-format = "grouped"
 show-fixes = true
@@ -219,3 +215,7 @@ extend-immutable-calls = ["typer.Argument"]
 
 [tool.ruff.lint.flake8-type-checking]
 runtime-evaluated-base-classes = ["pydantic.BaseModel"]
+
+[tool.uv]
+package = true
+trusted-publishing = "always"

tests/pipeline/test_generate.py

@@ -692,13 +692,15 @@ def test_validate_simple_graph(self):
             RESOURCE_PATH / "pipelines-with-graphs" / "simple-pipeline" / PIPELINE_YAML,
         )
         assert len(pipeline.components) == 2
-        assert len(pipeline._graph.nodes) == 3
-        assert len(pipeline._graph.edges) == 2
+        assert len(pipeline._graph.nodes()) == 3
+        assert len(pipeline._graph.edges()) == 2
         topic_nodes = [
-            node for node in pipeline._graph.nodes if node.startswith("topic-")
+            node for node in pipeline._graph.nodes() if node.startswith("topic-")
         ]
         assert len(topic_nodes) == 1
-        assert len(pipeline.components) == len(pipeline._graph.nodes) - len(topic_nodes)
+        assert len(pipeline.components) == len(pipeline._graph.nodes()) - len(
+            topic_nodes
+        )
 
     def test_validate_components_are_disabled_in_production_but_enabled_on_development(
         self,
@@ -719,27 +721,27 @@ def test_validate_components_are_disabled_in_production_but_enabled_on_developme
         )
 
         assert len(pipeline_production.components) == 1
-        assert len(pipeline_production._graph.edges) == 0
+        assert len(pipeline_production._graph.edges()) == 0
         topic_nodes_production = [
             node
-            for node in pipeline_production._graph.nodes
+            for node in pipeline_production._graph.nodes()
             if node.startswith("topic-")
         ]
         assert len(topic_nodes_production) == 0
         assert len(pipeline_production.components) == len(
-            pipeline_production._graph.nodes
+            pipeline_production._graph.nodes()
         ) - len(topic_nodes_production)
 
-        assert len(pipeline_development._graph.nodes) == 3
-        assert len(pipeline_development._graph.edges) == 2
+        assert len(pipeline_development._graph.nodes()) == 3
+        assert len(pipeline_development._graph.edges()) == 2
         topic_nodes_development = [
             node
-            for node in pipeline_development._graph.nodes
+            for node in pipeline_development._graph.nodes()
             if node.startswith("topic-")
         ]
         assert len(topic_nodes_development) == 1
         assert len(pipeline_development.components) == len(
-            pipeline_development._graph.nodes
+            pipeline_development._graph.nodes()
         ) - len(topic_nodes_development)
 
     def test_validate_topic_and_component_same_name(self):
@@ -749,8 +751,11 @@ def test_validate_topic_and_component_same_name(self):
             / "same-topic-and-component-name"
             / PIPELINE_YAML,
         )
-        component, topic = list(pipeline._graph.nodes)
-        edges = list(pipeline._graph.edges)
+        component, topic = list(pipeline._graph.nodes())
+        edges = [
+            (pipeline._graph[s], pipeline._graph[t])
+            for s, t in pipeline._graph.edge_list()
+        ]
         assert component == topic.removeprefix("topic-")
         assert (component, topic) in edges