From 263af996d8ed05e37ef5a62c6ed240216a6eb67b Mon Sep 17 00:00:00 2001
From: Andrew Godwin
Date: Sun, 27 Nov 2022 11:03:52 -0700
Subject: Add run-for option to console runstator

---
 stator/management/commands/runstator.py | 9 +++++++++
 1 file changed, 9 insertions(+)

(limited to 'stator/management')

diff --git a/stator/management/commands/runstator.py b/stator/management/commands/runstator.py
index 3030960..4d52520 100644
--- a/stator/management/commands/runstator.py
+++ b/stator/management/commands/runstator.py
@@ -33,6 +33,13 @@ class Command(BaseCommand):
             default=30,
             help="How often to run cleaning and scheduling",
         )
+        parser.add_argument(
+            "--run-for",
+            "-r",
+            type=int,
+            default=0,
+            help="How long to run for before exiting (defaults to infinite)",
+        )
         parser.add_argument("model_labels", nargs="*", type=str)
 
     def handle(
@@ -41,6 +48,7 @@ class Command(BaseCommand):
         concurrency: int,
         liveness_file: str,
         schedule_interval: int,
+        run_for: int,
         *args,
         **options
     ):
@@ -60,5 +68,6 @@ class Command(BaseCommand):
             concurrency=concurrency,
             liveness_file=liveness_file,
             schedule_interval=schedule_interval,
+            run_for=run_for,
         )
         async_to_sync(runner.run)()
-- 
cgit v1.2.3