runnable 0.38.0__tar.gz → 0.39.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (72) hide show
  1. {runnable-0.38.0 → runnable-0.39.0}/PKG-INFO +8 -1
  2. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/conditional.py +19 -1
  3. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/fail.py +7 -1
  4. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/map.py +16 -1
  5. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/parallel.py +7 -1
  6. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/stub.py +7 -1
  7. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/success.py +7 -1
  8. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/task.py +15 -1
  9. {runnable-0.38.0 → runnable-0.39.0}/pyproject.toml +13 -4
  10. {runnable-0.38.0 → runnable-0.39.0}/runnable/cli.py +52 -0
  11. {runnable-0.38.0 → runnable-0.39.0}/runnable/graph.py +271 -0
  12. {runnable-0.38.0 → runnable-0.39.0}/runnable/nodes.py +16 -0
  13. {runnable-0.38.0 → runnable-0.39.0}/runnable/tasks.py +79 -0
  14. {runnable-0.38.0 → runnable-0.39.0}/.gitignore +0 -0
  15. {runnable-0.38.0 → runnable-0.39.0}/LICENSE +0 -0
  16. {runnable-0.38.0 → runnable-0.39.0}/README.md +0 -0
  17. {runnable-0.38.0 → runnable-0.39.0}/extensions/README.md +0 -0
  18. {runnable-0.38.0 → runnable-0.39.0}/extensions/__init__.py +0 -0
  19. {runnable-0.38.0 → runnable-0.39.0}/extensions/catalog/README.md +0 -0
  20. {runnable-0.38.0 → runnable-0.39.0}/extensions/catalog/any_path.py +0 -0
  21. {runnable-0.38.0 → runnable-0.39.0}/extensions/catalog/file_system.py +0 -0
  22. {runnable-0.38.0 → runnable-0.39.0}/extensions/catalog/minio.py +0 -0
  23. {runnable-0.38.0 → runnable-0.39.0}/extensions/catalog/pyproject.toml +0 -0
  24. {runnable-0.38.0 → runnable-0.39.0}/extensions/catalog/s3.py +0 -0
  25. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/README.md +0 -0
  26. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/__init__.py +0 -0
  27. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/emulate.py +0 -0
  28. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/k8s.py +0 -0
  29. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/k8s_job_spec.yaml +0 -0
  30. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/local.py +0 -0
  31. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/local_container.py +0 -0
  32. {runnable-0.38.0 → runnable-0.39.0}/extensions/job_executor/pyproject.toml +0 -0
  33. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/README.md +0 -0
  34. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/__init__.py +0 -0
  35. {runnable-0.38.0 → runnable-0.39.0}/extensions/nodes/pyproject.toml +0 -0
  36. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/README.md +0 -0
  37. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/__init__.py +0 -0
  38. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/argo.py +0 -0
  39. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/emulate.py +0 -0
  40. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/local.py +0 -0
  41. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/local_container.py +0 -0
  42. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/mocked.py +0 -0
  43. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/pyproject.toml +0 -0
  44. {runnable-0.38.0 → runnable-0.39.0}/extensions/pipeline_executor/retry.py +0 -0
  45. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/README.md +0 -0
  46. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/__init__.py +0 -0
  47. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/any_path.py +0 -0
  48. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/chunked_fs.py +0 -0
  49. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/chunked_minio.py +0 -0
  50. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/db/implementation_FF.py +0 -0
  51. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/db/integration_FF.py +0 -0
  52. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/file_system.py +0 -0
  53. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/generic_chunked.py +0 -0
  54. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/minio.py +0 -0
  55. {runnable-0.38.0 → runnable-0.39.0}/extensions/run_log_store/pyproject.toml +0 -0
  56. {runnable-0.38.0 → runnable-0.39.0}/extensions/secrets/README.md +0 -0
  57. {runnable-0.38.0 → runnable-0.39.0}/extensions/secrets/dotenv.py +0 -0
  58. {runnable-0.38.0 → runnable-0.39.0}/extensions/secrets/pyproject.toml +0 -0
  59. {runnable-0.38.0 → runnable-0.39.0}/runnable/__init__.py +0 -0
  60. {runnable-0.38.0 → runnable-0.39.0}/runnable/catalog.py +0 -0
  61. {runnable-0.38.0 → runnable-0.39.0}/runnable/context.py +0 -0
  62. {runnable-0.38.0 → runnable-0.39.0}/runnable/datastore.py +0 -0
  63. {runnable-0.38.0 → runnable-0.39.0}/runnable/defaults.py +0 -0
  64. {runnable-0.38.0 → runnable-0.39.0}/runnable/entrypoints.py +0 -0
  65. {runnable-0.38.0 → runnable-0.39.0}/runnable/exceptions.py +0 -0
  66. {runnable-0.38.0 → runnable-0.39.0}/runnable/executor.py +0 -0
  67. {runnable-0.38.0 → runnable-0.39.0}/runnable/names.py +0 -0
  68. {runnable-0.38.0 → runnable-0.39.0}/runnable/parameters.py +0 -0
  69. {runnable-0.38.0 → runnable-0.39.0}/runnable/pickler.py +0 -0
  70. {runnable-0.38.0 → runnable-0.39.0}/runnable/sdk.py +0 -0
  71. {runnable-0.38.0 → runnable-0.39.0}/runnable/secrets.py +0 -0
  72. {runnable-0.38.0 → runnable-0.39.0}/runnable/utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: runnable
