builder.py
  1  # Python Imports
  2  from typing import List, Literal, Optional, Self
  3  
  4  from kubernetes.client import (
  5      V1ConfigMapVolumeSource,
  6      V1ContainerPort,
  7      V1ObjectMeta,
  8      V1Pod,
  9      V1PodDNSConfig,
 10      V1PolicyRule,
 11      V1ResourceRequirements,
 12      V1Role,
 13      V1RoleBinding,
 14      V1RoleRef,
 15      V1Volume,
 16      V1VolumeMount,
 17  )
 18  
 19  # Project Imports
 20  from src.deployments.core.builders import PodBuilder
 21  from src.deployments.core.configs.command import Command, CommandConfig
 22  from src.deployments.core.configs.container import ContainerConfig, Image
 23  from src.deployments.core.configs.helpers import find_container_config
 24  from src.deployments.core.configs.pod import PodConfig, PodSpecConfig, PodTemplateSpecConfig
 25  
 26  ScriptMode = Literal["server", "batch"]
 27  
 28  
 29  NAME = "pod-api-requester-container"
 30  
 31  
 32  class PodApiRequesterBuilder(PodBuilder):
 33      _mode: Optional[ScriptMode] = None
 34      """Sent to the --mode arg of the api-requester script.
 35      This must be set through `with_mode` prior to building."""
 36  
 37      def build_role(self) -> V1Role:
 38          return V1Role(
 39              api_version="rbac.authorization.k8s.io/v1",
 40              kind="Role",
 41              metadata=V1ObjectMeta(
 42                  name="pod-service-reader",
 43                  namespace=self.config.namespace,
 44              ),
 45              rules=[
 46                  V1PolicyRule(
 47                      api_groups=[""], resources=["pods", "services"], verbs=["get", "list", "watch"]
 48                  )
 49              ],
 50          )
 51  
 52      def build_rolebinding(self) -> V1RoleBinding:
 53          return V1RoleBinding(
 54              api_version="rbac.authorization.k8s.io/v1",
 55              kind="RoleBinding",
 56              metadata=V1ObjectMeta(
 57                  name="pod-service-reader-binding", namespace=self.config.namespace
 58              ),
 59              role_ref=V1RoleRef(
 60                  kind="Role",
 61                  name="pod-service-reader",
 62                  apiGroup="rbac.authorization.k8s.io",
 63              ),
 64              subjects=[
 65                  {
 66                      "kind": "ServiceAccount",
 67                      "name": "default",
 68                      "namespace": self.config.namespace,
 69                  },
 70              ],
 71          )
 72  
 73      def build(self) -> V1Pod:
 74          if not self.config.name:
 75              raise ValueError(f"Must configure node first. Config: `{self.config}`")
 76          if not self._mode:
 77              raise ValueError(
 78                  f"Script mode must be set using `with_mode` before building. Config: `{self.config}`"
 79              )
 80          return super().build()
 81  
 82      def with_namespace(self, namespace: str) -> Self:
 83          self.config = create_pod_config(namespace=namespace)
 84          return self
 85  
 86      def with_mode(self, mode: ScriptMode) -> Self:
 87          container = find_container_config(self.config.pod_spec_config, NAME)
 88          apply_command_config(container.command_config, mode=mode)
 89          self._mode = mode
 90          return self
 91  
 92      def with_command(self, command: str, args: List[str]) -> Self:
 93          container_config = find_container_config(self.config.pod_spec_config, NAME)
 94          container_config.command_config = CommandConfig(
 95              commands=[Command(command=command, args=args)]
 96          )
 97          return self
 98  
 99      def with_debug(self) -> Self:
100          self._mode = "debug"
101          return self.with_command("sleep", args=["infinity"])
102  
103  
104  def apply_command_config(config: CommandConfig, mode: ScriptMode = "server"):
105      try:
106          command = config.find_command("python")
107      except IndexError as e:
108          raise ValueError(f"pod-api-requester command not found. CommandConfig: `{config}`") from e
109  
110      command.add_args(
111          [
112              ("--mode", mode),
113              ("--config", "/mount/config.yaml"),
114          ],
115          on_duplicate="replace",
116      )
117  
118  
119  def create_container_config() -> ContainerConfig:
120      config = ContainerConfig(
121          name=NAME,
122          image=Image(
123              repo="pearsonwhite/pod-api-requester", tag="4575c70fd1efddabb7673ebe8a1f2b482473e0db"
124          ),
125          image_pull_policy="Always",
126      )
127      config.ports = [
128          V1ContainerPort(8645),
129          V1ContainerPort(8008),
130          V1ContainerPort(8080),
131      ]
132  
133      config.with_volume_mount(V1VolumeMount(name="api-requester-config-volume", mount_path="/mount"))
134      config.command_config = CommandConfig(
135          commands=[Command(command="python", args=["/app/api_requester.py"])]
136      )
137  
138      return config
139  
140  
141  def create_pod_spec_config(namespace: str) -> PodSpecConfig:
142      config = PodSpecConfig(
143          container_configs=[create_container_config()],
144          dns_config=V1PodDNSConfig(
145              searches=[f"zerotesting-publisher.{namespace}.svc.cluster.local"]
146          ),
147          namespace=namespace,
148      )
149      config.with_volume(
150          V1Volume(
151              name="api-requester-config-volume",
152              config_map=V1ConfigMapVolumeSource(name="api-requester-config"),
153          )
154      )
155      return config
156  
157  
158  def create_pod_template_spec_config(namespace: str) -> PodTemplateSpecConfig:
159      config = PodTemplateSpecConfig(
160          pod_spec_config=create_pod_spec_config(namespace),
161          namespace=namespace,
162          name="publisher",
163      )
164      config.with_app("zerotenkay-publisher")
165      return config
166  
167  
168  def create_pod_config(namespace: str) -> PodConfig:
169      config = PodConfig(
170          name="publisher",
171          namespace=namespace,
172          pod_spec_config=create_pod_spec_config(namespace=namespace),
173      )
174      config.with_app("zerotenkay-publisher")
175      return config
176  
177  
178  def create_resources() -> V1ResourceRequirements:
179      return V1ResourceRequirements(
180          requests={"memory": "64Mi", "cpu": "150m"}, limits={"memory": "600Mi", "cpu": "400m"}
181      )