Coverage for src/srunx/cli/workflow.py: 97%

1"""CLI interface for workflow management."""

2

3 import argparse

4 import sys

5 from pathlib import Path

6

7 from srunx.logging import configure_workflow_logging, get_logger

8 from srunx.models import Job, ShellJob, Workflow

9 from srunx.workflows.runner import WorkflowRunner

10

11 logger = get_logger(__name__)

12

13

14 def create_workflow_parser() -> argparse.ArgumentParser:

15 """Create argument parser for workflow commands."""

16 parser = argparse.ArgumentParser(

17 description="Execute YAML-defined workflows using SLURM and Prefect",

18 formatter_class=argparse.RawDescriptionHelpFormatter,

19 epilog="""

20 Example YAML workflow:

21 name: ml_pipeline

22 tasks:

23 - name: preprocess

24 command: ["python", "preprocess.py"]

25 nodes: 1

26

27 - name: train

28 command: ["python", "train.py"]

29 depends_on: [preprocess]

30 gpus_per_node: 1

31 conda: ml_env

32

33 - name: evaluate

34 command: ["python", "evaluate.py"]

35 depends_on: [train]

36

37 - name: notify

38 command: ["python", "notify.py"]

39 depends_on: [train, evaluate]

40 async: true

41 """,

42 )

43

44 parser.add_argument(

45 "yaml_file",

46 type=str,

47 help="Path to YAML workflow definition file",

48 )

49

50 parser.add_argument(

51 "--validate-only",

52 action="store_true",

53 help="Only validate the workflow file without executing",

54 )

55

56 parser.add_argument(

57 "--dry-run",

58 action="store_true",

59 help="Show what would be executed without running jobs",

60 )

61

62 parser.add_argument(

63 "--log-level",

64 choices=["DEBUG", "INFO", "WARNING", "ERROR"],

65 default="INFO",

66 help="Set logging level (default: %(default)s)",

67 )

68

69 return parser

70

71

72 def cmd_run_workflow(args: argparse.Namespace) -> None:

73 """Handle workflow execution command."""

74 # Configure logging for workflow execution

75 configure_workflow_logging(level=args.log_level)

76

77 try:

78 yaml_file = Path(args.yaml_file)

79 if not yaml_file.exists():

80 logger.error(f"Workflow file not found: {args.yaml_file}")

81 sys.exit(1)

82

83 runner = WorkflowRunner()

84

85 # Load workflow for validation

86 workflow = runner.load_from_yaml(yaml_file)

87 logger.info(

88 f"Loaded workflow '{workflow.name}' with {len(workflow.tasks)} tasks"

89 )

90

91 # Validate dependencies

92 _validate_workflow_dependencies(workflow)

93

94 if args.validate_only:

95 logger.info("Workflow validation successful")

96 return

97

98 if args.dry_run:

99 _show_workflow_plan(workflow)

100 return

101

102 # Execute workflow

103 logger.info("Starting workflow execution")

104 results = runner.execute_workflow(workflow)

105

106 logger.info("Workflow execution completed successfully")

107 logger.info("Job Results:")

108 for task_name, job in results.items():

109 if hasattr(job, "job_id") and job.job_id:

110 logger.info(f" {task_name}: Job ID {job.job_id}")

111 else:

112 logger.info(f" {task_name}: {job}")

113

114 except Exception as e:

115 logger.error(f"Workflow execution failed: {e}")

116 sys.exit(1)

117

118

119 def _validate_workflow_dependencies(workflow: Workflow) -> None:

120 """Validate workflow task dependencies."""

121 task_names = {task.name for task in workflow.tasks}

122

123 for task in workflow.tasks:

124 for dependency in task.depends_on:

125 if dependency not in task_names:

126 raise ValueError(

127 f"Task '{task.name}' depends on unknown task '{dependency}'"

128 )

129

130 # Check for circular dependencies (simple check)

131 visited = set()

132 rec_stack = set()

133

134 def has_cycle(task_name: str) -> bool:

135 if task_name in rec_stack:

136 return True

137 if task_name in visited:

138 return False

139

140 visited.add(task_name)

141 rec_stack.add(task_name)

142

143 task = workflow.get_task(task_name)

144 if task:

145 for dependency in task.depends_on:

146 if has_cycle(dependency):

147 return True

148

149 rec_stack.remove(task_name)

150 return False

151

152 for task in workflow.tasks:

153 if has_cycle(task.name):

154 raise ValueError(

155 f"Circular dependency detected involving task '{task.name}'"

156 )

157

158

159 def _show_workflow_plan(workflow: Workflow) -> None:

160 """Show workflow execution plan."""

161 msg = f"""\

162 Workflow execution plan:

163 Workflow: {workflow.name}

164 Tasks: {len(workflow.tasks)}

165"""

166

167 for task in workflow.tasks:

168 msg += f"\t\tTask: {task.name}\n"

169 if isinstance(task.job, Job):

170 msg += f"\t\t\tCommand: {' '.join(task.job.command or [])}\n"

171 msg += f"\t\t\tResources: {task.job.resources.nodes} nodes, {task.job.resources.gpus_per_node} GPUs/node\n"

172 if task.job.environment.conda:

173 msg += f"\t\t\tConda env: {task.job.environment.conda}\n"

174 if task.job.environment.sqsh:

175 msg += f"\t\t\tSqsh: {task.job.environment.sqsh}\n"

176 elif isinstance(task.job, ShellJob):

177 msg += f"\t\t\tPath: {task.job.path}\n"

178

179 if task.depends_on:

180 msg += f"\t\t\tDependencies: {', '.join(task.depends_on)}\n"

181 if task.async_execution:

182 msg += "\t\t\tExecution: asynchronous\n"

183

184 logger.info(msg)

185

186

187 def main() -> None:

188 """Main entry point for workflow CLI."""

189 parser = create_workflow_parser()

190 args = parser.parse_args()

191

192 cmd_run_workflow(args)

193

194

195 if __name__ == "__main__":

196 main()