kubetorch 0.2.0__py3-none-any.whl

kubetorch/resources/volumes/volume.py

@@ -0,0 +1,365 @@
+import subprocess
+import textwrap
+import uuid
+
+from functools import cached_property
+from typing import Dict
+
+from kubernetes import client
+from kubernetes.client import ApiException, V1PersistentVolumeClaim
+
+from kubetorch.constants import DEFAULT_VOLUME_ACCESS_MODE, KT_MOUNT_FOLDER
+from kubetorch.globals import config
+from kubetorch.logger import get_logger
+from kubetorch.utils import load_kubeconfig
+
+logger = get_logger(__name__)
+
+
+class Volume:
+    """
+    Manages persistent storage for Kubetorch services and deployments.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        size: str,
+        storage_class: str = None,
+        mount_path: str = None,
+        access_mode: str = None,
+        namespace: str = None,
+        core_v1: client.CoreV1Api = None,
+    ):
+        """
+        Kubetorch Volume object, specifying persistent storage properties.
+
+        Args:
+            name (str): Name of the volume.
+            size (str): Size of the volume.
+            storage_class (str, optional): Storage class to use for the volume.
+            mount_path (str, optional): Mount path for the volume.
+            access_mode (str, optional): Access mode for the volume.
+            namespace (str, optional): Namespace for the volume.
+
+        Example:
+
+        .. code-block:: python
+
+            import kubetorch as kt
+
+            kt.Volume(name="my-data", size="5Gi"),  # Standard volume (ReadWriteOnce)
+
+            # Shared volume (ReadWriteMany, requires JuiceFS or similar)
+            kt.Volume(name="shared-data", size="10Gi", storage_class="juicefs-sc-shared", access_mode="ReadWriteMany")
+
+            # uv cache
+            compute = kt.Compute(
+                cpus=".01",
+                env_vars={
+                    "UV_CACHE_DIR": "/ktfs/kt-global-cache/uv_cache",
+                    "HF_HOME": "/ktfs/kt-global-cache/hf_cache",
+                },
+                volumes=[kt.Volume("kt-global-cache", size="10Gi")],
+            )
+
+        """
+        self._storage_class = storage_class
+        if core_v1 is None:
+            load_kubeconfig()
+
+        self.size = size
+        self.access_mode = access_mode or DEFAULT_VOLUME_ACCESS_MODE
+        self.mount_path = mount_path or f"/{KT_MOUNT_FOLDER}/{name}"
+
+        self.name = name
+        self.namespace = namespace
+        self.core_v1 = core_v1 or client.CoreV1Api()
+
+    @property
+    def pvc_name(self) -> str:
+        return self.name
+
+    @cached_property
+    def storage_class(self) -> str:
+        """Get storage class - either specified or cluster default"""
+        if self._storage_class:
+            return self._storage_class
+
+        try:
+            storage_v1 = client.StorageV1Api()
+            storage_classes = storage_v1.list_storage_class().items
+
+            # If RWX is requested, prefer RWX-capable classes
+            if self.access_mode == "ReadWriteMany":
+                for sc in storage_classes:
+                    provisioner = getattr(sc, "provisioner", "")
+                    if provisioner in {
+                        "csi.juicefs.com",
+                        "nfs.csi.k8s.io",
+                        "cephfs.csi.ceph.com",
+                    }:
+                        return sc.metadata.name
+                raise ValueError("No RWX-capable storage class found")
+
+            # Otherwise, pick the default StorageClass
+            for sc in storage_classes:
+                annotations = sc.metadata.annotations or {}
+                if (
+                    annotations.get("storageclass.kubernetes.io/is-default-class")
+                    == "true"
+                ):
+                    logger.info(f"Using default storage class: {sc.metadata.name}")
+                    return sc.metadata.name
+
+            # No default found, fall back to first available
+            available_classes = [sc.metadata.name for sc in storage_classes]
+            first_sc = available_classes[0]
+            if len(available_classes) == 1:
+                logger.info(
+                    f"No default storage class found, using only available one: {first_sc}"
+                )
+            else:
+                logger.warning(
+                    f"No default storage class found, using first available: {first_sc}. "
+                    f"Available: {available_classes}. Consider setting a default or specifying storage_class parameter."
+                )
+            return first_sc
+
+        except Exception as e:
+            logger.error(f"Failed to get storage classes: {e}")
+            raise
+
+    @classmethod
+    def from_name(
+        cls,
+        name: str,
+        create_if_missing: bool = False,
+        namespace: str = None,
+        core_v1: client.CoreV1Api = None,
+    ) -> "Volume":
+        """Get existing volume or optionally create it"""
+        if core_v1 is None:
+            load_kubeconfig()
+            core_v1 = client.CoreV1Api()
+
+        namespace = namespace or config.namespace
+        pvc_name = name
+
+        try:
+            pvc = core_v1.read_namespaced_persistent_volume_claim(pvc_name, namespace)
+
+            storage_class = pvc.spec.storage_class_name
+            size = pvc.spec.resources.requests.get("storage")
+            access_mode = (
+                pvc.spec.access_modes[0]
+                if pvc.spec.access_modes
+                else DEFAULT_VOLUME_ACCESS_MODE
+            )
+
+            annotations = pvc.metadata.annotations or {}
+            mount_path = annotations.get(
+                "kubetorch.com/mount-path", f"/{KT_MOUNT_FOLDER}/{name}"
+            )
+
+            # Create Volume with actual attributes from PVC
+            vol = cls(
+                name=name,
+                storage_class=storage_class,
+                mount_path=mount_path,
+                size=size,
+                access_mode=access_mode,
+                namespace=namespace,
+                core_v1=core_v1,
+            )
+
+            logger.debug(
+                f"Loaded existing PVC {pvc_name} with storage_class={storage_class}"
+            )
+            return vol
+
+        except ApiException as e:
+            if e.status == 404:
+                # PVC doesn't exist
+                if create_if_missing:
+                    vol = cls(name, namespace=namespace, core_v1=core_v1)
+                    vol.create()
+                    return vol
+                else:
+                    raise ValueError(
+                        f"Volume '{name}' (PVC: {pvc_name}) does not exist in namespace '{namespace}'"
+                    )
+            else:
+                # Some other API error
+                raise
+
+    def config(self) -> Dict[str, str]:
+        """Get configuration for this volume"""
+        return {
+            "name": self.name,
+            "size": self.size,
+            "access_mode": self.access_mode,
+            "mount_path": self.mount_path,
+            "storage_class": self.storage_class,
+            "namespace": self.namespace,
+        }
+
+    def pod_template_spec(self) -> dict:
+        """Convert to Kubernetes volume spec for pod template"""
+        return {
+            "name": self.name,
+            "persistentVolumeClaim": {"claimName": self.pvc_name},
+        }
+
+    def create(self) -> V1PersistentVolumeClaim:
+        """Create PVC if it doesn't exist"""
+        try:
+            try:
+                # Check if PVC already exists
+                existing_pvc = self.core_v1.read_namespaced_persistent_volume_claim(
+                    name=self.pvc_name, namespace=self.namespace
+                )
+                logger.debug(
+                    f"PVC {self.pvc_name} already exists in namespace {self.namespace}"
+                )
+                return existing_pvc
+            except ApiException as e:
+                if e.status != 404:
+                    # Some other error occurred
+                    raise
+
+            logger.info(f"Creating new PVC with name: {self.pvc_name}")
+
+            storage_class_name = self.storage_class
+
+            pvc_spec = client.V1PersistentVolumeClaimSpec(
+                access_modes=[self.access_mode],
+                resources=client.V1ResourceRequirements(
+                    requests={"storage": self.size}
+                ),
+                storage_class_name=storage_class_name,
+            )
+
+            pvc_metadata = client.V1ObjectMeta(
+                name=self.pvc_name,
+                labels={
+                    "app": "kubetorch",
+                    "kubetorch.com/volume": self.name,
+                },
+                annotations={"kubetorch.com/mount-path": self.mount_path},
+            )
+
+            pvc = client.V1PersistentVolumeClaim(
+                api_version="v1",
+                kind="PersistentVolumeClaim",
+                metadata=pvc_metadata,
+                spec=pvc_spec,
+            )
+
+            created_pvc = self.core_v1.create_namespaced_persistent_volume_claim(
+                namespace=self.namespace, body=pvc
+            )
+
+            logger.info(
+                f"Successfully created PVC {self.pvc_name} in namespace {self.namespace} with "
+                f"storage class {storage_class_name}"
+            )
+            return created_pvc
+
+        except Exception as e:
+            logger.error(f"Failed to create PVC {self.pvc_name}: {e}")
+            raise
+
+    def delete(self) -> None:
+        """Delete the PVC"""
+        try:
+            self.core_v1.delete_namespaced_persistent_volume_claim(
+                name=self.pvc_name, namespace=self.namespace
+            )
+            logger.debug(f"Successfully deleted PVC {self.pvc_name}")
+        except ApiException as e:
+            if e.status == 404:
+                logger.warning(f"PVC {self.pvc_name} not found")
+            else:
+                logger.error(f"Failed to delete PVC {self.pvc_name}: {e}")
+                raise
+
+    def exists(self) -> bool:
+        """Check if the PVC exists"""
+        try:
+            self.core_v1.read_namespaced_persistent_volume_claim(
+                name=self.pvc_name, namespace=self.namespace
+            )
+            return True
+        except ApiException as e:
+            if e.status == 404:
+                return False
+            else:
+                # Some other API error, re-raise
+                raise
+
+    def ssh(self, image: str = "alpine:latest"):
+        """
+        Launch an interactive debug shell with this volume mounted.
+
+        This method creates a temporary Kubernetes pod that mounts the
+        PersistentVolumeClaim (PVC) backing this Volume at the same path
+        (`self.mount_path`) used by Kubetorch services.
+
+        Args:
+            image (str, optional): Container image to use for the debug pod.
+                Must include a shell (e.g., `alpine:3.18`, `ubuntu:22.04`,
+                or a custom tools image). Defaults to `alpine:latest`.
+
+        Example:
+
+        .. code-block:: python
+
+            import kubetorch as kt
+
+            vol = kt.Volume.from_name("kt-global-cache")
+            vol.ssh()
+        """
+        pod_name = f"debug-{self.name}-{uuid.uuid4().hex[:6]}"
+
+        cmd = [
+            "kubectl",
+            "run",
+            pod_name,
+            "--rm",
+            "-it",
+            "--namespace",
+            self.namespace,
+            "--image",
+            image,
+            "--restart=Never",
+            "--overrides",
+            textwrap.dedent(
+                f"""
+            {{
+              "apiVersion": "v1",
+              "spec": {{
+                "containers": [{{
+                  "name": "debug",
+                  "image": "{image}",
+                  "stdin": true,
+                  "tty": true,
+                  "volumeMounts": [{{
+                    "name": "vol",
+                    "mountPath": "{self.mount_path}"
+                  }}]
+                }}],
+                "volumes": [{{
+                  "name": "vol",
+                  "persistentVolumeClaim": {{
+                    "claimName": "{self.pvc_name}"
+                  }}
+                }}]
+              }}
+            }}
+            """
+            ).strip(),
+        ]
+
+        # Suppress noisy "write on closed stream" error when exiting
+        subprocess.run(cmd, stderr=subprocess.DEVNULL)