add env vars for ewms worker resources

WIPACrepo · Jan 31, 2025 · 2dcea7f · 2dcea7f
1 parent 8139628
commit 2dcea7f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 14 deletions.
diff --git a/skydriver/config.py b/skydriver/config.py
@@ -5,7 +5,6 @@
 import logging
 from typing import Any, Optional
 
-import humanfriendly
 from wipac_dev_tools import from_environment_as_dataclass, logging_tools
 
 sdict = dict[str, Any]
@@ -14,10 +13,6 @@
 # Constants
 
 
-DEFAULT_WORKER_MEMORY_BYTES: int = humanfriendly.parse_size("8GB")
-DEFAULT_WORKER_DISK_BYTES: int = humanfriendly.parse_size("1GB")
-DEFAULT_MAX_WORKER_RUNTIME = 4 * 60 * 60
-
 SCAN_MIN_PRIORITY_TO_START_ASAP = 100
 
 # WARNING: these values must remain constant, they are cross-referenced in the db
@@ -80,7 +75,7 @@ class EnvConfig:
     K8S_TTL_SECONDS_AFTER_FINISHED: int = 10 * 60
     K8S_ACTIVE_DEADLINE_SECONDS: int = 24 * 60 * 60
     #
-    K8S_SCANNER_MEM_REQUEST: str = "1024M"  # note: this is also used as the limit
+    K8S_SCANNER_MEM_REQUEST__DEFAULT: str = "1024M"  # note: also used as the limit def.
     K8S_SCANNER_CPU_LIMIT: float = 1.0
     K8S_SCANNER_CPU_REQUEST: float = 0.10
     #
@@ -94,6 +89,11 @@ class EnvConfig:
     K8S_SCANNER_SIDECAR_S3_MEM_REQUEST: str = "1Mi"
     K8S_SCANNER_SIDECAR_S3_CPU_REQUEST: float = 0.10
 
+    # EWMS optional config
+    EWMS_WORKER_MEMORY__DEFAULT: str = "8GB"
+    EWMS_WORKER_DISK__DEFAULT: int = "1GB"
+    EWMS_MAX_WORKER_RUNTIME__DEFAULT: int = 4 * 60 * 60  # 4 hours
+
     # keycloak
     KEYCLOAK_OIDC_URL: str = ""
     KEYCLOAK_CLIENT_ID_SKYDRIVER_REST: str = ""

diff --git a/skydriver/ewms.py b/skydriver/ewms.py
@@ -74,7 +74,7 @@ async def request_workflow_on_ewms(
                 },
                 "worker_config": {
                     "do_transfer_worker_stdouterr": True,  # toggle?
-                    "max_worker_runtime": 6 * 60 * 60,  # 6 hours
+                    "max_worker_runtime": scan_request_obj["max_worker_runtime"],
                     "n_cores": 1,
                     "priority": scan_request_obj["priority"],
                     "worker_disk": scan_request_obj["worker_disk_bytes"],

diff --git a/skydriver/rest_handlers.py b/skydriver/rest_handlers.py
@@ -28,9 +28,6 @@
 
 from . import database, ewms, images, k8s
 from .config import (
-    DEFAULT_MAX_WORKER_RUNTIME,
-    DEFAULT_WORKER_DISK_BYTES,
-    DEFAULT_WORKER_MEMORY_BYTES,
     DebugMode,
     ENV,
     KNOWN_CLUSTERS,
@@ -353,13 +350,13 @@ async def post(self) -> None:
         arghand.add_argument(
             "scanner_server_memory",
             type=_data_size_parse,
-            default=humanfriendly.parse_size(ENV.K8S_SCANNER_MEM_REQUEST),
+            default=humanfriendly.parse_size(ENV.K8S_SCANNER_MEM_REQUEST__DEFAULT),
         )
         # client worker args
         arghand.add_argument(
             "worker_memory",
             type=_data_size_parse,
-            default=DEFAULT_WORKER_MEMORY_BYTES,
+            default=humanfriendly.parse_size(ENV.EWMS_WORKER_MEMORY__DEFAULT),
         )
         arghand.add_argument(  # NOTE - DEPRECATED
             "memory",
@@ -375,7 +372,7 @@ async def post(self) -> None:
         arghand.add_argument(
             "worker_disk",
             type=_data_size_parse,
-            default=DEFAULT_WORKER_DISK_BYTES,
+            default=humanfriendly.parse_size(ENV.EWMS_WORKER_DISK__DEFAULT),
         )
         arghand.add_argument(
             "cluster",
@@ -415,7 +412,7 @@ async def post(self) -> None:
         arghand.add_argument(
             "max_worker_runtime",
             type=int,
-            default=DEFAULT_MAX_WORKER_RUNTIME,
+            default=ENV.EWMS_MAX_WORKER_RUNTIME__DEFAULT,
         )
         arghand.add_argument(
             # TODO - remove when TMS is handling workforce-scaling