3
- Version: 0.38.0
3
+ Version: 0.39.0
4
4
  Summary: Add your description here
5
5
  Author-email: "Vammi, Vijay" <vijay.vammi@astrazeneca.com>
6
6
  License-File: LICENSE
@@ -18,12 +18,19 @@ Provides-Extra: docker
18
18
  Requires-Dist: docker>=7.1.0; extra == 'docker'
19
19
  Provides-Extra: examples
20
20
  Requires-Dist: pandas>=2.2.3; extra == 'examples'
21
+ Provides-Extra: examples-torch
22
+ Requires-Dist: torch>=2.7.1; extra == 'examples-torch'
21
23
  Provides-Extra: k8s
22
24
  Requires-Dist: kubernetes>=31.0.0; extra == 'k8s'
23
25
  Provides-Extra: notebook
24
26
  Requires-Dist: ploomber-engine>=0.0.33; extra == 'notebook'
25
27
  Provides-Extra: s3
26
28
  Requires-Dist: cloudpathlib[s3]; extra == 's3'
29
+ Provides-Extra: ui
30
+ Requires-Dist: fastapi>=0.95.0; extra == 'ui'
31
+ Requires-Dist: jinja2>=3.1.2; extra == 'ui'
32
+ Requires-Dist: python-multipart>=0.0.5; extra == 'ui'
33
+ Requires-Dist: uvicorn>=0.22.0; extra == 'ui'
27
34
  Description-Content-Type: text/markdown
28
35
 
29
36
 
@@ -7,7 +7,7 @@ from pydantic import Field, field_serializer, field_validator
7
7
  from runnable import console, defaults
8
8
  from runnable.datastore import Parameter
9
9
  from runnable.graph import Graph, create_graph
10
- from runnable.nodes import CompositeNode, MapVariableType
10
+ from runnable.nodes import CompositeNode, MapVariableType, NodeInD3
11
11
 
12
12
  logger = logging.getLogger(defaults.LOGGER_NAME)
13
13
 
@@ -241,3 +241,21 @@ class ConditionalNode(CompositeNode):
241
241
  step_log.status = defaults.FAIL
242
242
 
243
243
  self._context.run_log_store.add_step_log(step_log, self._context.run_id)
244
+
245
+ def to_d3_node(self) -> NodeInD3:
246
+ def get_display_string() -> str:
247
+ display = f"match {self.parameter}:\n"
248
+ for case in self.branches.keys():
249
+ display += f' case "{case}":\n ...\n'
250
+ if self.default:
251
+ display += " case _:\n ...\n"
252
+ return display
253
+
254
+ return NodeInD3(
255
+ id=self.internal_name,
256
+ label="conditional",
257
+ metadata={
258
+ "conditioned on": self.parameter,
259
+ "display": get_display_string(),
260
+ },
261
+ )
@@ -6,7 +6,7 @@ from pydantic import Field
6
6
  from runnable import datastore, defaults
