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
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-24 15:16 +0000
1"""CLI interface for workflow management."""
3import argparse
4import os
5import sys
6from pathlib import Path
8from srunx.callbacks import SlackCallback
9from srunx.logging import configure_workflow_logging, get_logger
10from srunx.runner import WorkflowRunner
12logger = get_logger(__name__)
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
30 - name: train
31 path: /path/to/train.sh
32 depends_on:
33 - preprocess
35 - name: evaluate
36 command: ["python", "evaluate.py"]
37 depends_on:
38 - train
39 environment:
40 conda: ml_env
42 - name: upload
43 command: ["python", "upload_model.py"]
44 depends_on:
45 - train
46 environment:
47 venv: /path/to/venv
49 - name: notify
50 command: ["python", "notify.py"]
51 depends_on:
52 - evaluate
53 - upload
54 environment:
55 venv: /path/to/venv
56 """,
57 )
59 parser.add_argument(
60 "yaml_file",
61 type=str,
62 help="Path to YAML workflow definition file",
63 )
65 parser.add_argument(
66 "--validate-only",
67 action="store_true",
68 help="Only validate the workflow file without executing",
69 )
71 parser.add_argument(
72 "--dry-run",
73 action="store_true",
74 help="Show what would be executed without running jobs",
75 )
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 )
84 parser.add_argument(
85 "--slack",
86 action="store_true",
87 help="Send notifications to Slack",
88 )
90 return parser
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)
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)
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))
112 runner = WorkflowRunner.from_yaml(yaml_file, callbacks=callbacks)
114 # Validate dependencies
115 runner.workflow.validate()
117 if args.validate_only:
118 logger.info("Workflow validation successful")
119 return
121 if args.dry_run:
122 runner.workflow.show()
123 return
125 # Execute workflow
126 results = runner.run()
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}")
135 except Exception as e:
136 logger.error(f"Workflow execution failed: {e}")
137 sys.exit(1)
140def main() -> None:
141 """Main entry point for workflow CLI."""
142 parser = create_workflow_parser()
143 args = parser.parse_args()
145 cmd_run_workflow(args)
148if __name__ == "__main__":
149 main()