File indexing completed on 2026-09-28 09:37:32
0001 """
0002 Snakemake Agent: listens on /topic/epictopic for stf_processed / run_snakemake
0003 messages and runs a Snakemake workflow, forwarding every message field as a
0004 Snakemake config value.
0005 """
0006
0007 import traceback
0008 from pathlib import Path
0009 from datetime import datetime
0010
0011 from swf_common_lib.base_agent import BaseAgent
0012
0013 from snakemake.api import SnakemakeApi
0014 from snakemake.settings.types import (
0015 ResourceSettings,
0016 ConfigSettings,
0017 OutputSettings,
0018 )
0019
0020
0021 class SnakemakeAgent(BaseAgent):
0022 """Agent that triggers Snakemake workflows in response to ActiveMQ messages."""
0023
0024 TRIGGER_MSG_TYPES = {"stf_processed", "run_snakemake"}
0025
0026 def __init__(self, snakefile, workdir=None, cores=1, config_path=None, debug=False):
0027 super().__init__(
0028 agent_type="SNAKEMAKE",
0029 subscription_queue="/topic/epictopic",
0030 debug=debug,
0031 config_path=config_path,
0032 )
0033 self.snakefile = Path(snakefile).resolve()
0034 self.workdir = Path(workdir).resolve() if workdir else self.snakefile.parent
0035 self.cores = cores
0036 self.logger.info(f"SnakemakeAgent initialised, snakefile={self.snakefile}")
0037
0038 def on_message(self, frame):
0039 message_data, msg_type = self.log_received_message(frame)
0040 if message_data is None:
0041 return
0042
0043 msg_namespace = message_data.get("namespace")
0044 if msg_namespace is not None and msg_namespace != self.namespace:
0045 return
0046
0047 if "execution_id" in message_data:
0048 self.current_execution_id = message_data["execution_id"]
0049 if "run_id" in message_data:
0050 self.current_run_id = message_data["run_id"]
0051
0052 if msg_type not in self.TRIGGER_MSG_TYPES:
0053 return
0054
0055 self.set_processing()
0056 try:
0057 self._run_workflow(message_data)
0058 self._publish_result(message_data, success=True)
0059 except Exception as exc:
0060 self.logger.error(f"Workflow failed: {exc}")
0061 self.logger.error(traceback.format_exc())
0062 self._publish_result(message_data, success=False, error=str(exc))
0063 finally:
0064 self.set_ready()
0065
0066 def _run_workflow(self, message_data):
0067 snake_config = {k: str(v) for k, v in message_data.items()}
0068
0069 with SnakemakeApi(
0070 output_settings=OutputSettings(verbose=self.DEBUG),
0071 ) as api:
0072 workflow_api = api.workflow(
0073 snakefile=self.snakefile,
0074 workdir=self.workdir,
0075 resource_settings=ResourceSettings(cores=self.cores),
0076 config_settings=ConfigSettings(config=snake_config),
0077 )
0078 dag_api = workflow_api.dag()
0079 dag_api.execute_workflow()
0080
0081 def _publish_result(self, trigger_message, success, error=None):
0082 msg = {
0083 "msg_type": "snakemake_complete",
0084 "namespace": self.namespace,
0085 "success": success,
0086 "snakefile": str(self.snakefile),
0087 "trigger_msg_type": trigger_message.get("msg_type"),
0088 "run_id": trigger_message.get("run_id"),
0089 "timestamp": datetime.now().isoformat(),
0090 }
0091 if trigger_message.get("execution_id"):
0092 msg["execution_id"] = trigger_message["execution_id"]
0093 if error:
0094 msg["error"] = error
0095
0096 try:
0097 self.send_message("/topic/epictopic", msg)
0098 except Exception as exc:
0099 self.logger.error(f"Failed to publish snakemake_complete: {exc}")
0100
0101
0102 if __name__ == "__main__":
0103 import argparse
0104
0105 parser = argparse.ArgumentParser(description="Snakemake Agent")
0106 parser.add_argument("--snakefile", required=True)
0107 parser.add_argument("--workdir", default=None)
0108 parser.add_argument("--cores", type=int, default=1)
0109 parser.add_argument("--debug", action="store_true")
0110 parser.add_argument("--testbed-config", default=None)
0111 args = parser.parse_args()
0112
0113 SnakemakeAgent(
0114 snakefile=args.snakefile,
0115 workdir=args.workdir,
0116 cores=args.cores,
0117 debug=args.debug,
0118 config_path=args.testbed_config,
0119 ).run()
