WIP: Add nodegroup annotation to the pod

docs/CHANGELOG.rst

@@ -10,6 +10,11 @@ This project adheres to `Semantic Versioning <http://semver.org/>`_.
 Unreleased
 ==========
 
+Added
+-----
+- Add node affinity for instance-category using the ratio between instance to
+  process resource requirements.
+
 Changed
 -------
 - Change fallback for storage resource from 200GB to 10GB

resolwe/flow/managers/workload_connectors/kubernetes.py

@@ -106,6 +106,29 @@ def sanitize_kubernetes_label(label: str, trim_end: bool = True) -> str:
     return sanitized_label
 
 
+def memory_to_mb(memory_str):
+    # If the input is an integer, assume it's in gigabytes (GB)
+    if isinstance(memory_str, int):
+        return memory_str * 1024  # Default to gigabytes
+
+    # If no unit is provided, assume it's in gigabytes (GB)
+    if memory_str[-1].isdigit():
+        return int(memory_str) * 1024  # Default to gigabytes
+
+    # Check the last character for the unit (M for MB, G for GB)
+    unit = memory_str[-1].upper()
+    value = int(memory_str[:-1])  # The numerical part
+
+    if unit == "G":
+        return value * 1024  # Convert GB to MB
+    elif unit == "M":
+        return value  # It's already in MB
+    else:
+        raise ValueError(
+            "Invalid memory unit. Use 'M' for megabytes or 'G' for gigabytes."
+        )
+
+
 class ConfigLocation(Enum):
     """The enum specifying where to read the configuration from."""
 
@@ -610,6 +633,104 @@ def optimize_job_scheduling(self, data: Data, job_description: dict):
                     }
                 }
 
+    def append_node_affinity(self, job_description: dict, process_requirements: dict):
+        """
+        Appends nodeAffinity based on CPU and memory ratio to an existing job description if instance_types are defined
+        within FLOW_KUBERNETES_AFFINITY.
+
+        :param job_description: dict containing the job spec where the affinity will be appended.
+        :param process_requirements: dict containing 'cpu' and 'memory' requirements of the process.
+        :return: None (modifies job_description in place).
+        """
+
+        # Extract kubernetes_affinity and instance_types from settings
+        if kubernetes_affinity := getattr(settings, "FLOW_KUBERNETES_AFFINITY", None):
+            instance_types = kubernetes_affinity.get("instance_types", {})
+        else:
+            instance_types = {}
+
+        # Proceed only if instance_types is not empty
+        if not instance_types:
+            return
+
+        best_match_category = None
+        lowest_process_hourly_rate = float("inf")
+
+        process_cpu = process_requirements["cpu"]
+        process_memory = process_requirements["memory"]
+
+        # Sort instance types by hourly_rate to prefer cheaper options.
+        # This should avoid spawning 128GB instance type for a single 32GB
+        # process, but support larger instances when needed. This is not fully
+        # optimized if smaller instance types are available, since it would
+        # create a lot of instance fragmentation. Smaller nodes should have
+        # at least 8 cpus in this configuration.
+        sorted_instance_types = sorted(
+            instance_types.items(), key=lambda x: x[1].get("hourly_rate", float("inf"))
+        )
+
+        for instance_type, attributes in sorted_instance_types:
+            instance_cpu = int(attributes["cpu"])
+            instance_memory = memory_to_mb(attributes["memory"])
+            instance_category = attributes["instance-category"]
+            instance_hourly_rate = float(attributes["hourly_rate"])
+
+            # Skip instance types that don't meet the process requirements
+            # This needs to be here, since some instance category may not have
+            # a sufficiently large instance type available.
+            if instance_cpu < process_cpu or instance_memory < process_memory:
+                continue
+
+            # In practice this should be equivalent to:
+            # compute-optimized for ratio <~ 2.25
+            # general purpose for 2.25 <~ ratio <= 5.25
+            # memory-optimized for ratio > 5.25
+            process_instance_density = min(
+                instance_cpu / process_cpu, instance_memory / process_memory
+            )
+            process_hourly_rate = instance_hourly_rate / process_instance_density
+
+            if process_hourly_rate < lowest_process_hourly_rate:
+                lowest_process_hourly_rate = process_hourly_rate
+                best_match_category = instance_category
+
+        if not best_match_category:
+            raise ValueError(
+                "No suitable instance type found for the given process requirements."
+            )
+
+        # Create the new match expression
+        new_match_expression = {
+            "key": "instance-category",
+            "operator": "In",
+            "values": [best_match_category],
+        }
+
+        job_spec = job_description["spec"]["template"]["spec"]
+
+        # Ensure nodeAffinity structure exists in job_description
+        if "affinity" not in job_spec:
+            job_spec["affinity"] = {}
+
+        if "nodeAffinity" not in job_spec["affinity"]:
+            job_spec["affinity"]["nodeAffinity"] = {
+                "requiredDuringSchedulingIgnoredDuringExecution": {
+                    "nodeSelectorTerms": []
+                }
+            }
+
+        node_affinity = job_spec["affinity"]["nodeAffinity"]
+        node_selector_terms = node_affinity[
+            "requiredDuringSchedulingIgnoredDuringExecution"
+        ]["nodeSelectorTerms"]
+
+        # If there are no existing terms, create one
+        if not node_selector_terms:
+            node_selector_terms.append({"matchExpressions": []})
+
+        # Append the new match expression to the first nodeSelectorTerm's matchExpressions
+        node_selector_terms[0]["matchExpressions"].append(new_match_expression)
+
     def _get_overcommit_factors(self, data: Data) -> dict:
         """Get the overcommit settings for CPU and memory.
 
@@ -670,6 +791,8 @@ def start(self, data: Data, listener_connection: Tuple[str, str, str]):
         requests["cpu"] = limits["cores"] * overcommit_factors["cpu"]
         limits["cpu"] = limits.pop("cores") + 1
 
+        process_requirements = {"cpu": limits["cpu"], "memory": limits["memory"]}
+
         # The memory in the database is stored in megabytes but the kubertenes
         # requires memory in bytes.
         # We request less memory than stored in the database and set limit at 10% more
@@ -852,6 +975,19 @@ def start(self, data: Data, listener_connection: Tuple[str, str, str]):
             },
         }
         self.optimize_job_scheduling(data, job_description)
+        try:
+            self.append_node_affinity(job_description, process_requirements)
+        except Exception as error:
+            # This operation is not critical for the job to run, so we log the
+            # error and continue. The job will still run, but the node affinity
+            # may be sub-optimal.
+            logger.exception(
+                __(
+                    "Kubernetes instance category affinity assignemt failed: '{}'.",
+                    error,
+                )
+            )
+
         start_time = time.time()
 
         processing_name = constants.PROCESSING_VOLUME_NAME