kugl 0.3.0__py3-none-any.whl

kugl/api.py

@@ -0,0 +1,28 @@
+"""
+Imports usable by user-defined tables in Python (once we have those.)
+"""
+
+from kugl.impl.registry import Registry
+
+from kugl.util import (
+    fail,
+    parse_age,
+    parse_utc,
+    to_age,
+    to_utc,
+)
+
+
+def schema(name: str):
+    def wrap(cls):
+        Registry.get().add_schema(name, cls)
+        return cls
+    return wrap
+
+
+def table(**kwargs):
+    def wrap(cls):
+        Registry.get().add_table(cls, **kwargs)
+        return cls
+    return wrap
+

kugl/builtins/helpers.py

@@ -0,0 +1,156 @@
+"""
+Wrappers to make JSON returned by kubectl easier to work with.
+"""
+
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+import funcy as fn
+
+from kugl.util import parse_size, parse_cpu
+
+# What container name is considered the "main" container, if present
+MAIN_CONTAINERS = ["main", "notebook", "app"]
+
+
+@dataclass
+class Limits:
+    """
+    A class to hold CPU, GPU and memory resources. This is called "Limits" although it's used for both requests
+    and limits, so as not to confuse "resources" with Kubernetes resources in general.
+    """
+    cpu: Optional[float]
+    gpu: Optional[float]
+    mem: Optional[int]
+
+    def __add__(self, other):
+        if self.cpu is None and other.cpu is None:
+            cpu = None
+        else:
+            cpu = (self.cpu or 0) + (other.cpu or 0)
+        if self.gpu is None and other.gpu is None:
+            gpu = None
+        else:
+            gpu = (self.gpu or 0) + (other.gpu or 0)
+        if self.mem is None and other.mem is None:
+            mem = None
+        else:
+            mem = (self.mem or 0) + (other.mem or 0)
+        return Limits(cpu, gpu, mem)
+
+    def __radd__(self, other):
+        """Needed to support sum() -- handles 0 as a starting value"""
+        return self if other == 0 else self.__add__(other)
+
+    def as_tuple(self):
+        return (self.cpu, self.gpu, self.mem)
+
+    @classmethod
+    def extract(cls, obj):
+        """Extract a Limits object from a dictionary, or return an empty one if the dictionary is None.
+
+        :param obj: A dictionary with keys "cpu", "nvidia.com/gpu" and "memory" """
+        if obj is None:
+            return Limits(None, None, None)
+        cpu = parse_cpu(obj.get("cpu"))
+        gpu = parse_cpu(obj.get("nvidia.com/gpu"))
+        mem = parse_size(obj.get("memory"))
+        return Limits(cpu, gpu, mem)
+
+
+class ItemHelper:
+    """Some common code for wrappers on JSON for pods, nodes et cetera."""
+
+    def __init__(self, obj):
+        self.obj = obj
+        self.metadata = self.obj.get("metadata", {})
+        self.labels = self.metadata.get("labels", {})
+
+    def __getitem__(self, key):
+        """Return a key from the object; no default, will error if not present"""
+        return self.obj[key]
+
+    @property
+    def name(self):
+        """Return the name of the object from the metadata, or none if unavailable."""
+        return self.metadata.get("name") or self.obj.get("name")
+
+    @property
+    def namespace(self):
+        """Return the name of the object from the metadata, or none if unavailable."""
+        return self.metadata.get("namespace")
+
+    def label(self, name):
+        """Return one of the labels from the object, or None if it doesn't have that label."""
+        return self.labels.get(name)
+
+
+class Containerized:
+
+    @abstractmethod
+    def containers(self):
+        raise NotImplementedError()
+
+    def resources(self, tag):
+        return sum(Limits.extract(c.get("resources", {}).get(tag)) for c in self.containers)
+
+
+class PodHelper(ItemHelper, Containerized):
+
+    @property
+    def command(self):
+        return " ".join((self.main or {}).get("command", []))
+
+    @property
+    def is_daemon(self):
+        return any(ref.get("kind") == "DaemonSet" for ref in self.metadata.get("ownerReferences", []))
+
+    @property
+    def containers(self):
+        """Return the containers in the pod, if any, else an empty list."""
+        return self["spec"].get("containers", [])
+
+    @property
+    def main(self):
+        """Return the main container in the pod, if any, defined as the first container with a name
+        in MAIN_CONTAINERS.  If there are none of those, return the first one.
+        """
+        if not self.containers:
+            return None
+        main = fn.first(fn.filter(lambda c: c["name"] in MAIN_CONTAINERS, self.containers))
+        return main or self.containers[0]
+
+
+class JobHelper(ItemHelper, Containerized):
+
+    @property
+    def status(self):
+        status = self.obj.get("status", {})
+        if len(status) == 0:
+            return "Unknown"
+        # Per
+        # https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1JobStatus.md
+        # and https://kubernetes.io/docs/concepts/workloads/controllers/job/
+        for c in status.get("conditions", []):
+            if c["status"] == "True":
+                if c["type"] == "Failed":
+                    # TODO use a separate column
+                    return c.get("reason") or "Failed"
+                if c["type"] == "Suspended":
+                    return "Suspended"
+                if c["type"] == "Complete":
+                    return "Complete"
+            if c["type"] == "FailureTarget":
+                return "Failed"
+            if c["type"] == "SuccessCriteriaMet":
+                return "Complete"
+        if status.get("active", 0) > 0:
+            return "Running"
+        return "Unknown"
+
+    @property
+    def containers(self):
+        """Return the containers in the job, if any, else an empty list."""
+        return self["spec"]["template"]["spec"].get("containers", [])
+