7
7
  from runnable.datastore import StepLog
8
8
  from runnable.defaults import MapVariableType
9
- from runnable.nodes import TerminalNode
9
+ from runnable.nodes import NodeInD3, TerminalNode
10
10
 
11
11
 
12
12
  class FailNode(TerminalNode):
@@ -70,3 +70,9 @@ class FailNode(TerminalNode):
70
70
  step_log.attempts.append(attempt_log)
71
71
 
72
72
  return step_log
73
+
74
+ def to_d3_node(self) -> NodeInD3:
75
+ return NodeInD3(
76
+ id=self.internal_name,
77
+ label="fail",
78
+ )
@@ -18,7 +18,7 @@ from runnable.datastore import (
18
18
  )
19
19
  from runnable.defaults import MapVariableType
20
20
  from runnable.graph import Graph, create_graph
21
- from runnable.nodes import CompositeNode
21
+ from runnable.nodes import CompositeNode, NodeInD3
22
22
 
23
23
  logger = logging.getLogger(defaults.LOGGER_NAME)
24
24
 
@@ -348,3 +348,18 @@ class MapNode(CompositeNode):
348
348
  self._context.run_log_store.set_parameters(
349
349
  parameters=params, run_id=self._context.run_id
350
350
  )
351
+
352
+ def to_d3_node(self) -> NodeInD3:
353
+ return NodeInD3(
354
+ id=self.internal_name,
355
+ label="map",
356
+ metadata={
357
+ "node_type": "map",
358
+ "iterate_on": self.iterate_on, # Parameter name containing the iterable
359
+ "iterate_as": self.iterate_as, # Name used for each iteration
360
+ "map_branch_id": self.internal_name
361
+ + "."
362
+ + defaults.MAP_PLACEHOLDER, # The branch identifier pattern
363
+ "is_composite": True, # Flag indicating this is a composite node
364
+ },
365
+ )
@@ -6,7 +6,7 @@ from pydantic import Field, field_serializer
6
6
  from runnable import defaults
7
7
  from runnable.defaults import MapVariableType
8
8
  from runnable.graph import Graph, create_graph
9
- from runnable.nodes import CompositeNode
9
+ from runnable.nodes import CompositeNode, NodeInD3
10
10
 
11
11
 
12
12
  class ParallelNode(CompositeNode):
@@ -157,3 +157,9 @@ class ParallelNode(CompositeNode):
157
157
  step_log.status = defaults.FAIL
158
158
 
159
159
  self._context.run_log_store.add_step_log(step_log, self._context.run_id)
160
+
161
+ def to_d3_node(self) -> NodeInD3:
162
+ return NodeInD3(
163
+ id=self.internal_name,
164
+ label="parallel",
165
+ )
@@ -7,7 +7,7 @@ from pydantic import ConfigDict, Field
7
7
  from runnable import datastore, defaults
8
8
  from runnable.datastore import StepLog
9
9
  from runnable.defaults import MapVariableType
10
- from runnable.nodes import ExecutableNode
10
+ from runnable.nodes import ExecutableNode, NodeInD3
11
11
 
12
12
  logger = logging.getLogger(defaults.LOGGER_NAME)
13
13
 
@@ -87,3 +87,9 @@ class StubNode(ExecutableNode):
87
87
  step_log.attempts.append(attempt_log)
88
88
 
89
89
  return step_log
90
+
91
+ def to_d3_node(self) -> NodeInD3:
92
+ return NodeInD3(
93
+ id=self.internal_name,
94
+ label="stub",
95
+ )
@@ -6,7 +6,7 @@ from pydantic import Field
6
6
  from runnable import datastore, defaults
7
7
  from runnable.datastore import StepLog
8
8
  from runnable.defaults import MapVariableType
9
- from runnable.nodes import TerminalNode
9
+ from runnable.nodes import NodeInD3, TerminalNode
10
10
 
11
11
 
12
12
  class SuccessNode(TerminalNode):
