Skip to content

Commit

Permalink
feat: SIGINT handling; pause job pop for big jobs; better logs
Browse files Browse the repository at this point in the history
  • Loading branch information
tazlin committed Oct 3, 2023
1 parent 2dd8d98 commit 525a047
Show file tree
Hide file tree
Showing 6 changed files with 256 additions and 83 deletions.
1 change: 1 addition & 0 deletions bridgeData_template.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,7 @@ models_to_load:
# This is to avoid loading models which you do not want either due to VRAM constraints, or due to NSFW content
models_to_skip:
- "pix2pix"
- "SDXL_beta::stability.ai#6901" # Do not remove this, as this model would never work
#- "stable_diffusion_inpainting" # Inpainting is generally quite heavy along with other models for smaller GPUs.
#- "stable_diffusion_2.1", # Stable diffusion 2.1 has bigger memory requirements than 1.5, so if your card cannot lift, it, disable it
#- "stable_diffusion_2.0", # Same as Stable diffusion 2.1
Expand Down
28 changes: 25 additions & 3 deletions horde_worker_regen/process_management/horde_process.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import abc
import enum
import signal
import time
from abc import abstractmethod
from enum import auto
Expand Down Expand Up @@ -155,6 +156,7 @@ def receive_and_handle_control_messages(self) -> None:

if message.control_flag == HordeControlFlag.END_PROCESS:
self._end_process = True
logger.info("Received end process message")
return

self._receive_and_handle_control_message(message)
Expand All @@ -165,11 +167,31 @@ def worker_cycle(self) -> None:

def main_loop(self) -> None:
"""The main loop of the worker process."""
signal.signal(signal.SIGINT, signal_handler)

while not self._end_process:
time.sleep(self._loop_interval)
try:
time.sleep(self._loop_interval)

self.receive_and_handle_control_messages()
self.receive_and_handle_control_messages()

self.worker_cycle()
self.worker_cycle()
except KeyboardInterrupt:
logger.info("Keyboard interrupt received")

self.send_process_state_change_message(
process_state=HordeProcessState.PROCESS_ENDING,
info="Process ending",
)

self.cleanup_and_exit()

logger.info("Process ended")
self.send_process_state_change_message(
process_state=HordeProcessState.PROCESS_ENDED,
info="Process ended",
)


def signal_handler(sig: int, frame: object) -> None:
print("You pressed Ctrl+C!")
7 changes: 0 additions & 7 deletions horde_worker_regen/process_management/messages.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,13 +62,6 @@ class HordeProcessState(enum.Enum):

EVALUATING_SAFETY = auto()

def can_accept_job(self) -> bool:
return (
self == HordeProcessState.WAITING_FOR_JOB
or self == HordeProcessState.INFERENCE_COMPLETE
or self == HordeProcessState.INFERENCE_FAILED
)


class HordeProcessMessage(BaseModel):
"""Process messages are sent from the child processes to the main process."""
Expand Down
Loading

0 comments on commit 525a047

Please sign in to comment.