kugl/builtins/kubernetes.py

@@ -0,0 +1,220 @@
+"""
+Built-in table definitions for Kubernetes.
+
+NOTE: This is not a good example of how to write user-defined tables.
+FIXME: Remove references to non-API imports.
+FIXME: Don't use ArgumentParser in the API.
+"""
+import json
+from argparse import ArgumentParser
+from threading import Thread
+
+from .helpers import Limits, ItemHelper, PodHelper, JobHelper
+from kugl.api import schema, table, fail
+from kugl.util import parse_utc, run, WHITESPACE
+
+
+@schema("kubernetes")
+class KubernetesData:  # FIXME: this should be a resource type, not a schema
+
+    def add_cli_options(self, ap: ArgumentParser):
+        ap.add_argument("-a", "--all-namespaces", default=False, action="store_true")
+        ap.add_argument("-n", "--namespace", type=str)
+
+    def handle_cli_options(self, args):
+        if args.all_namespaces and args.namespace:
+            fail("Cannot use both -a/--all-namespaces and -n/--namespace")
+        self.set_namespace(args.all_namespaces, args.namespace)
+
+    def set_namespace(self, all_namespaces: bool, namespace: str):
+        if all_namespaces:
+            # FIXME: engine.py and testing.py still use this
+            self.ns = "__all"
+            self.all_ns = True
+        else:
+            self.ns = namespace or "default"
+            self.all_ns = False
+
+    def get_objects(self, kind: str, namespaced: bool)-> dict:
+        """Fetch resources from Kubernetes using kubectl.
+
+        :param kind: Kubernetes resource type e.g. "pods"
+        :return: JSON as output by "kubectl get {kind} -o json"
+        """
+        namespace_flag = ["--all-namespaces"] if self.ns else ["-n", self.ns]
+        if kind == "pods":
+            pod_statuses = {}
+            # Kick off a thread to get pod statuses
+            def _fetch():
+                _, output, _ = run(["kubectl", "get", "pods", *namespace_flag])
+                pod_statuses.update(self._pod_status_from_pod_list(output))
+            status_thread = Thread(target=_fetch, daemon=True)
+            status_thread.start()
+        if namespaced:
+            _, output, _= run(["kubectl", "get", kind, *namespace_flag, "-o", "json"])
+        else:
+            _, output, _ = run(["kubectl", "get", kind, "-o", "json"])
+        data = json.loads(output)
+        if kind == "pods":
+            # Add pod status to pods
+            status_thread.join()
+            def pod_with_updated_status(pod):
+                metadata = pod["metadata"]
+                status = pod_statuses.get(f"{metadata['namespace']}/{metadata['name']}")
+                if status:
+                    pod["kubectl_status"] = status
+                    return pod
+                return None
+            data["items"] = list(filter(None, map(pod_with_updated_status, data["items"])))
+        return data
+
+    def _pod_status_from_pod_list(self, output) -> dict[str, str]:
+        """
+        Convert the tabular output of 'kubectl get pods' to JSON.
+        :return: a dict mapping "namespace/name" to status
+        """
+        rows = [WHITESPACE.split(line.strip()) for line in output.strip().split("\n")]
+        if len(rows) < 2:
+            return {}
+        header, rows = rows[0], rows[1:]
+        name_index = header.index("NAME")
+        status_index = header.index("STATUS")
+        # It would be nice if 'kubectl get pods' printed the UID, but it doesn't, so use
+        # "namespace/name" as the key.  (Can't use a tuple since this has to be JSON-dumped.)
+        if self.all_ns:
+            namespace_index = header.index("NAMESPACE")
+            return {f"{row[namespace_index]}/{row[name_index]}": row[status_index] for row in rows}
+        else:
+            return {f"{self.ns}/{row[name_index]}": row[status_index] for row in rows}
+
+
+@table(schema="kubernetes", name="nodes", resource="nodes")
+class NodesTable:
+
+    @property
+    def schema(self):
+        return """
+            name TEXT,
+            uid TEXT,
+            cpu_alloc REAL,
+            gpu_alloc REAL,
+            mem_alloc INTEGER,
+            cpu_cap REAL,
+            gpu_cap REAL,
+            mem_cap INTEGER
+        """
+
+    def make_rows(self, context) -> list[tuple[dict, tuple]]:
+        for item in context.data["items"]:
+            node = ItemHelper(item)
+            yield item, (
+                node.name,
+                node.metadata.get("uid"),
+                *Limits.extract(node["status"]["allocatable"]).as_tuple(),
+                *Limits.extract(node["status"]["capacity"]).as_tuple(),
+            )
+
+
+@table(schema="kubernetes", name="pods", resource="pods")
+class PodsTable:
+
+    @property
+    def schema(self):
+        return """
+            name TEXT,
+            uid TEXT,
+            is_daemon INTEGER,
+            namespace TEXT,
+            node_name TEXT,
+            creation_ts INTEGER,
+            command TEXT,
+            phase TEXT,
+            status TEXT,
+            cpu_req REAL,
+            gpu_req REAL,
+            mem_req INTEGER,
+            cpu_lim REAL,
+            gpu_lim REAL,
+            mem_lim INTEGER
+        """
+
+    def make_rows(self, context) -> list[tuple[dict, tuple]]:
+        for item in context.data["items"]:
+            pod = PodHelper(item)
+            yield item, (
+                pod.name,
+                pod.metadata.get("uid"),
+                1 if pod.is_daemon else 0,
+                pod.namespace,
+                pod["spec"].get("nodeName"),
+                parse_utc(pod.metadata["creationTimestamp"]),
+                pod.command,
+                pod["status"]["phase"],
+                pod["kubectl_status"],
+                *pod.resources("requests").as_tuple(),
+                *pod.resources("limits").as_tuple(),
+            )
+
+
+@table(schema="kubernetes", name="jobs", resource="jobs")
+class JobsTable:
+
+    @property
+    def schema(self):
+        return """
+            name TEXT,
+            uid TEXT,
+            namespace TEXT,
+            status TEXT,
+            cpu_req REAL,
+            gpu_req REAL,
+            mem_req INTEGER,
+            cpu_lim REAL,
+            gpu_lim REAL,
+            mem_lim INTEGER
+        """
+
+    def make_rows(self, context) -> list[tuple[dict, tuple]]:
+        for item in context.data["items"]:
+            job = JobHelper(item)
+            yield item, (
+                job.name,
+                job.metadata.get("uid"),
+                job.namespace,
+                job.status,
+                *job.resources("requests").as_tuple(),
+                *job.resources("limits").as_tuple(),
+            )
+
+
+class LabelsTable:
+    """Base class for all built-in label tables; subclasses need only define UID_FIELD."""
+
+    @property
+    def schema(self):
+        return f"""
+            {self.UID_FIELD} TEXT,
+            key TEXT,
+            value TEXT
+        """
+
+    def make_rows(self, context) -> list[tuple[dict, tuple]]:
+        for item in context.data["items"]:
+            thing = ItemHelper(item)
+            for key, value in thing.labels.items():
+                yield item, (thing.metadata.get("uid"), key, value)
+
+
+@table(schema="kubernetes", name="node_labels", resource="nodes")
+class NodeLabelsTable(LabelsTable):
+    UID_FIELD = "node_uid"
+
+
+@table(schema="kubernetes", name="pod_labels", resource="pods")
+class PodLabelsTable(LabelsTable):
+    UID_FIELD = "pod_uid"
+
+
+@table(schema="kubernetes", name="job_labels", resource="jobs")
+class JobLabelsTable(LabelsTable):
+    UID_FIELD = "job_uid"