@@ -70,3 +70,9 @@ class SuccessNode(TerminalNode):
70
70
  step_log.attempts.append(attempt_log)
71
71
 
72
72
  return step_log
73
+
74
+ def to_d3_node(self) -> NodeInD3:
75
+ return NodeInD3(
76
+ id=self.internal_name,
77
+ label="success",
78
+ )
@@ -7,7 +7,7 @@ from pydantic import ConfigDict, Field
7
7
  from runnable import datastore, defaults
8
8
  from runnable.datastore import StepLog
9
9
  from runnable.defaults import MapVariableType
10
- from runnable.nodes import ExecutableNode
10
+ from runnable.nodes import ExecutableNode, NodeInD3
11
11
  from runnable.tasks import BaseTaskType, create_task
12
12
 
13
13
  logger = logging.getLogger(defaults.LOGGER_NAME)
@@ -90,3 +90,17 @@ class TaskNode(ExecutableNode):
90
90
  step_log.attempts.append(attempt_log)
91
91
 
92
92
  return step_log
93
+
94
+ def to_d3_node(self) -> NodeInD3:
95
+ """
96
+ Convert the task node to a D3 node representation.
97
+
98
+ Returns:
99
+ NodeInD3: The D3 node representation of the task node.
100
+ """
101
+ return NodeInD3(
102
+ id=self.internal_name,
103
+ label="task",
104
+ task_type=self.executable.task_type,
105
+ metadata=self.executable.get_d3_metadata(),
106
+ )
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "runnable"
3
- version = "0.38.0"
3
+ version = "0.39.0"
4
4
  description = "Add your description here"
5
5
  readme = "README.md"
