Coverage for src/srunx/workflows/runner.py: 99%

1"""Workflow runner for executing YAML-defined workflows with SLURM and Prefect."""

2

3 from pathlib import Path

4 from typing import Any

5

6 import yaml

7 from prefect import flow

8 from prefect.states import State

9 from srunx.logging import get_logger

10 from srunx.models import (

11 Job,

12 JobEnvironment,

13 JobResource,

14 ShellJob,

15 Workflow,

16 WorkflowTask,

17)

18 from srunx.workflows.tasks import submit_and_monitor_job, submit_job_async

19

20 logger = get_logger(__name__)

21

22

23 class WorkflowRunner:

24 """Runner for executing workflows defined in YAML."""

25

26 def __init__(self) -> None:

27 """Initialize workflow runner."""

28 self.executed_tasks: dict[str, State[Job | ShellJob]] = {}

29

30 def load_from_yaml(self, yaml_path: str | Path) -> Workflow:

31 """Load and validate a workflow from a YAML file.

32

33 Args:

34 yaml_path: Path to the YAML workflow definition file.

35

36 Returns:

37 Validated Workflow object.

38

39 Raises:

40 FileNotFoundError: If the YAML file doesn't exist.

41 yaml.YAMLError: If the YAML is malformed.

42 ValidationError: If the workflow structure is invalid.

43 """

44 yaml_file = Path(yaml_path)

45 if not yaml_file.exists():

46 raise FileNotFoundError(f"Workflow file not found: {yaml_path}")

47

48 with open(yaml_file, encoding="utf-8") as f:

49 data = yaml.safe_load(f)

50

51 return self._parse_workflow_data(data)

52

53 def _parse_workflow_data(self, data: dict) -> Workflow:

54 """Parse workflow data from dictionary."""

55 workflow_name = data.get("name", "unnamed_workflow")

56 tasks_data = data.get("tasks", [])

57

58 tasks = []

59 for task_data in tasks_data:

60 task = self._parse_task_data(task_data)

61 tasks.append(task)

62

63 return Workflow(name=workflow_name, tasks=tasks)

64

65 def _parse_task_data(self, task_data: dict) -> WorkflowTask:

66 """Parse a single task from dictionary."""

67 # Basic task properties

68 name = task_data["name"]

69 file = task_data.get("file")

70 async_execution = task_data.get("async", False)

71 depends_on = task_data.get("depends_on", [])

72

73 job_data: dict[str, Any] = {"name": name}

74 if task_data.get("log_dir") is not None:

75 job_data["log_dir"] = task_data.get("log_dir")

76 if task_data.get("work_dir") is not None:

77 job_data["work_dir"] = task_data.get("work_dir")

78

79 if file:

80 job_data |= {

81 "file": file,

82 "resources": None,

83 "environment": None,

84 }

85 else:

86 command = task_data.get("command")

87

88 # Resource configuration

89 resources = JobResource(

90 nodes=task_data.get("nodes", 1),

91 gpus_per_node=task_data.get("gpus_per_node", 0),

92 ntasks_per_node=task_data.get("ntasks_per_node", 1),

93 cpus_per_task=task_data.get("cpus_per_task", 1),

94 memory_per_node=task_data.get("memory_per_node"),

95 time_limit=task_data.get("time_limit"),

96 )

97

98 # Environment configuration

99 environment = JobEnvironment(

100 conda=task_data.get("conda"),

101 venv=task_data.get("venv"),

102 sqsh=task_data.get("sqsh") or task_data.get("container"),

103 env_vars=task_data.get("env_vars", {}),

104 )

105

106 job_data |= {

107 "command": command,

108 "resources": resources,

109 "environment": environment,

110 }

111

112 # Create job

113 job = Job.model_validate(job_data)

114

115 return WorkflowTask(

116 name=name,

117 job=job,

118 depends_on=depends_on,

119 async_execution=async_execution,

120 )

121

122 def execute_workflow(self, workflow: Workflow) -> dict[str, State[Job | ShellJob]]:

123 """Execute a workflow using Prefect.

124

125 Args:

126 workflow: Workflow to execute.

127

128 Returns:

129 Dictionary mapping task names to Job instances.

130 """

131 task_map = {task.name: task for task in workflow.tasks}

132

133 @flow(name=workflow.name)

134 def workflow_flow() -> dict[str, State[Job | ShellJob]]:

135 """Prefect flow for workflow execution."""

136

137 def execute_task(task_name: str) -> State[Job | ShellJob]:

138 """Execute a task and its dependencies recursively."""

139 if task_name in self.executed_tasks:

140 return self.executed_tasks[task_name]

141

142 task = task_map[task_name]

143

144 # Execute dependencies first

145 for dependency in task.depends_on:

146 execute_task(dependency)

147

148 # Execute the task

149 if task.async_execution:

150 job_future = submit_job_async(task.job)

151 else:

152 job_future = submit_and_monitor_job(task.job)

153

154 self.executed_tasks[task_name] = job_future

155 return job_future

156

157 # Execute all tasks

158 results: dict[str, State[Job | ShellJob]] = {}

159 for task in workflow.tasks:

160 results[task.name] = execute_task(task.name)

161

162 return results

163

164 return workflow_flow()

165

166 def execute_from_yaml(

167 self, yaml_path: str | Path

168 ) -> dict[str, State[Job | ShellJob]]:

169 """Load and execute a workflow from YAML file.

170

171 Args:

172 yaml_path: Path to YAML workflow file.

173

174 Returns:

175 Dictionary mapping task names to Job instances.

176 """

177 logger.info(f"Loading workflow from {yaml_path}")

178 workflow = self.load_from_yaml(yaml_path)

179

180 logger.info(

181 f"Executing workflow '{workflow.name}' with {len(workflow.tasks)} tasks"

182 )

183 results = self.execute_workflow(workflow)

184

185 logger.info("Workflow execution completed")

186 return results

187

188

189 def run_workflow_from_file(yaml_path: str | Path) -> dict[str, State[Job | ShellJob]]:

190 """Convenience function to run workflow from YAML file.

191

192 Args:

193 yaml_path: Path to YAML workflow file.

194

195 Returns:

196 Dictionary mapping task names to Job instances.

197 """

198 runner = WorkflowRunner()

199 return runner.execute_from_yaml(yaml_path)