kugl/builtins/kubernetes.yaml

@@ -0,0 +1,27 @@
+
+resources:
+  - name: pods
+    namespaced: true
+  - name: pod_statuses
+    namespaced: true
+  - name: jobs
+    namespaced: true
+  - name: nodes
+    namespaced: false
+
+# node_taints builtin is defined here because it doesn't have any special column extraction
+# logic, and because it serves as a good unit test.
+
+create:
+  - table: node_taints
+    resource: nodes
+    row_source:
+      - items
+      - spec.taints
+    columns:
+      - name: node_uid
+        path: ^metadata.uid
+      - name: key
+        path: key
+      - name: effect
+        path: effect

kugl/impl/config.py

@@ -0,0 +1,207 @@
+"""
+Pydantic models for configuration files.
+"""
+import json
+import re
+from typing import Literal, Optional, Tuple, Callable, Union
+
+import jmespath
+from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic.functional_validators import model_validator
+
+from kugl.util import Age, parse_utc, parse_size, KPath, ConfigPath, parse_age, parse_cpu, fail
+
+PARENTED_PATH = re.compile(r"^(\^*)(.*)")
+
+
+class Settings(BaseModel):
+    """Holds the settings: entry from a user config file."""
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+    cache_timeout: Age = Age(120)
+    reckless: bool = False
+
+
+class UserInit(BaseModel):
+    """The root model for init.yaml; holds the entire file content."""
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+    settings: Optional[Settings] = Settings()
+    shortcuts: dict[str, list[str]] = {}
+
+
+class ColumnDef(BaseModel):
+    """Holds one entry from a columns: list in a user config file."""
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+    name: str
+    type: Literal["text", "integer", "real", "date", "age", "size", "cpu"] = "text"
+    path: Optional[str] = None
+    label: Optional[Union[str, list[str]]] = None
+    # Function to extract a column value from an object.
+    _extract: Callable[[object], object]
+    # Function to convert the extracted value to the SQL type
+    _convert: type
+    # Parsed value of self.path
+    _finder: jmespath.parser.Parser
+    # Number of ^ in self.path
+    _parents: int
+    # SQL type for this column
+    _sqltype: str
+
+    @model_validator(mode="after")
+    @classmethod
+    def gen_extractor(cls, config: 'ColumnDef') -> 'ColumnDef':
+        """
+        Generate the extract function for a column definition; given an object, it will
+        return a column value of the appropriate type.
+        """
+        if config.path and config.label:
+            raise ValueError("cannot specify both path and label")
+        elif config.path:
+            m = PARENTED_PATH.match(config.path)
+            config._parents = len(m.group(1))
+            try:
+                config._finder = jmespath.compile(m.group(2))
+            except jmespath.exceptions.ParseError as e:
+                raise ValueError(f"invalid JMESPath expression {m.group(2)} in column {config.name}") from e
+            config._extract = config._extract_jmespath
+        elif config.label:
+            if not isinstance(config.label, list):
+                config.label = [config.label]
+            config._extract = config._extract_label
+        else:
+            raise ValueError("must specify either path or label")
+        config._sqltype = KUGL_TYPE_TO_SQL_TYPE[config.type]
+        config._convert = KUGL_TYPE_CONVERTERS[config.type]
+        return config
+
+    def extract(self, obj: object, context) -> object:
+        """Extract the column value from an object and convert to the correct type."""
+        if obj is None:
+            if context.debug:
+                print(f"No object provided to extractor {self}")
+            return None
+        if context.debug:
+            print(f"Extract {self} from {self._abbreviate(obj)}")
+        value = self._extract(obj, context)
+        if context.debug:
+            print(f"Extracted {value}")
+        return None if value is None else self._convert(value)
+
+    def _extract_jmespath(self, obj: object, context) -> object:
+        """Extract a value from an object using a JMESPath finder."""
+        if self._parents > 0:
+            obj = context.get_parent(obj, self._parents)
+        if obj is None:
+            fail(f"Missing parent or too many ^ while evaluating {self.path}")
+        return self._finder.search(obj)
+
+    def _extract_label(self, obj: object, context) -> object:
+        """Extract a value from an object using a label."""
+        obj = context.get_root(obj)
+        if available := obj.get("metadata", {}).get("labels", {}):
+            for label in self.label:
+                if (value := available.get(label)) is not None:
+                    return value
+
+    def __str__(self):
+        if self.path:
+            return f"{self.name} path={self.path}"
+        return f"{self.name} label={','.join(self.label)}"
+
+    def _abbreviate(self, obj):
+        text = json.dumps(obj)
+        if len(text) > 100:
+            return text[:100] + "..."
+        return text
+
+
+KUGL_TYPE_CONVERTERS = {
+    "integer": int,
+    "real" : float,
+    "text": str,
+    "date": parse_utc,
+    "age": parse_age,
+    "size": parse_size,
+    "cpu": parse_cpu,
+}
+
+KUGL_TYPE_TO_SQL_TYPE = {
+    "integer": "integer",
+    "real": "real",
+    "text": "text",
+    "date": "integer",
+    "age": "integer",
+    "size": "integer",
+    "cpu": "real",
+}
+
+
+class ExtendTable(BaseModel):
+    """Holds the extend: section from a user config file."""
+    model_config = ConfigDict(extra="forbid")
+    table: str
+    columns: list[ColumnDef] = []
+
+
+class ResourceDef(BaseModel):
+    """Holds one entry from the resources: list in a user config file."""
+    name: str
+    # FIXME: Don't conflate all resource attributes in one class
+    namespaced: bool = True
+    cacheable: bool = True
+    file: Optional[str] = None
+    exec: Optional[Union[str, list[str]]] = None
+
+    @model_validator(mode="after")
+    @classmethod
+    def validate(cls, config: 'ResourceDef') -> 'ResourceDef':
+        if config.file and config.exec:
+            raise ValueError("Resource cannot specify both file and exec")
+        if config.file:
+            config.cacheable = False
+        return config
+
+    def __hash__(self):
+        return hash(self.name)
+
+    def __eq__(self, other):
+        return self.name == other.name
+
+
+class CreateTable(ExtendTable):
+    """Holds the create: section from a user config file."""
+    resource: str
+    row_source: Optional[list[str]] = None
+
+
+class UserConfig(BaseModel):
+    """The root model for a user config file; holds the complete file content."""
+    model_config = ConfigDict(extra="forbid")
+    resources: list[ResourceDef] = []
+    extend: list[ExtendTable] = []
+    create: list[CreateTable] = []
+
+
+# FIXME use typevars
+def parse_model(model_class, root: dict) -> Tuple[object, list[str]]:
+    """Parse a dict into a model instance (typically a UserConfig).
+
+    :return: A tuple of (parsed object, list of errors).  On success, the error list is None.
+        On failure, the parsed object is None.
+    """
+    try:
+        return model_class.model_validate(root), None
+    except ValidationError as e:
+        error_location = lambda err: '.'.join(str(x) for x in err['loc'])
+        return None, [f"{error_location(err)}: {err['msg']}" for err in e.errors()]
+
+# FIXME use typevars
+def parse_file(model_class, path: ConfigPath) -> Tuple[object, list[str]]:
+    """Parse a configuration file into a model instance, handling edge cases.
+
+    :return: Same as parse_model."""
+    if not path.exists():
+        return model_class(), None
+    if path.is_world_writeable():
+        return None, [f"{path} is world writeable, refusing to run"]
+    return parse_model(model_class, path.parse_yaml() or {})
+