refactoring

Author: sanderegg

services/autoscaling/src/simcore_service_autoscaling/modules/cluster_scaling/_auto_scaling_core.py

@@ -49,7 +49,6 @@
     ec2_startup_script,
     find_selected_instance_type_for_task,
     get_machine_buffer_type,
-    node_host_name_from_ec2_private_dns,
     sort_drained_nodes,
 )
 from ...utils.rabbitmq import (
@@ -285,7 +284,7 @@ async def _try_attach_pending_ec2s(
     assert app_settings.AUTOSCALING_EC2_INSTANCES  # nosec
     for instance_data in cluster.pending_ec2s:
         try:
-            node_host_name = node_host_name_from_ec2_private_dns(
+            node_host_name = utils_ec2.node_host_name_from_ec2_private_dns(
                 instance_data.ec2_instance
             )
             if new_node := await utils_docker.find_node_with_name(

services/autoscaling/src/simcore_service_autoscaling/modules/dask.py

@@ -21,7 +21,7 @@
     DaskWorkerNotFoundError,
 )
 from ..models import AssociatedInstance, DaskTask, DaskTaskId
-from ..utils.cluster_scaling import (
+from ..utils.utils_ec2 import (
     node_host_name_from_ec2_private_dns,
     node_ip_from_ec2_private_dns,
 )

services/autoscaling/src/simcore_service_autoscaling/utils/cluster_scaling.py

@@ -1,7 +1,6 @@
 import functools
 import logging
-import re
-from typing import Final, TypeAlias
+from typing import TypeAlias
 
 from aws_library.ec2 import (
     EC2InstanceBootSpecific,
@@ -19,39 +18,11 @@
 )
 from ..core.settings import ApplicationSettings
 from ..models import AssociatedInstance
-from . import utils_docker
+from . import utils_docker, utils_ec2
 
-_EC2_INTERNAL_DNS_RE: Final[re.Pattern] = re.compile(r"^(?P<host_name>ip-[^.]+)\..+$")
 _logger = logging.getLogger(__name__)
 
 
-def node_host_name_from_ec2_private_dns(
-    ec2_instance_data: EC2InstanceData,
-) -> str:
-    """returns the node host name 'ip-10-2-3-22' from the ec2 private dns
-    Raises:
-        Ec2InvalidDnsNameError: if the dns name does not follow the expected pattern
-    """
-    if match := re.match(_EC2_INTERNAL_DNS_RE, ec2_instance_data.aws_private_dns):
-        host_name: str = match.group("host_name")
-        return host_name
-    raise Ec2InvalidDnsNameError(aws_private_dns_name=ec2_instance_data.aws_private_dns)
-
-
-def node_ip_from_ec2_private_dns(
-    ec2_instance_data: EC2InstanceData,
-) -> str:
-    """returns the node ipv4 from the ec2 private dns string
-    Raises:
-        Ec2InvalidDnsNameError: if the dns name does not follow the expected pattern
-    """
-    return (
-        node_host_name_from_ec2_private_dns(ec2_instance_data)
-        .removeprefix("ip-")
-        .replace("-", ".")
-    )
-
-
 async def associate_ec2_instances_with_nodes(
     nodes: list[Node], ec2_instances: list[EC2InstanceData]
 ) -> tuple[list[AssociatedInstance], list[EC2InstanceData]]:
@@ -65,7 +36,9 @@ def _find_node_with_name(node: Node) -> bool:
 
     for instance_data in ec2_instances:
         try:
-            docker_node_name = node_host_name_from_ec2_private_dns(instance_data)
+            docker_node_name = utils_ec2.node_host_name_from_ec2_private_dns(
+                instance_data
+            )
         except Ec2InvalidDnsNameError:
             _logger.exception("Unexpected EC2 private dns name")
             non_associated_instances.append(instance_data)

services/autoscaling/src/simcore_service_autoscaling/utils/utils_ec2.py

@@ -1,18 +1,27 @@
 """Free helper functions for AWS API"""
 
 import logging
+import re
 from collections import OrderedDict
 from collections.abc import Callable
 from textwrap import dedent
+from typing import Final
 
 from aws_library.ec2 import AWSTagKey, AWSTagValue, EC2InstanceType, EC2Tags, Resources
+from aws_library.ec2._models import EC2InstanceData
 from common_library.json_serialization import json_dumps
 
 from .._meta import VERSION
-from ..core.errors import ConfigurationError, TaskBestFittingInstanceNotFoundError
+from ..core.errors import (
+    ConfigurationError,
+    Ec2InvalidDnsNameError,
+    TaskBestFittingInstanceNotFoundError,
+)
 from ..core.settings import ApplicationSettings
 
-logger = logging.getLogger(__name__)
+_logger = logging.getLogger(__name__)
+
+_EC2_INTERNAL_DNS_RE: Final[re.Pattern] = re.compile(r"^(?P<host_name>ip-[^.]+)\..+$")
 
 
 def get_ec2_tags_dynamic(app_settings: ApplicationSettings) -> EC2Tags:
@@ -105,3 +114,30 @@ def find_best_fitting_ec2_instance(
         raise TaskBestFittingInstanceNotFoundError(needed_resources=resources)
 
     return instance
+
+
+def node_host_name_from_ec2_private_dns(
+    ec2_instance_data: EC2InstanceData,
+) -> str:
+    """returns the node host name 'ip-10-2-3-22' from the ec2 private dns
+    Raises:
+        Ec2InvalidDnsNameError: if the dns name does not follow the expected pattern
+    """
+    if match := re.match(_EC2_INTERNAL_DNS_RE, ec2_instance_data.aws_private_dns):
+        host_name: str = match.group("host_name")
+        return host_name
+    raise Ec2InvalidDnsNameError(aws_private_dns_name=ec2_instance_data.aws_private_dns)
+
+
+def node_ip_from_ec2_private_dns(
+    ec2_instance_data: EC2InstanceData,
+) -> str:
+    """returns the node ipv4 from the ec2 private dns string
+    Raises:
+        Ec2InvalidDnsNameError: if the dns name does not follow the expected pattern
+    """
+    return (
+        node_host_name_from_ec2_private_dns(ec2_instance_data)
+        .removeprefix("ip-")
+        .replace("-", ".")
+    )

services/autoscaling/tests/unit/test_modules_cluster_scaling_dynamic.py

@@ -66,15 +66,15 @@
     AutoscalingDocker,
     get_docker_client,
 )
-from simcore_service_autoscaling.utils.cluster_scaling import (
-    node_host_name_from_ec2_private_dns,
-)
 from simcore_service_autoscaling.utils.utils_docker import (
     _OSPARC_NODE_EMPTY_DATETIME_LABEL_KEY,
     _OSPARC_NODE_TERMINATION_PROCESS_LABEL_KEY,
     _OSPARC_SERVICE_READY_LABEL_KEY,
     _OSPARC_SERVICES_READY_DATETIME_LABEL_KEY,
 )
+from simcore_service_autoscaling.utils.utils_ec2 import (
+    node_host_name_from_ec2_private_dns,
+)
 from types_aiobotocore_ec2.client import EC2Client
 from types_aiobotocore_ec2.literals import InstanceStateNameType, InstanceTypeType
 from types_aiobotocore_ec2.type_defs import FilterTypeDef, InstanceTypeDef

services/autoscaling/tests/unit/test_utils_cluster_scaling.py

@@ -19,14 +19,12 @@
 from pydantic import TypeAdapter
 from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
 from pytest_simcore.helpers.typing_env import EnvVarsDict
-from simcore_service_autoscaling.core.errors import Ec2InvalidDnsNameError
 from simcore_service_autoscaling.core.settings import ApplicationSettings
 from simcore_service_autoscaling.models import AssociatedInstance, EC2InstanceData
 from simcore_service_autoscaling.utils.cluster_scaling import (
     associate_ec2_instances_with_nodes,
     ec2_startup_script,
     get_machine_buffer_type,
-    node_host_name_from_ec2_private_dns,
     sort_drained_nodes,
 )
 from simcore_service_autoscaling.utils.utils_docker import (
@@ -52,34 +50,6 @@ def _creator(**overrides) -> DockerNode:
     return _creator
 
 
-@pytest.mark.parametrize(
-    "aws_private_dns, expected_host_name",
-    [
-        ("ip-10-12-32-3.internal-data", "ip-10-12-32-3"),
-        ("ip-10-12-32-32.internal-data", "ip-10-12-32-32"),
-        ("ip-10-0-3-129.internal-data", "ip-10-0-3-129"),
-        ("ip-10-0-3-12.internal-data", "ip-10-0-3-12"),
-    ],
-)
-def test_node_host_name_from_ec2_private_dns(
-    fake_ec2_instance_data: Callable[..., EC2InstanceData],
-    aws_private_dns: str,
-    expected_host_name: str,
-):
-    instance = fake_ec2_instance_data(
-        aws_private_dns=aws_private_dns,
-    )
-    assert node_host_name_from_ec2_private_dns(instance) == expected_host_name
-
-
-def test_node_host_name_from_ec2_private_dns_raises_with_invalid_name(
-    fake_ec2_instance_data: Callable[..., EC2InstanceData], faker: Faker
-):
-    instance = fake_ec2_instance_data(aws_private_dns=faker.name())
-    with pytest.raises(Ec2InvalidDnsNameError):
-        node_host_name_from_ec2_private_dns(instance)
-
-
 @pytest.mark.parametrize("valid_ec2_dns", [True, False])
 async def test_associate_ec2_instances_with_nodes_with_no_correspondence(
     fake_ec2_instance_data: Callable[..., EC2InstanceData],

services/autoscaling/tests/unit/test_utils_ec2.py

@@ -3,18 +3,24 @@
 # pylint: disable=unused-variable
 
 
+from collections.abc import Callable
+
 import pytest
 from aws_library.ec2 import EC2InstanceType, Resources
+from aws_library.ec2._models import EC2InstanceData
 from faker import Faker
 from pydantic import ByteSize
 from simcore_service_autoscaling.core.errors import (
     ConfigurationError,
+    Ec2InvalidDnsNameError,
     TaskBestFittingInstanceNotFoundError,
 )
 from simcore_service_autoscaling.utils.utils_ec2 import (
     closest_instance_policy,
     compose_user_data,
     find_best_fitting_ec2_instance,
+    node_host_name_from_ec2_private_dns,
+    node_ip_from_ec2_private_dns,
 )
 
 
@@ -70,3 +76,59 @@ def test_compose_user_data(faker: Faker):
     user_data = compose_user_data(command)
     assert user_data.startswith("#!/bin/bash")
     assert command in user_data
+
+
+@pytest.mark.parametrize(
+    "aws_private_dns, expected_host_name",
+    [
+        ("ip-10-12-32-3.internal-data", "ip-10-12-32-3"),
+        ("ip-10-12-32-32.internal-data", "ip-10-12-32-32"),
+        ("ip-10-0-3-129.internal-data", "ip-10-0-3-129"),
+        ("ip-10-0-3-12.internal-data", "ip-10-0-3-12"),
+    ],
+)
+def test_node_host_name_from_ec2_private_dns(
+    fake_ec2_instance_data: Callable[..., EC2InstanceData],
+    aws_private_dns: str,
+    expected_host_name: str,
+):
+    instance = fake_ec2_instance_data(
+        aws_private_dns=aws_private_dns,
+    )
+    assert node_host_name_from_ec2_private_dns(instance) == expected_host_name
+
+
+def test_node_host_name_from_ec2_private_dns_raises_with_invalid_name(
+    fake_ec2_instance_data: Callable[..., EC2InstanceData], faker: Faker
+):
+    instance = fake_ec2_instance_data(aws_private_dns=faker.name())
+    with pytest.raises(Ec2InvalidDnsNameError):
+        node_host_name_from_ec2_private_dns(instance)
+
+
+@pytest.mark.parametrize(
+    "aws_private_dns, expected_host_name",
+    [
+        ("ip-10-12-32-3.internal-data", "10.12.32.3"),
+        ("ip-10-12-32-32.internal-data", "10.12.32.32"),
+        ("ip-10-0-3-129.internal-data", "10.0.3.129"),
+        ("ip-10-0-3-12.internal-data", "10.0.3.12"),
+    ],
+)
+def test_node_ip_from_ec2_private_dns(
+    fake_ec2_instance_data: Callable[..., EC2InstanceData],
+    aws_private_dns: str,
+    expected_host_name: str,
+):
+    instance = fake_ec2_instance_data(
+        aws_private_dns=aws_private_dns,
+    )
+    assert node_ip_from_ec2_private_dns(instance) == expected_host_name
+
+
+def test_node_ip_from_ec2_private_dns_raises_with_invalid_name(
+    fake_ec2_instance_data: Callable[..., EC2InstanceData], faker: Faker
+):
+    instance = fake_ec2_instance_data(aws_private_dns=faker.name())
+    with pytest.raises(Ec2InvalidDnsNameError):
+        node_ip_from_ec2_private_dns(instance)