rawfile-localpv/orchestrator/k8s.py

71 lines
1.8 KiB
Python
Raw Normal View History

2020-04-24 11:48:09 +00:00
import json
import uuid
from pathlib import Path
from time import sleep
import pykube
import yaml
from munch import Munch
from consts import CONFIG
2020-04-24 11:48:09 +00:00
api = pykube.HTTPClient(pykube.KubeConfig.from_env())
def volume_to_node(volume_id):
pv = pykube.PersistentVolume.objects(api).get_by_name(name=volume_id)
pv = Munch.fromDict(pv.obj)
node_name = pv.spec.nodeAffinity.required.nodeSelectorTerms[0].matchExpressions[0][
"values"
][0]
2020-05-29 14:43:47 +00:00
expected_node_affinity = yaml.safe_load(
2020-04-24 11:48:09 +00:00
f"""
required:
nodeSelectorTerms:
- matchExpressions:
- key: hostname
operator: In
values:
- {node_name}
"""
)
2020-05-29 14:43:47 +00:00
assert pv.spec.nodeAffinity == expected_node_affinity
2020-04-24 11:48:09 +00:00
return node_name
def wait_for(pred, desc=""):
print(f"Waiting for {desc}", end="", flush=True)
while not pred():
print(".", end="", flush=True)
sleep(0.5)
print(" done")
def run_on_node(fn, node):
name = f"task-{uuid.uuid4()}"
ctx = {
"name": name,
"namespace": "kube-system", # FIXME
"nodeSelector": json.dumps({"kubernetes.io/hostname": node}),
"cmd": json.dumps(fn),
"image_repository": CONFIG["image_repository"],
"image_tag": CONFIG["image_tag"],
2020-04-24 11:48:09 +00:00
}
template = Path("./templates/task.yaml").read_bytes().decode()
manifest = template.format(**ctx)
obj = yaml.safe_load(manifest)
task_pod = pykube.Pod(api, obj)
task_pod.create()
def is_finished():
task_pod.reload()
status = task_pod.obj["status"]
if status["phase"] in ["Succeeded", "Failed"]:
return True
return False
wait_for(is_finished, "task to finish")
if task_pod.obj["status"]["phase"] != "Succeeded":
raise Exception(f"Task {name} failed")
task_pod.delete()