idds/tests/test_prompt_transformer.py

0001 #!/usr/bin/env python
0002 """Small, deterministic unit tests for prompt message construction.
0003
0004 The original test file attempted to exercise live brokers and contained
0005 syntax errors. These tests focus on the pure construction/contract of
0006 messages so they remain reliable in CI without requiring an ActiveMQ
0007 instance.
0008 """
0009 import argparse
0010 import logging
0011 import os
0012 import sys
0013 import time
0014
0015 from idds.common.version import release_version
0016 from idds.common.utils import setup_logging
0017
0018 from idds.common.utils import get_prompt_broker_config
0019 from idds.prompt.transformer import Transformer
0020
0021
0022 setup_logging(__name__)
0023
0024
0025 def get_parser():
0026     """
0027     Return the argparse parser.
0028     """
0029     oparser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]), add_help=True)
0030
0031     # common items
0032     oparser.add_argument(
0033         "--version", action="version", version="%(prog)s " + release_version
0034     )
0035     oparser.add_argument(
0036         "--verbose",
0037         "-v",
0038         default=False,
0039         action="store_true",
0040         help="Print more verbose output.",
0041     )
0042     oparser.add_argument(
0043         "--run_id", type=int, required=True, help="Run ID for the prompt run."
0044     )
0045     oparser.add_argument(
0046         "--workdir", type=str, required=False, default=None, help="Working directory for the prompt run."
0047     )
0048     oparser.add_argument(
0049         "--namespace", type=str, required=False, default=None, help="Namespace for the prompt run."
0050     )
0051
0052     oparser.add_argument(
0053         "--idle_timeout",
0054         type=int,
0055         default=600,
0056         help="Idle timeout in seconds to stop the transformer if no messages are received.",
0057     )
0058
0059     return oparser
0060
0061
0062 def run():
0063
0064     arguments = sys.argv[1:]
0065
0066     oparser = get_parser()
0067     # argcomplete.autocomplete(oparser)
0068
0069     args, unknown = oparser.parse_known_args(arguments)
0070
0071     start_time = time.time()
0072     logging.info("Starting prompt with arguments: {0}".format(arguments))
0073     logging.info("Starting at: %-0.4f" % start_time)
0074
0075     logging.info(f"called with args: {str(args)}")
0076
0077     broker_info = get_prompt_broker_config()
0078     transformer_broker = broker_info.get("transformer_broker", None)
0079     transformer_broadcast_broker = broker_info.get(
0080         "transformer_broadcast_broker", None
0081     )
0082     result_broker = broker_info.get("result_broker", None)
0083
0084     transformer = Transformer(run_id=args.run_id, workdir=args.workdir, namespace=args.namespace, idle_timeout=args.idle_timeout)
0085     transformer._transformer_broker = transformer_broker
0086     transformer._transformer_broadcast_broker = transformer_broadcast_broker
0087     transformer._result_broker = result_broker
0088     transformer._broker_initialized = True
0089
0090     logging.info("Initialized transformer brokers")
0091     ret = transformer.run()
0092     logging.info("Transformer run returned: %s" % str(ret))
0093     return ret
0094
0095
0096 if __name__ == "__main__":
0097     run()