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

50 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-24 15:16 +0000

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

2 

3import argparse 

4import os 

5import sys 

6from pathlib import Path 

7 

8from srunx.callbacks import SlackCallback 

9from srunx.logging import configure_workflow_logging, get_logger 

10from srunx.runner import WorkflowRunner 

11 

12logger = get_logger(__name__) 

13 

14 

15def create_workflow_parser() -> argparse.ArgumentParser: 

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

17 parser = argparse.ArgumentParser( 

18 description="Execute YAML-defined workflows using SLURM", 

19 formatter_class=argparse.RawDescriptionHelpFormatter, 

20 epilog=""" 

21Example YAML workflow: 

22 name: ml_pipeline 

23 jobs: 

24 - name: preprocess 

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

26 resources: 

27 nodes: 1 

28 gpus_per_node: 2 

29 

30 - name: train 

31 path: /path/to/train.sh 

32 depends_on: 

33 - preprocess 

34 

35 - name: evaluate 

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

37 depends_on: 

38 - train 

39 environment: 

40 conda: ml_env 

41 

42 - name: upload 

43 command: ["python", "upload_model.py"] 

44 depends_on: 

45 - train 

46 environment: 

47 venv: /path/to/venv 

48 

49 - name: notify 

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

51 depends_on: 

52 - evaluate 

53 - upload 

54 environment: 

55 venv: /path/to/venv 

56 """, 

57 ) 

58 

59 parser.add_argument( 

60 "yaml_file", 

61 type=str, 

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

63 ) 

64 

65 parser.add_argument( 

66 "--validate-only", 

67 action="store_true", 

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

69 ) 

70 

71 parser.add_argument( 

72 "--dry-run", 

73 action="store_true", 

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

75 ) 

76 

77 parser.add_argument( 

78 "--log-level", 

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

80 default="INFO", 

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

82 ) 

83 

84 parser.add_argument( 

85 "--slack", 

86 action="store_true", 

87 help="Send notifications to Slack", 

88 ) 

89 

90 return parser 

91 

92 

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

94 """Handle workflow execution command.""" 

95 # Configure logging for workflow execution 

96 configure_workflow_logging(level=args.log_level) 

97 

98 try: 

99 yaml_file = Path(args.yaml_file) 

100 if not yaml_file.exists(): 

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

102 sys.exit(1) 

103 

104 # Setup callbacks if requested 

105 callbacks = [] 

106 if getattr(args, "slack", False): 

107 webhook_url = os.getenv("SLACK_WEBHOOK_URL") 

108 if not webhook_url: 

109 raise ValueError("SLACK_WEBHOOK_URL environment variable is not set") 

110 callbacks.append(SlackCallback(webhook_url=webhook_url)) 

111 

112 runner = WorkflowRunner.from_yaml(yaml_file, callbacks=callbacks) 

113 

114 # Validate dependencies 

115 runner.workflow.validate() 

116 

117 if args.validate_only: 

118 logger.info("Workflow validation successful") 

119 return 

120 

121 if args.dry_run: 

122 runner.workflow.show() 

123 return 

124 

125 # Execute workflow 

126 results = runner.run() 

127 

128 logger.info("Job Results:") 

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

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

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

132 else: 

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

134 

135 except Exception as e: 

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

137 sys.exit(1) 

138 

139 

140def main() -> None: 

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

142 parser = create_workflow_parser() 

143 args = parser.parse_args() 

144 

145 cmd_run_workflow(args) 

146 

147 

148if __name__ == "__main__": 

149 main()