diff --git a/examples/archive_experiments.py b/examples/archive_experiments.py new file mode 100644 index 0000000..cba9a83 --- /dev/null +++ b/examples/archive_experiments.py @@ -0,0 +1,59 @@ +#!/usr/bin/python3 +""" +An example script that cleans up failed experiments by moving them to the archive +""" +import argparse +from datetime import datetime + +from trains_agent import APIClient + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--project", "-P", help="Project ID. Only clean up experiments from this project") +parser.add_argument("--user", "-U", help="User ID. Only clean up experiments assigned to this user") +parser.add_argument( + "--status", "-S", default="failed", + help="Experiment status. Only clean up experiments with this status (default %(default)s)" +) +parser.add_argument( + "--iterations", "-I", type=int, + help="Number of iterations. Only clean up experiments with less or equal number of iterations" +) +parser.add_argument( + "--sec-from-start", "-T", type=int, + help="Seconds from start time. " + "Only clean up experiments if less or equal number of seconds have elapsed since started" +) + +args = parser.parse_args() + +client = APIClient() + +tasks = client.tasks.get_all( + project=[args.project] if args.project else None, + user=[args.user] if args.user else None, + status=[args.status] if args.status else None, + system_tags=["-archived"] +) + +count = 0 + +for task in tasks: + if args.iterations and (task.last_iteration or 0) > args.iterations: + continue + if args.sec_from_start: + if not task.started: + continue + if (datetime.utcnow() - task.started.replace(tzinfo=None)).total_seconds() > args.sec_from_start: + continue + + try: + client.tasks.edit( + task=task.id, + system_tags=(task.system_tags or []) + ["archived"], + force=True + ) + count += 1 + except Exception as ex: + print("Failed editing experiment: {}".format(ex)) + +print("Cleaned up {} experiments".format(count))