Skip to content

✨♻️Autoscaling/dask nthreads2 🚨 🚨 🚨 #8558

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: master
Choose a base branch
from
+1,103 −229
Open

✨♻️Autoscaling/dask nthreads2 🚨 🚨 🚨 #8558

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
f8b7f33
🎨Computational backend: Make sure the number of threads of a dask-wor…
sanderegg Oct 24, 2025
915cefa
🐛Autoscaling: fixes unknown passing type to dask-scheduler (#8556)
sanderegg Oct 24, 2025
a579a82
better logs
sanderegg Oct 25, 2025
e3d2a06
fix so we do not drain too fast
sanderegg Oct 25, 2025
5e0a39c
Revert "fix so we do not drain too fast"
sanderegg Oct 25, 2025
49b89da
fix
sanderegg Oct 25, 2025
ffbbfb8
fix
sanderegg Oct 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions packages/aws-library/src/aws_library/ec2/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
EC2InstanceData,
EC2InstanceType,
EC2Tags,
GenericResourceValueType,
Resources,
)

Expand All @@ -36,6 +37,7 @@
"EC2NotConnectedError",
"EC2RuntimeError",
"EC2Tags",
"GenericResourceValueType",
"Resources",
"SimcoreEC2API",
)
Expand Down
160 changes: 144 additions & 16 deletions packages/aws-library/src/aws_library/ec2/_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import contextlib
import datetime
import re
import tempfile
Expand All @@ -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"))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MIONR: just as a precaution in case one day you add another object-like field , here i would do it in a more general fashion. I know right now is perhaps an overkill, but it might be handy to add some helpers like this in common_library.dict_tools so that afterwards we could do something like

common_library.dict_tools.flatten(self.model_dump(mode='json'))

A draft courtesy of chatty (double check first!)

def _flatten(obj, prefix=""):
        items = {}
        if isinstance(obj, dict):
            for k, v in obj.items():
                items.update(_flatten(v, f"{prefix}{k}." if prefix else k))
        elif isinstance(obj, list):
            for i, v in enumerate(obj):
                items.update(_flatten(v, f"{prefix}{i}." if prefix else f"{i}"))
        else:
            items[prefix] = obj
        return items

def as_flat_dict(self) -> dict[str, int | float | str]:
    """Recursively flattens all nested dicts/lists into a single-level dict."""
    return _flatten(self.model_dump(mode='json'))

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")
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading