feat: returns ttl during job pop

Closes #394
Haidra-Org · Oct 5, 2024 · ea2d840 · ea2d840
1 parent 28f46b4
commit ea2d840
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 0 deletions.
diff --git a/horde/apis/models/kobold_v2.py b/horde/apis/models/kobold_v2.py
@@ -252,6 +252,7 @@ def __init__(self, api):
                         example="00000000-0000-0000-0000-000000000000",
                     ),
                 ),
+                "ttl": fields.Integer(description="The amount of seconds before this job is considered stale and aborted."),
                 "extra_source_images": fields.List(fields.Nested(self.model_extra_source_images)),
                 "skipped": fields.Nested(self.response_model_generations_skipped, skip_none=True),
                 "softprompt": fields.String(description="The soft prompt requested for this generation."),

diff --git a/horde/apis/models/stable_v2.py b/horde/apis/models/stable_v2.py
@@ -494,6 +494,7 @@ def __init__(self, api):
                         example="00000000-0000-0000-0000-000000000000",
                     ),
                 ),
+                "ttl": fields.Integer(description="The amount of seconds before this job is considered stale and aborted."),
                 "skipped": fields.Nested(self.response_model_generations_skipped, skip_none=True),
                 "model": fields.String(description="Which of the available models to use for this request."),
                 "source_image": fields.String(description="The Base64-encoded webp to use for img2img."),

diff --git a/horde/apis/models/v2.py b/horde/apis/models/v2.py
@@ -438,6 +438,7 @@ def __init__(self, api):
             {
                 "payload": fields.Nested(self.response_model_generation_payload, skip_none=True),
                 "id": fields.String(description="The UUID for this generation."),
+                "ttl": fields.Integer(description="The amount of seconds before this job is considered stale and aborted."),
                 "skipped": fields.Nested(self.response_model_generations_skipped, skip_none=True),
             },
         )

diff --git a/horde/classes/base/waiting_prompt.py b/horde/classes/base/waiting_prompt.py
@@ -284,6 +284,7 @@ def get_pop_payload(self, procgen_list, payload):
             "id": procgen_list[0].id,
             "model": procgen_list[0].model,
             "ids": [g.id for g in procgen_list],
+            "ttl": procgen_list[0].job_ttl,
         }
         if self.extra_source_images and check_bridge_capability("extra_source_images", procgen_list[0].worker.bridge_agent):
             prompt_payload["extra_source_images"] = self.extra_source_images["esi"]

diff --git a/horde/classes/stable/waiting_prompt.py b/horde/classes/stable/waiting_prompt.py
@@ -197,6 +197,7 @@ def get_pop_payload(self, procgen_list, payload):
                 "id": procgen.id,
                 "model": procgen.model,
                 "ids": [g.id for g in procgen_list],
+                "ttl": procgen_list[0].job_ttl,
             }
             if self.source_image and check_bridge_capability("img2img", procgen.worker.bridge_agent):
                 if check_bridge_capability("r2_source", procgen.worker.bridge_agent):