🎨Computational backend: Make sure the number of threads of a dask-worker is computed for autoscaling 🚨🚨🚨 (#8423)

Co-authored-by: Copilot <[email protected]>

Author: sanderegg

packages/aws-library/src/aws_library/ec2/__init__.py

@@ -17,6 +17,7 @@
     EC2InstanceData,
     EC2InstanceType,
     EC2Tags,
+    GenericResourceValueType,
     Resources,
 )
 
@@ -36,6 +37,7 @@
     "EC2NotConnectedError",
     "EC2RuntimeError",
     "EC2Tags",
+    "GenericResourceValueType",
     "Resources",
     "SimcoreEC2API",
 )

packages/aws-library/src/aws_library/ec2/_models.py

@@ -1,3 +1,4 @@
+import contextlib
 import datetime
 import re
 import tempfile
@@ -14,45 +15,171 @@
     Field,
     NonNegativeFloat,
     NonNegativeInt,
+    StrictFloat,
+    StrictInt,
     StringConstraints,
+    TypeAdapter,
+    ValidationError,
     field_validator,
 )
 from pydantic.config import JsonDict
 from types_aiobotocore_ec2.literals import InstanceStateNameType, InstanceTypeType
 
+GenericResourceValueType: TypeAlias = StrictInt | StrictFloat | str
+
 
 class Resources(BaseModel, frozen=True):
     cpus: NonNegativeFloat
     ram: ByteSize
+    generic_resources: Annotated[
+        dict[str, GenericResourceValueType],
+        Field(
+            default_factory=dict,
+            description=(
+                "Arbitrary additional resources (e.g. {'threads': 8}). "
+                "Numeric values are treated as quantities and participate in add/sub/compare."
+            ),
+        ),
+    ] = DEFAULT_FACTORY
 
     @classmethod
     def create_as_empty(cls) -> "Resources":
         return cls(cpus=0, ram=ByteSize(0))
 
     def __ge__(self, other: "Resources") -> bool:
-        return self.cpus >= other.cpus and self.ram >= other.ram
+        """operator for >= comparison
+        if self has greater or equal resources than other, returns True
+        This will return True only if all of the resources in self are greater or equal to other
+
+        Note that generic_resources are compared only if they are numeric
+        Non-numeric generic resources must be equal in both or only defined in self
+        to be considered greater or equal
+        """
+        if self == other:
+            return True
+        return self > other
 
     def __gt__(self, other: "Resources") -> bool:
-        return self.cpus > other.cpus or self.ram > other.ram
+        """operator for > comparison
+        if self has resources greater than other, returns True
+        This will return True only if all of the resources in self are greater than other
+
+        Note that generic_resources are compared only if they are numeric
+        Non-numeric generic resources must only be defined in self
+        to be considered greater
+        """
+        if (self.cpus < other.cpus) or (self.ram < other.ram):
+            return False
+
+        keys = set(self.generic_resources) | set(other.generic_resources)
+        for k in keys:
+            a = self.generic_resources.get(k)
+            b = other.generic_resources.get(k)
+            if a is None:
+                return False
+            if b is None:
+                # a is greater as b is not defined
+                continue
+            if isinstance(a, int | float) and isinstance(b, int | float):
+                if a < b:
+                    return False
+            else:
+                # remaining options is a is str and b is str or mixed types
+                # NOTE: we cannot compare strings unless they are equal or some kind of boolean (e.g. "true", "false", "yes", "no", "1", "0")
+                assert isinstance(a, str)  # nosec
+                assert isinstance(b, int | float | str)  # nosec
+                # let's try to get a boolean out of the values to compare them
+                with contextlib.suppress(ValidationError):
+                    a_as_boolean = TypeAdapter(bool).validate_python(a)
+                    b_as_boolean = TypeAdapter(bool).validate_python(b)
+                    if not a_as_boolean and b_as_boolean:
+                        return False
+
+        # here we have either everything greater or equal or non-comparable strings
+
+        return self != other
 
     def __add__(self, other: "Resources") -> "Resources":
