From d0b09f2cc36635757e9f5dbef04e511fe4ee223e Mon Sep 17 00:00:00 2001
From: tazlin <tazlin.on.github@gmail.com>
Date: Sun, 4 Aug 2024 19:55:13 -0400
Subject: [PATCH] feat: allow setting worker name via env var or cli arg (#245)

* fix: support setting worker name from bridge script
* fix: check for worker name not none in startup script
* fix: log dreamer name
* fix: use field validator for worker name override instead
* fix: don't use missing function
* style: fix
* fix: check worker name in type safe way
---
 horde_worker_regen/bridge_data/data_model.py | 15 ++++++++++++++-
 horde_worker_regen/run_worker.py             | 10 ++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/horde_worker_regen/bridge_data/data_model.py b/horde_worker_regen/bridge_data/data_model.py
index 992ad4bc..3713f002 100644
--- a/horde_worker_regen/bridge_data/data_model.py
+++ b/horde_worker_regen/bridge_data/data_model.py
@@ -7,7 +7,7 @@
 
 from horde_sdk.ai_horde_worker.bridge_data import CombinedHordeBridgeData
 from loguru import logger
-from pydantic import Field, model_validator
+from pydantic import Field, field_validator, model_validator
 from ruamel.yaml import YAML
 
 from horde_worker_regen.consts import TOTAL_LORA_DOWNLOAD_TIMEOUT
@@ -131,6 +131,19 @@ def validate_performance_modes(self) -> reGenBridgeData:
 
         return self
 
+    @field_validator("dreamer_worker_name", mode="after")
+    def validate_dreamer_worker_name(cls, value: str) -> str:
+        """Apply the environment variable override for the `dreamer_worker_name` field."""
+        AIWORKER_DREAMER_WORKER_NAME = os.getenv("AIWORKER_DREAMER_WORKER_NAME")
+        if AIWORKER_DREAMER_WORKER_NAME:
+            logger.warning(
+                "AIWORKER_DREAMER_WORKER_NAME environment variable is set. This will override the value for "
+                "`dreamer_worker_name` in the config file.",
+            )
+            return AIWORKER_DREAMER_WORKER_NAME
+
+        return value
+
     def prepare_custom_models(self) -> None:
         """Prepare the custom models."""
         if os.getenv("HORDELIB_CUSTOM_MODELS"):
diff --git a/horde_worker_regen/run_worker.py b/horde_worker_regen/run_worker.py
index 3e1a9ea5..55b02c9e 100644
--- a/horde_worker_regen/run_worker.py
+++ b/horde_worker_regen/run_worker.py
@@ -153,11 +153,21 @@ def init() -> None:
         default=False,
         help="Enable AMD GPU-specific optimisations",
     )
+    parser.add_argument(
+        "-n",
+        "--worker-name",
+        type=str,
+        default=None,
+        help="Override the worker name from the config file, for running multiple workers on one machine",
+    )
 
     args = parser.parse_args()
 
     os.environ["HORDE_SDK_DISABLE_CUSTOM_SINKS"] = "1"
 
+    if args.worker_name:
+        os.environ["AIWORKER_DREAMER_WORKER_NAME"] = args.worker_name
+
     from horde_worker_regen.load_env_vars import load_env_vars_from_config
 
     if not args.load_config_from_env_vars: