Replace raw config dicts with TypedDicts (#18)

* Add Typing for custom ResourceOptions

by introducing TypedDicts for K8s, Docker, and Ray ResourceOptions.

their filling value can be `None`. Then, we remove the `None`s by
using the `remove_none_values` function.

We add generic typing to the `remove_none_value` function and coerce back to the input type for consisten type hints.

As the evaluation of the functional form is not deferred and using
future imports has no effect, we use Union instead of | (pipe) to be compatible with
Python 3.9

* Use single K8sResourceOption type

as there is no conceptual difference within the two. As they may use
different resources (in the limit and request) we do make the parameters
optional. This also reflects the code logic which filters `None` values.

We now also handle LIMIT and REQUEST types the same. This commit leaves
a TODO that we have to decide if separate logic is warranted or not.
Until then we do not remove the logical branches from the code, but
right as of this commit both branches act the same.

Author: maxmynter

src/jobs/job.py

@@ -7,14 +7,14 @@
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Literal
+from typing import Any, Callable, TypedDict, Union
 
 import docker.types
 
 from jobs.assembler import config
 from jobs.assembler.renderers import RENDERERS
 from jobs.image import Image
-from jobs.types import AnyPath
+from jobs.types import AnyPath, K8sResourceKind
 from jobs.util import remove_none_values, run_command, to_rational
 
 
@@ -23,56 +23,73 @@ class BuildMode(enum.Enum):
     DOCKERFILE = "dockerfile"
 
 
+class DockerResourceOptions(TypedDict):
+    mem_limit: str | None
+    nano_cpus: float | None
+    device_requests: list[docker.types.DeviceRequest] | None
+
+
+# Functional definition of TypedDict to enable special characters in dict keys
+K8sResourceOptions = TypedDict(
+    "K8sResourceOptions",
+    {
+        "cpu": Union[str, None],
+        "memory": Union[str, None],
+        "nvidia.com/gpu": Union[int, None],
+    },
+    total=False,
+)
+
+
+class RayResourceOptions(TypedDict):
+    entrypoint_memory: int | None
+    entrypoint_num_cpus: int | None
+    entrypoint_num_gpus: int | None
+
+
 @dataclass(frozen=True)
 class ResourceOptions:
     memory: str | None = None
     cpu: str | None = None
     gpu: int | None = None
 
-    def to_docker(self) -> dict[str, Any]:
-        return remove_none_values(
-            {
-                "mem_limit": (
-                    str(int(to_rational(self.memory))) if self.memory else None
-                ),
-                "nano_cpus": int(to_rational(self.cpu) * 10**9) if self.cpu else None,
-                "device_requests": (
-                    [
-                        docker.types.DeviceRequest(
-                            capabilities=[["gpu"]],
-                            count=self.gpu,
-                        )
-                    ]
-                    if self.gpu
-                    else None
-                ),
-            }
-        )
+    def to_docker(self) -> DockerResourceOptions:
+        options: DockerResourceOptions = {
+            "mem_limit": str(int(to_rational(self.memory))) if self.memory else None,
+            "nano_cpus": int(to_rational(self.cpu) * 10**9) if self.cpu else None,
+            "device_requests": (
+                [
+                    docker.types.DeviceRequest(
+                        capabilities=[["gpu"]],
+                        count=self.gpu,
+                    )
+                ]
+                if self.gpu
+                else None
+            ),
+        }
+        return remove_none_values(options)
 
     def to_kubernetes(
-        self, kind: Literal["requests", "limits"] = "requests"
-    ) -> dict[str, str]:
-        if kind == "requests":
-            return remove_none_values(
-                {
-                    "cpu": self.cpu,
-                    "memory": self.memory,
-                    "nvidia.com/gpu": self.gpu,
-                }
-            )
-        elif kind == "limits":
-            return remove_none_values({"nvidia.com/gpu": self.gpu})
-
-    def to_ray(self) -> dict[str, Any]:
-        return remove_none_values(
-            {
-                "entrypoint_memory": (
-                    int(to_rational(self.memory)) if self.memory else None
-                ),
-                "entrypoint_num_cpus": int(to_rational(self.cpu)) if self.cpu else None,
-                "entrypoint_num_gpus": self.gpu,
-            }
-        )
+        self, kind: K8sResourceKind = K8sResourceKind.REQUESTS
+    ) -> K8sResourceOptions:
+        # TODO: Currently kind is not accessed and the logic for "request" and "limit" is the same.
+        # Down the road we have to decide if we want to keep it that way (and get rid of the distinction and arguments),
+        # or if it makes sense for us to distinguish both cases.
+        options: K8sResourceOptions = {
+            "cpu": self.cpu or None,
+            "memory": self.memory or None,
+            "nvidia.com/gpu": self.gpu or None,
+        }
+        return remove_none_values(options)
+
+    def to_ray(self) -> RayResourceOptions:
+        options: RayResourceOptions = {
+            "entrypoint_memory": int(to_rational(self.memory)) if self.memory else None,
+            "entrypoint_num_cpus": int(to_rational(self.cpu)) if self.cpu else None,
+            "entrypoint_num_gpus": self.gpu or None,
+        }
+        return remove_none_values(options)
 
 
 @dataclass(frozen=True)

src/jobs/runner.py

@@ -18,7 +18,9 @@
 
 import jobs
 from jobs import Image, Job
-from jobs.util import sanitize_rfc1123_domain_name
+from jobs.job import DockerResourceOptions
+from jobs.types import K8sResourceKind
+from jobs.util import remove_none_values, sanitize_rfc1123_domain_name
 
 JOBS_EXECUTE_CMD = "jobs_execute"
 
@@ -44,15 +46,19 @@ def __init__(self):
     def run(self, job: Job, image: Image) -> None:
         command = _make_executor_command(job)
 
-        resource_kwargs = {}
+        resource_kwargs: DockerResourceOptions = {
+            "mem_limit": None,
+            "nano_cpus": None,
+            "device_requests": None,
+        }
         if job.options and (res := job.options.resources):
             resource_kwargs = res.to_docker()
 
         container: docker.api.client.ContainerApiMixin = self._client.containers.run(
             image=image.tag,
             command=command,
             detach=True,
-            **resource_kwargs,
+            **remove_none_values(resource_kwargs),
         )
 
         exit_code = container.wait()
@@ -83,12 +89,14 @@ def _make_job_crd(self, job: Job, image: Image) -> client.V1Job:
             image_pull_policy="IfNotPresent",
             name="dummy-job",
             command=_make_executor_command(job),
-            resources={
-                "requests": res.to_kubernetes(kind="requests"),
-                "limits": res.to_kubernetes(kind="limits"),
-            }
-            if job.options and (res := job.options.resources)
-            else None,
+            resources=(
+                {
+                    "requests": res.to_kubernetes(kind=K8sResourceKind.REQUESTS),
+                    "limits": res.to_kubernetes(kind=K8sResourceKind.LIMITS),
+                }
+                if job.options and (res := job.options.resources)
+                else None
+            ),
         )
 
         # Job template

src/jobs/types.py

@@ -1,4 +1,10 @@
 import os
+from enum import Enum
 from typing import Union
 
 AnyPath = Union[os.PathLike[str], str]
+
+
+class K8sResourceKind(Enum):
+    REQUESTS = "requests"
+    LIMITS = "limits"

src/jobs/util.py

@@ -7,10 +7,12 @@
 import threading
 import time
 from io import TextIOBase
-from typing import IO, AnyStr, Iterable, Mapping, TextIO
+from typing import IO, Any, AnyStr, Iterable, Mapping, TextIO, TypeVar, cast
 
 from jobs.types import AnyPath
 
+T = TypeVar("T", bound=Mapping[str, Any])
+
 
 def to_rational(s: str) -> float:
     """Convert a number with optional SI/binary unit to floating-point"""
@@ -42,9 +44,10 @@ def to_rational(s: str) -> float:
     return factor * magnitude
 
 
-def remove_none_values(d: dict) -> dict:
+def remove_none_values(d: T) -> T:
     """Remove all keys with a ``None`` value from a dict."""
-    return {k: v for k, v in d.items() if v is not None}
+    filtered_dict = {k: v for k, v in d.items() if v is not None}
+    return cast(T, filtered_dict)
 
 
 def sanitize_rfc1123_domain_name(s: str) -> str: