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

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

2

3 import argparse

4 import os

5 import sys

6 from pathlib import Path

7

8 from srunx.callbacks import SlackCallback

9 from srunx.logging import configure_workflow_logging, get_logger

10 from srunx.runner import WorkflowRunner

11

12 logger = get_logger(__name__)

13

14

15 def 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="""

21 Example 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

93 def 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

140 def 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

148 if __name__ == "__main__":

149 main()