add refs

Author: alon-dotan-starkware

.github/workflows/main.yml

@@ -8,6 +8,13 @@ jobs:
   build-image:
     runs-on: ubuntu-latest
     steps:
+    - name: Login to registry ${{ env.REGISTRY }}
+        uses: docker/[email protected]
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+          logout: true
     - uses: actions/checkout@v4
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3

deployment/rbac.yml

@@ -0,0 +1,24 @@
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRole
+metadata:
+  name: simple-python-scheduler
+rules:
+  - apiGroups: ["*"]
+    resources:
+      - "pods"
+      - "nodes"
+      - "bindings"
+    verbs: ["*"]
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRoleBinding
+metadata:
+  name: simple-python-scheduler
+subjects:
+  - kind: ServiceAccount
+    name: default
+    namespace: simple-python-scheduler
+roleRef:
+  kind: ClusterRole
+  name: simple-python-scheduler
+  apiGroup: rbac.authorization.k8s.io

src/main.py

@@ -6,10 +6,11 @@
 
 from kubernetes import client, config, watch
 
-config.load_kube_config()
+config.load_incluster_config()
 v1 = client.CoreV1Api()
 
-scheduler_name = "foobar"
+scheduler_name = "simple-python-scheduler"
+NAMESPACE = "test-ns"
 
 
 def nodes_available():
@@ -21,33 +22,56 @@ def nodes_available():
     return ready_nodes
 
 
-def scheduler(name, node, namespace="default"):
-    body = client.V1Binding()
+def get_running_pods():
+    pods = []
+    for pod in v1.list_namespaced_pod(namespace=NAMESPACE).items:
+        if pod.status.phase == "Running":
+            pods.append(pod)
+    return pods
 
-    target = client.V1ObjectReference()
-    target.kind = "Node"
-    target.apiVersion = "v1"
-    target.name = node
 
-    meta = client.V1ObjectMeta()
-    meta.name = name
+def preemption(priority):
+    for pod in get_running_pods():
+        try:
+            pod_priority = pod.metadata.annotations["priority"]
+        except:
+            print(f"Preempting pod {pod.metadata.name}")
+            return v1.delete_namespaced_pod(name=pod.metadata.name, namespace=NAMESPACE)
 
-    body.target = target
-    body.metadata = meta
+        if int(pod_priority) > int(priority):
+            print(f"Preempting pod {pod.metadata.name}")
+            return v1.delete_namespaced_pod(name=pod.metadata.name, namespace=NAMESPACE)
+    return None
 
-    return v1.create_namespaced_binding_binding(name, namespace, body)
+
+def scheduler(name, priority, node, namespace=NAMESPACE):
+    while len(v1.list_node().items) <= len(get_running_pods()):
+        preemption(priority)
+
+    target = client.V1ObjectReference(kind="Node", api_version="v1", name=node)
+    meta = client.V1ObjectMeta(name=name)
+    body = client.V1Binding(target=target, metadata=meta)
+    return v1.create_namespaced_binding(
+        namespace=namespace, body=body, _preload_content=False
+    )
 
 
 def main():
     w = watch.Watch()
-    for event in w.stream(v1.list_namespaced_pod, "default"):
+    for event in w.stream(v1.list_namespaced_pod, NAMESPACE):
         if (
             event["object"].status.phase == "Pending"
             and event["object"].spec.scheduler_name == scheduler_name
         ):
             try:
+                try:
+                    priority = event["object"].metadata.annotations["priority"]
+                except TypeError:
+                    priority = 1000
                 res = scheduler(
-                    event["object"].metadata.name, random.choice(nodes_available())
+                    name=event["object"].metadata.name,
+                    priority=priority,
+                    node=random.choice(nodes_available()),
                 )
             except client.rest.ApiException as e:
                 print(json.loads(e.body)["message"])

src/resources/example-job.yaml

@@ -0,0 +1,19 @@
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: pi
+  namespace: test-ns
+spec:
+  template:
+    metadata:
+      annotations:
+        priority: "200"
+    spec:
+      schedulerName: simple-python-scheduler
+      containers:
+      - name: pi
+        image: perl:5.34.0
+        command: ["perl",  "-Mbignum=bpi", "-wle", "print bpi(2000)"]
+      restartPolicy: Never
+  parallelism: 4
+  backoffLimit: 4

src/resources/example-pod.yaml

@@ -0,0 +1,14 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: foobar-sched-2
+  namespace: test-ns
+  annotations:
+    priority: "2"
+spec:
+  schedulerName: simple-python-scheduler
+  containers:
+  - name: nginx
+    image: nginx:latest
+    ports:
+    - containerPort: 80