konduktor-nightly 0.1.0.dev20250209104336__py3-none-any.whl → 0.1.0.dev20250313070642__py3-none-any.whl

konduktor/logging.py

@@ -1,12 +1,22 @@
 """Logging utilities."""
 
+import contextlib
 import logging
 import os
+import threading
+from datetime import datetime
 
 import colorama
 
-_FORMAT = "[%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s"
-_DATE_FORMAT = "%m-%d %H:%M:%S"
+from konduktor import constants
+
+CHECK_MARK_EMOJI = '\U00002714'  # Heavy check mark unicode
+PARTY_POPPER_EMOJI = '\U0001f389'  # Party popper unicode
+
+_FORMAT = '[%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
+_DATE_FORMAT = '%m-%d %H:%M:%S'
+
+_logging_config = threading.local()
 
 
 class NewLineFormatter(logging.Formatter):
@@ -18,9 +28,9 @@ class NewLineFormatter(logging.Formatter):
 
     def format(self, record):
         msg = logging.Formatter.format(self, record)
-        if record.message != "":
+        if record.message != '':
             parts = msg.partition(record.message)
-            msg = msg.replace("\n", "\r\n" + parts[0])
+            msg = msg.replace('\n', '\r\n' + parts[0])
             if self.dim:
                 msg = colorama.Style.DIM + msg + colorama.Style.RESET_ALL
         return msg
@@ -29,10 +39,21 @@ class NewLineFormatter(logging.Formatter):
 FORMATTER = NewLineFormatter(_FORMAT, datefmt=_DATE_FORMAT)
 
 
+@contextlib.contextmanager
+def set_logging_level(logger: str, level: int):
+    logger = logging.getLogger(logger)
+    original_level = logger.level
+    logger.setLevel(level)
+    try:
+        yield
+    finally:
+        logger.setLevel(original_level)
+
+
 def get_logger(name: str):
     # Determine the logging level based on the KONDUKTOR_DEBUG environment variable
     log_level = logging.INFO
-    if os.environ.get("KONDUKTOR_DEBUG", None) == "1":
+    if os.environ.get('KONDUKTOR_DEBUG', None) == '1':
         log_level = logging.DEBUG
 
     # Configure the logger
@@ -45,3 +66,26 @@ def get_logger(name: str):
         logger.addHandler(ch)
     logger.propagate = False
     return logger
+
+
+def is_silent():
+    if not hasattr(_logging_config, 'is_silent'):
+        # Should not set it globally, as the global assignment
+        # will be executed only once if the module is imported
+        # in the main thread, and will not be executed in other
+        # threads.
+        _logging_config.is_silent = False
+    return _logging_config.is_silent
+
+
+def get_run_timestamp() -> str:
+    return 'konduktor-' + datetime.now().strftime('%Y-%m-%d-%H-%M-%S-%f')
+
+
+def generate_tmp_logging_file_path(file_name: str) -> str:
+    """Generate an absolute path of a tmp file for logging."""
+    run_timestamp = get_run_timestamp()
+    log_dir = os.path.join(constants.KONDUKTOR_LOGS_DIRECTORY, run_timestamp)
+    log_path = os.path.expanduser(os.path.join(log_dir, file_name))
+
+    return log_path

konduktor/manifests/dmesg_daemonset.yaml

@@ -26,6 +26,14 @@ spec:
         operator: "Equal"
         value: "present"
         effect: "NoSchedule"
+      - key: "trainy.konduktor.ai/faulty"
+        operator: "Equal"
+        value: "true"
+        effect: "NoSchedule"
+      - key: "cloud.google.com/gke-queued"
+        operator: "Equal"
+        value: "true"
+        effect: "NoSchedule"
       containers:
       - name: dmesg
         image: ubuntu:22.04

konduktor/manifests/pod_cleanup_controller.yaml

@@ -0,0 +1,129 @@
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: pod-cleanup-controller
+  namespace: default
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+  name: pod-cleanup-controller
+  namespace: default
+rules:
+- apiGroups: [""]
+  resources: ["pods", "pods/status", "events"]
+  verbs: ["get", "list", "watch", "delete", "patch", "update"]
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: RoleBinding
+metadata:
+  name: pod-cleanup-controller
+  namespace: default
+subjects:
+- kind: ServiceAccount
+  name: pod-cleanup-controller
+  namespace: default
+roleRef:
+  kind: Role
+  name: pod-cleanup-controller
+  apiGroup: rbac.authorization.k8s.io
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: pod-cleanup-controller
+  namespace: default
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: pod-cleanup-controller
+  template:
+    metadata:
+      labels:
+        app: pod-cleanup-controller
+    spec:
+      serviceAccountName: pod-cleanup-controller
+      containers:
+      - name: controller
+        image: python:3.10
+        command: ["/bin/sh", "-c"]
+        args: ["pip install kubernetes && echo 'starting controller' && python /app/controller.py"]
+        env:
+          - name: PYTHONUNBUFFERED
+            value: "0"
+        volumeMounts:
+        - name: controller-code
+          mountPath: /app
+      volumes:
+      - name: controller-code
+        configMap:
+          name: pod-cleanup-controller-code
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: pod-cleanup-controller-code
+  namespace: default
+data:
+  controller.py: |
+    from kubernetes import client, config, watch
+    from collections import defaultdict
+    from datetime import datetime
+    import time
+
+    FAILURE_MODES = ['ErrImagePull', 'InvalidImageName']
+
+    def check_failure_mode(message):
+        for mode in FAILURE_MODES:
+            if mode in message:
+                return mode
+        return ''
+
+    def main():
+        # Load kube config
+        try:
+            config.load_incluster_config()
+        except:
+            config.load_kube_config()
+
+        v1 = client.CoreV1Api()
+        error_counts = defaultdict(int)
+
+        w = watch.Watch()
+        while True:
+          for event in w.stream(v1.list_namespaced_event, namespace="default"):
+              if event['object'].type == 'Warning' and event['object'].reason == 'Failed' and  check_failure_mode(event['object'].message):
+                  pod_name = event['object'].involved_object.name
+                  pod_namespace = event['object'].involved_object.namespace
+                  print(f"Pod {pod_namespace}/{pod_name} has failed with ErrImagePull. Patching and deleting...")
+                  try:
+                      # Get current time in UTC
+                      current_time = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")
+                      
+                      # Create the status patch
+                      body = {
+                          "status": {
+                              "conditions": [
+                                  {
+                                      "type": "ConfigIssue",
+                                      "status": "True",
+                                      "reason": "ErrImagePull",
+                                      "lastTransitionTime": current_time
+                                  }
+                              ]
+                          }
+                      }
+                      
+                      # Patch pod status
+                      v1.patch_namespaced_pod_status(pod_name, pod_namespace, body)
+                      
+                      # Delete the pod
+                      v1.delete_namespaced_pod(pod_name, pod_namespace)
+                  except Exception as e:
+                      print(f"Error handling pod: {e}")
+          print("Finished event stream... waiting for another stream...")
+          time.sleep(5)
+
+    if __name__ == '__main__':
+        main()