+        """operator for adding two Resources
+        Note that only numeric generic resources are added
+        Non-numeric generic resources are ignored
+        """
+        merged: dict[str, GenericResourceValueType] = {}
+        keys = set(self.generic_resources) | set(other.generic_resources)
+        for k in keys:
+            a = self.generic_resources.get(k)
+            b = other.generic_resources.get(k)
+            # adding non numeric values does not make sense, so we skip those for the resulting resource
+            if isinstance(a, int | float) and isinstance(b, int | float):
+                merged[k] = a + b
+            elif a is None and isinstance(b, int | float):
+                merged[k] = b
+            elif b is None and isinstance(a, int | float):
+                merged[k] = a
+
         return Resources.model_construct(
-            **{
-                key: a + b
-                for (key, a), b in zip(
-                    self.model_dump().items(), other.model_dump().values(), strict=True
-                )
-            }
+            cpus=self.cpus + other.cpus,
+            ram=self.ram + other.ram,
+            generic_resources=merged,
         )
 
     def __sub__(self, other: "Resources") -> "Resources":
+        """operator for subtracting two Resources
+        Note that only numeric generic resources are subtracted
+        Non-numeric generic resources are ignored
+        """
+        merged: dict[str, GenericResourceValueType] = {}
+        keys = set(self.generic_resources) | set(other.generic_resources)
+        for k in keys:
+            a = self.generic_resources.get(k)
+            b = other.generic_resources.get(k)
+            # subtracting non numeric values does not make sense, so we skip those for the resulting resource
+            if isinstance(a, int | float) and isinstance(b, int | float):
+                merged[k] = a - b
+            elif a is None and isinstance(b, int | float):
+                merged[k] = -b
+            elif b is None and isinstance(a, int | float):
+                merged[k] = a
+
         return Resources.model_construct(
-            **{
-                key: a - b
-                for (key, a), b in zip(
-                    self.model_dump().items(), other.model_dump().values(), strict=True
-                )
-            }
+            cpus=self.cpus - other.cpus,
+            ram=self.ram - other.ram,
+            generic_resources=merged,
+        )
+
+    def __hash__(self) -> int:
+        """Deterministic hash including cpus, ram (in bytes) and generic_resources."""
+        # sort generic_resources items to ensure order-independent hashing
+        generic_items: tuple[tuple[str, GenericResourceValueType], ...] = tuple(
+            sorted(self.generic_resources.items())
+        )
+        return hash((self.cpus, self.ram, generic_items))
+
+    def as_flat_dict(self) -> dict[str, int | float | str]:
+        """Like model_dump, but flattens generic_resources to top level keys"""
+        base = self.model_dump()
+        base.update(base.pop("generic_resources"))
+        return base
+
+    @classmethod
+    def from_flat_dict(
+        cls,
+        data: dict[str, int | float | str],
+        *,
+        mapping: dict[str, str] | None = None,
+    ) -> "Resources":
+        """Inverse of as_flat_dict with optional key mapping"""
+        mapped_data = data
+        if mapping:
+            mapped_data = {mapping.get(k, k): v for k, v in data.items()}
+        generic_resources = {
+            k: v for k, v in mapped_data.items() if k not in {"cpus", "ram"}
+        }
+
+        return cls(
+            cpus=float(mapped_data.get("cpus", 0)),
+            ram=ByteSize(mapped_data.get("ram", 0)),
+            generic_resources=generic_resources,
         )
 
     @field_validator("cpus", mode="before")
@@ -174,8 +301,9 @@ def validate_bash_calls(cls, v):
                 temp_file.flush()
                 # NOTE: this will not capture runtime errors, but at least some syntax errors such as invalid quotes
                 sh.bash(
-                    "-n", temp_file.name
-                )  # pyright: ignore[reportCallIssue]  # sh is untyped, but this call is safe for bash syntax checking
+                    "-n",
+                    temp_file.name,  # pyright: ignore[reportCallIssue] - sh is untyped but safe for bash syntax checking
+                )
         except sh.ErrorReturnCode as exc:
             msg = f"Invalid bash call in custom_boot_scripts: {v}, Error: {exc.stderr}"
             raise ValueError(msg) from exc