6
6
  authors = [
@@ -35,6 +35,15 @@ k8s = [
35
35
  s3 = [
36
36
  "cloudpathlib[s3]"
37
37
  ]
38
+ examples-torch = [
39
+ "torch>=2.7.1",
40
+ ]
41
+ ui = [
42
+ "fastapi>=0.95.0",
43
+ "jinja2>=3.1.2",
44
+ "uvicorn>=0.22.0",
45
+ "python-multipart>=0.0.5",
46
+ ]
38
47
 
39
48
 
40
49
  [dependency-groups]
@@ -56,9 +65,7 @@ docs = [
56
65
  release = [
57
66
  "python-semantic-release>=9.15.2",
58
67
  ]
59
- examples-torch = [
60
- "torch>=2.7.1",
61
- ]
68
+
62
69
 
63
70
  [tool.uv.workspace]
64
71
  members = ["extensions/catalog",
@@ -67,6 +74,7 @@ members = ["extensions/catalog",
67
74
  "extensions/pipeline_executor",
68
75
  "extensions/run_log_store",
69
76
  "extensions/secrets",
77
+ "visualization"
70
78
  ]
71
79
 
72
80
  [tool.uv.sources]
@@ -76,6 +84,7 @@ catalog = {workspace = true}
76
84
  run_log_store = {workspace = true}
77
85
  pipeline_executor = {workspace = true}
78
86
  job_executor = {workspace = true}
87
+ visualization = {workspace = true}
79
88
 
80
89
 
81
90
  [project.scripts]
@@ -274,5 +274,57 @@ def execute_job(
274
274
  )
275
275
 
276
276
 
277
+ @app.command()
278
+ def ui(
279
+ host: Annotated[
280
+ str,
281
+ typer.Option(
282
+ "--host",
283
+ "-h",
284
+ help="The host to bind the server to",
285
+ ),
286
+ ] = "127.0.0.1",
287
+ port: Annotated[
288
+ int,
289
+ typer.Option(
290
+ "--port",
291
+ "-p",
292
+ help="The port to bind the server to",
293
+ ),
294
+ ] = 8000,
295
+ reload: Annotated[
296
+ bool,
297
+ typer.Option(
298
+ "--reload",
299
+ help="Enable auto-reload for development",
300
+ ),
301
+ ] = False,
302
+ ):
303
+ """
304
+ Start the web UI for pipeline visualization.
305
+
306
+ This command starts a FastAPI web server that provides a user interface
307
+ for visualizing and exploring runnable pipelines.
308
+ """
309
+ try:
310
+ import uvicorn
311
+
312
+ from visualization.main import app as web_app
313
+ except ImportError:
314
+ typer.echo(
315
+ "UI dependencies not installed. Install with: pip install runnable[ui]",
316
+ err=True,
317
+ )
318
+ raise typer.Exit(1)
319
+
320
+ typer.echo(f"Starting web UI at http://{host}:{port}")
321
+ uvicorn.run(
322
+ web_app,
323
+ host=host,
324
+ port=port,
325
+ reload=reload,
326
+ )
327
+
328
+
277
329
  if __name__ == "__main__":
278
330
  app()
@@ -499,3 +499,274 @@ def search_branch_by_internal_name(dag: Graph, internal_name: str):
499
499
  return current_branch
500
500
 
501
501
  raise exceptions.BranchNotFoundError(internal_name)
502
+
503
+
504
+ def get_visualization_data(graph: Graph) -> Dict[str, Any]:
505
+ """
506
+ Convert the graph into a D3 visualization friendly format with nodes and links.
507
+ Handles composite nodes (parallel, map, conditional) by recursively processing their embedded graphs.
508
+
509
+ Args:
510
+ graph: The Graph object to convert
511
+
512
+ Returns:
513
+ Dict with two keys:
514
+ - nodes: List of node objects with id, type, name, and alias
515
+ - links: List of edge objects with source and target node ids
516
+ """
517
+ import rich.console
518
+
519
+ from extensions.nodes.conditional import ConditionalNode
520
+ from extensions.nodes.map import MapNode
521
+ from extensions.nodes.parallel import ParallelNode
522
+ from runnable.nodes import ExecutableNode
523
+
524
+ rich_print = rich.console.Console().print
525
+
526
+ rich_print(graph)
527
+
528
+ nodes = []
529
+ links = []
530
+ processed_nodes = set()
531
+
532
+ def process_node(
533
+ node: BaseNode,
534
+ parent_id: Optional[str] = None,
535
+ current_graph: Graph = graph,
536
+ map_node_id: Optional[str] = None,
537
+ conditional_node_id: Optional[str] = None,
538
+ ) -> str:
539
+ node_id = f"{node.internal_name}"
540
+ node_alias = node.name # Alias based on the node's name
541
+
542
+ if node_id not in processed_nodes:
543
+ node_data = node.to_d3_node().model_dump(exclude_none=True)
544
+ node_data["alias"] = node_alias # Add alias to the node data
545
+ node_data["display_name"] = node_alias # Use alias as the display name
546
+
547
+ # Add map or parallel related metadata if this node is part of a map branch or parallel branch
548
+ if map_node_id:
549
+ if "metadata" not in node_data:
550
+ node_data["metadata"] = {}
551
+
552
+ # Mark this node as being part of a map branch
553
+ node_data["metadata"]["belongs_to_node"] = map_node_id
554
+
555
+ # If this is the map node itself, add a special attribute
556
+ if node_id == map_node_id:
557
+ node_data["metadata"]["is_map_root"] = True
558
+
559
+ # Add conditional related metadata if this node is part of a conditional branch
560
+ if conditional_node_id:
561
+ if "metadata" not in node_data:
562
+ node_data["metadata"] = {}
563
+
564
+ # Mark this node as being part of a conditional branch
565
+ node_data["metadata"]["belongs_to_node"] = conditional_node_id
566
+
567
+ # If this is the conditional node itself, add a special attribute
568
+ if node_id == conditional_node_id:
569
+ node_data["metadata"]["is_conditional_root"] = True
570
+
571
+ # Mark parallel nodes with special metadata
572
+ if isinstance(node, ParallelNode):
573
+ if "metadata" not in node_data:
574
+ node_data["metadata"] = {}
575
+
576
+ # Add parallel node type to metadata
577
+ node_data["metadata"]["node_type"] = "parallel"
578
+ node_data["metadata"]["parallel_branch_id"] = node_id
579
+
580
+ # Mark conditional nodes with special metadata
581
+ if isinstance(node, ConditionalNode):
582
+ if "metadata" not in node_data:
583
+ node_data["metadata"] = {}
584
+
585
+ # Add conditional node type to metadata
586
+ node_data["metadata"]["node_type"] = "conditional"
587
+ node_data["metadata"]["conditional_branch_id"] = node_id
588
+
589
+ nodes.append(node_data)
590
+ processed_nodes.add(node_id)
591
+
592
+ # Add link from parent if it exists
593
+ if parent_id:
594
+ links.append({"source": parent_id, "target": node_id})
595
+
596
+ # Handle composite nodes with embedded graphs
597
+ if isinstance(node, (ParallelNode, MapNode, ConditionalNode)):
598
+ if isinstance(node, ParallelNode):
599
+ # Process each parallel branch
600
+ for _, branch in node.branches.items():
601
+ branch_start = branch.get_node_by_name(branch.start_at)
602
+ process_node(
603
+ branch_start,
604
+ node_id,
605
+ branch,
606
+ map_node_id=node_id,
607
+ conditional_node_id=conditional_node_id,
608
+ )
609
+
610
+ # Handle next node connection after parallel branches complete
611
+ if hasattr(node, "next_node") and node.next_node:
612
+ try:
613
+ next_node = current_graph.get_node_by_name(node.next_node)
614
+ next_id = process_node(
615
+ next_node,
616
+ None,
617
+ current_graph=current_graph,
618
+ map_node_id=map_node_id,
619
+ conditional_node_id=conditional_node_id,
620
+ )
621
+ links.append(
622
+ {
623
+ "source": node_id,
624
+ "target": next_id,
625
+ "type": "success",
626
+ }
627
+ )
628
+ except exceptions.NodeNotFoundError as e:
629
+ rich_print(
630
+ f"Warning: Next node '{node.next_node}' not found for parallel node '{node.name}': {e}"
631
+ )
632
+
633
+ elif isinstance(node, MapNode):
634
+ # Process map branch
635
+ branch_start = node.branch.get_node_by_name(node.branch.start_at)
636
+ # Process the branch with additional context about the map node
637
+ process_node(
638
+ branch_start,
639
+ node_id,
640
+ node.branch,
641
+ map_node_id=node_id,
642
+ conditional_node_id=conditional_node_id,
643
+ )
644
+
645
+ elif isinstance(node, ConditionalNode):
646
+ # Process each conditional branch
647
+ for _, branch in node.branches.items():
648
+ branch_start = branch.get_node_by_name(branch.start_at)
649
+ process_node(
650
+ branch_start,
651
+ node_id,
652
+ branch,
653
+ map_node_id=map_node_id,
654
+ conditional_node_id=node_id,
655
+ )
656
+ if node.default:
657
+ default_start = node.default.get_node_by_name(
658
+ node.default.start_at
659
+ )
660
+ process_node(
661
+ default_start,
662
+ node_id,
663
+ node.default,
664
+ map_node_id=map_node_id,
665
+ conditional_node_id=node_id,
666
+ )
667
+
668
+ # Handle next node connection after conditional branches complete
669
+ if hasattr(node, "next_node") and node.next_node:
670
+ try:
671
+ next_node = current_graph.get_node_by_name(node.next_node)
672
+ next_id = process_node(
673
+ next_node,
674
+ None,
675
+ current_graph=current_graph,
676
+ map_node_id=map_node_id,
677
+ conditional_node_id=conditional_node_id,
678
+ )
679
+ links.append(
680
+ {
681
+ "source": node_id,
682
+ "target": next_id,
683
+ "type": "success",
684
+ }
685
+ )
686
+ except exceptions.NodeNotFoundError as e:
687
+ rich_print(
688
+ f"Warning: Next node '{node.next_node}' not found for conditional node '{node.name}': {e}"
689
+ )
690
+
691
+ # Add links to next and on_failure nodes if they exist
692
+ if isinstance(node, ExecutableNode):
693
+ # Handle normal "next" links (success path)
694
+ if hasattr(node, "next_node") and node.next_node:
695
+ try:
696
+ next_node = current_graph.get_node_by_name(node.next_node)
697
+ next_id = process_node(
698
+ next_node,
699
+ None,
700
+ current_graph=current_graph,
701
+ map_node_id=map_node_id,
702
+ conditional_node_id=conditional_node_id,
703
+ )
704
+ links.append(
705
+ {"source": node_id, "target": next_id, "type": "success"}
706
+ )
707
+ except exceptions.NodeNotFoundError as e:
708
+ rich_print(
709
+ f"Warning: Next node '{node.next_node}' not found for node '{node.name}': {e}"
710
+ )
711
+
712
+ # Handle on_failure links (failure path)
713
+ if hasattr(node, "on_failure") and node.on_failure:
714
+ try:
715
+ failure_node = current_graph.get_node_by_name(node.on_failure)
716
+ failure_id = process_node(
717
+ failure_node,
718
+ None,
719
+ current_graph=current_graph,
720
+ map_node_id=map_node_id,
721
+ conditional_node_id=conditional_node_id,
722
+ )
723
+ links.append(
724
+ {"source": node_id, "target": failure_id, "type": "failure"}
725
+ )
726
+ except exceptions.NodeNotFoundError as e:
727
+ rich_print(
728
+ f"Warning: On-failure node '{node.on_failure}' not found for node '{node.name}': {e}"
729
+ )
730
+
731
+ # For backward compatibility, also process all neighbors
732
+ # This handles cases where node might have other connection types
733
+ next_nodes = node._get_neighbors()
734
+ for next_node_name in next_nodes:
735
+ # Skip nodes we've already handled explicitly
736
+ if (
737
+ hasattr(node, "next_node") and node.next_node == next_node_name
738
+ ) or (
739
+ hasattr(node, "on_failure")
740
+ and node.on_failure == next_node_name
741
+ ):
742
+ continue
743
+
744
+ try:
745
+ next_node = current_graph.get_node_by_name(next_node_name)
746
+ next_id = process_node(
747
+ next_node,
748
+ None,
749
+ current_graph=current_graph,
750
+ map_node_id=map_node_id,
751
+ conditional_node_id=conditional_node_id,
752
+ )
753
+ links.append(
754
+ {"source": node_id, "target": next_id, "type": "default"}
755
+ )
756
+ except exceptions.NodeNotFoundError as e:
757
+ rich_print(
758
+ f"Warning: Neighbor node '{next_node_name}' not found for node '{node.name}': {e}"
759
+ )
760
+
761
+ return node_id
762
+
763
+ # Start processing from the start node
764
+ start_node = graph.get_node_by_name(graph.start_at)
765
+ try:
766
+ process_node(
767
+ start_node, None, graph, map_node_id=None, conditional_node_id=None
768
+ )
769
+ except (exceptions.NodeNotFoundError, AttributeError, KeyError) as e:
770
+ rich_print(f"Error processing node {start_node}: {e}")
771
+
772
+ return {"nodes": nodes, "links": links}
@@ -15,6 +15,13 @@ logger = logging.getLogger(defaults.LOGGER_NAME)
15
15
  # --8<-- [start:docs]
16
16
 
17
17
 
18
+ class NodeInD3(BaseModel):
19
+ id: str
20
+ label: str
21
+ task_type: Optional[str] = None
22
+ metadata: Optional[Dict[str, Any]] = None
23
+
24
+
18
25
  class BaseNode(ABC, BaseModel):
19
26
  """
20
27
  Base class with common functionality provided for a Node of a graph.
@@ -369,6 +376,15 @@ class BaseNode(ABC, BaseModel):
369
376
  Dict[str, Any]: _description_
370
377
  """
371
378
 
379
+ @abstractmethod
380
+ def to_d3_node(self) -> NodeInD3:
381
+ """
382
+ Convert the node to a D3 node representation.
383
+
384
+ Returns:
385
+ NodeInD3: The D3 node representation of the current node.
386
+ """
387
+
372
388
 
373
389
  # --8<-- [end:docs]
374
390
  class TraversalNode(BaseNode):
@@ -209,6 +209,9 @@ class BaseTaskType(BaseModel):
209
209
  parameters=diff_parameters, run_id=self._context.run_id
210
210
  )
211
211
 
212
+ def get_d3_metadata(self) -> dict[str, str]:
213
+ raise NotImplementedError
214
+
212
215
 
213
216
  def task_return_to_parameter(task_return: TaskReturns, value: Any) -> Parameter:
214
217
  # implicit support for pydantic models
@@ -287,6 +290,72 @@ class PythonTaskType(BaseTaskType): # pylint: disable=too-few-public-methods
287
290
  task_type: str = Field(default="python", serialization_alias="command_type")
288
291
  command: str
289
292
 
293
+ def get_d3_metadata(self) -> Dict[str, str]:
294
+ module, func = utils.get_module_and_attr_names(self.command)
295
+
296
+ # Import inspect module to get function signature
297
+ import inspect
298
+
299
+ def format_type_annotation(annotation):
300
+ """Format type annotation in a more readable way"""
301
+ if annotation == inspect._empty:
302
+ return "Any"
303
+ elif hasattr(annotation, "__name__"):
304
+ return annotation.__name__
305
+ elif hasattr(annotation, "__origin__"):
306
+ # Handle typing types like List, Dict, etc.
307
+ origin = (
308
+ annotation.__origin__.__name__
309
+ if hasattr(annotation.__origin__, "__name__")
310
+ else str(annotation.__origin__)
311
+ )
312
+ args = ", ".join(
313
+ format_type_annotation(arg) for arg in annotation.__args__
314
+ )
315
+ return f"{origin}[{args}]"
316
+ else:
317
+ # Fall back to string representation without 'typing.'
318
+ return str(annotation).replace("typing.", "")
319
+
320
+ # Import the module and get the function
321
+ sys.path.insert(0, os.getcwd())
322
+ try:
323
+ imported_module = importlib.import_module(module)
324
+ f = getattr(imported_module, func)
325
+
326
+ # Get function signature
327
+ sig = inspect.signature(f)
328
+
329
+ # Format parameters with type annotations
330
+ params_list = []
331
+ for param_name, param in sig.parameters.items():
332
+ type_annotation = format_type_annotation(param.annotation)
333
+ params_list.append(f"{param_name}: {type_annotation}")
334
+
335
+ params_str = ", ".join(params_list)
336
+
337
+ # Format returns based on self.returns or use "None" if no returns specified
338
+ if self.returns:
339
+ returns_str = ", ".join([f"{r.name}({r.kind})" for r in self.returns])
340
+ else:
341
+ returns_str = "None"
342
+
343
+ # Format function signature
344
+ signature = f"def {func}({params_str}) -> {returns_str}:"
345
+
346
+ return {
347
+ "module": module,
348
+ "function": func,
349
+ "signature": signature,
350
+ }
351
+ except Exception as e:
352
+ logger.warning(f"Could not inspect function {self.command}: {str(e)}")
353
+ return {
354
+ "module": module,
355
+ "function": func,
356
+ "signature": f"def {func}(...) -> ...",
357
+ }
358
+
290
359
  def execute_command(
291
360
  self,
292
361
  map_variable: MapVariableType = None,
@@ -440,6 +509,11 @@ class NotebookTaskType(BaseTaskType):
440
509
  command: str
441
510
  optional_ploomber_args: dict = {}
442
511
 
512
+ def get_d3_metadata(self) -> Dict[str, str]:
513
+ return {
514
+ "notebook": self.command,
515
+ }
516
+
443
517
  @field_validator("command")
444
518
  @classmethod
445
519
  def notebook_should_end_with_ipynb(cls, command: str) -> str:
@@ -640,6 +714,11 @@ class ShellTaskType(BaseTaskType):
640
714
  task_type: str = Field(default="shell", serialization_alias="command_type")
641
715
  command: str
642
716
 
717
+ def get_d3_metadata(self) -> Dict[str, str]:
718
+ return {
719
+ "command": self.command[:50],
720
+ }
721
+
643
722
  @field_validator("returns")
644
723
  @classmethod
645
724
  def returns_should_be_json(cls, returns: List[TaskReturns]):
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes