kubetorch 0.2.0__py3-none-any.whl

kubetorch/resources/secrets/provider_secrets/ssh_secret.py

@@ -0,0 +1,12 @@
+from .. import Secret
+
+
+class SSHSecret(Secret):
+    """
+    .. note::
+            To create a SSHSecret, please use the factory method :func:`secret` with ``provider="ssh"``.
+    """
+
+    _DEFAULT_PATH = "~/.ssh"
+    _DEFAULT_FILENAMES = ["id_rsa"]
+    _PROVIDER = "ssh"

kubetorch/resources/secrets/provider_secrets/wandb_secret.py

@@ -0,0 +1,11 @@
+from .. import Secret
+
+
+class WandBSecret(Secret):
+    """
+    .. note::
+            To create an WandBSecret, please use the factory method :func:`secret` with ``provider="wandb"``.
+    """
+
+    _PROVIDER = "wandb"
+    _DEFAULT_ENV_VARS = {"api_key": "WANDB_API_KEY"}

kubetorch/resources/secrets/secret.py

@@ -0,0 +1,238 @@
+import os
+from typing import Dict, List, Optional, Tuple
+
+from kubetorch.globals import config
+
+from kubetorch.resources.secrets.utils import read_files_as_secrets_dict
+
+
+class Secret:
+    _DEFAULT_PATH = None
+    _DEFAULT_FILENAMES = None
+    _DEFAULT_ENV_VARS = {}
+    _MAP_FILENAMES_TO_ENV_VARS = {}
+    _PROVIDER = None
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        provider: Optional[str] = None,
+        values: Dict = None,
+        path: str = None,
+        env_vars: Dict = None,
+        override: bool = False,
+        **kwargs,
+    ):
+        """
+        Secret class. Built-in provider classes contain default path and/or environment variable mappings,
+        based on it's expected usage.
+
+        Note:
+            Currently supported built-in providers:
+            anthropic, aws, azure, gcp, github, huggingface, lambda, langchain, openai, pinecone, ssh, wandb.
+
+        Args:
+            name (str, optional): Name to assign the Kubetorch secret.
+            provider (str, optional): Provider corresponding to the secret (e.g. "aws", "gcp").
+            values (Dict, optional): Dictionary mapping secret keys to the corresponding secret values.
+            path (str, optional): Path where the secret values are held.
+            env_vars (Dict, optional): Dictionary mapping secret keys to the corresponding environment variable key.
+            override (bool, optional): If True, override the secret's values in Kubernetes if a secret with the same
+                name already exists.
+        """
+        name_prefix = (
+            f"{config.username}-" if config.username else ""
+        )  # we need the username as prefix in case diffrent users will create the same provider secret
+        self._name = (
+            name or f"{name_prefix}{provider}" or f"{name_prefix}{self._PROVIDER}"
+        )
+        self._name = self._name.replace(
+            "_", "-"
+        )  # cleanup so the name will match k8 standards.
+        self._namespace = kwargs.get("namespace", None) or config.namespace
+        self._values = values
+
+        self.provider = provider or self._PROVIDER
+        self.path = path
+        if path:
+            filenames = kwargs.get(
+                "filenames", None
+            )  # we might get filenames as kwarg if we load the secret from name or form config
+            updated_path, filenames = self._split_path_if_needed(
+                path=path, filenames=filenames
+            )
+            self.path = updated_path
+            self.filenames = filenames
+        self.env_vars = env_vars
+        self._override = override
+
+        if not any([values, path, env_vars]):
+            if self._values_from_path():
+                pass
+            elif self._values_from_env(self._DEFAULT_ENV_VARS):
+                self.env_vars = self._DEFAULT_ENV_VARS
+            else:
+                raise ValueError(
+                    "Secrets values not provided and could not be extracted from default file "
+                    f"({self._DEFAULT_PATH}) or env vars ({self._DEFAULT_ENV_VARS.values()}) locations."
+                )
+
+    @property
+    def name(self):
+        """Name of the secret."""
+        return self._name
+
+    @property
+    def override(self):
+        """Should we override secret's values in Kubernetes if a secret with the same name already exists"""
+        return self._override
+
+    @property
+    def values(self):
+        """Secret values."""
+        if self._values:
+            return self._values
+        if self.path:
+            return self._values_from_path(self.path)
+        if self.env_vars:
+            return self._values_from_env(self.env_vars)
+        return {}
+
+    def _values_from_env(self, env_vars: Dict = None):
+        env_vars = env_vars or self.env_vars
+        if not env_vars:
+            return {}
+        return {key: os.environ[key] for key in env_vars}
+
+    def _values_from_path(self, path: str = None):
+        path = path or self.path or self._DEFAULT_PATH
+        if not path:
+            return {}
+
+        # Double-check that the path is a directory
+        path, filenames = self._split_path_if_needed(path)
+
+        values = read_files_as_secrets_dict(path=path, filenames=filenames)
+        if values:
+            # Only set if the values were successfully found
+            if self._MAP_FILENAMES_TO_ENV_VARS:
+                env_vars = []
+                for filename, env_var in self._MAP_FILENAMES_TO_ENV_VARS.items():
+                    if filename in values:
+                        values[env_var] = values[filename].strip()
+                        del values[filename]
+                        env_vars.append(env_var)
+                if env_vars:
+                    self.env_vars = env_vars
+                    self._values = values
+                    return values
+
+            self.path = path
+            self.filenames = filenames
+        return values
+
+    def _split_path_if_needed(
+        self, path: str, filenames: list = None
+    ) -> Tuple[str, List[str]]:
+        """Split path into path and filesnames if a single file is specified as a full path"""
+        updated_path = path
+        is_default_path = updated_path == self._DEFAULT_PATH
+        updated_filenames = getattr(self, "filenames", None) or filenames
+        if not updated_filenames:
+            if not is_default_path or not self._DEFAULT_FILENAMES:
+                # Reform single-file path a directory and filenames list
+                updated_filenames = [os.path.basename(path)]
+                updated_path = os.path.dirname(path)
+            else:
+                updated_filenames = self._DEFAULT_FILENAMES
+        return updated_path, updated_filenames
+
+    @classmethod
+    def from_config(cls, config: dict):
+        override_value = config.get("override", "False").lower()
+        bool_override_value = override_value == "true"
+        config["override"] = bool_override_value
+        if "provider" in config:
+            from .provider_secrets.providers import _get_provider_class
+
+            provider_class = _get_provider_class(config["provider"])
+            return provider_class.from_config(config)
+        return cls(**config)
+
+    @classmethod
+    def from_name(cls, name, namespace: str = config.namespace):
+
+        from kubetorch.resources.secrets.kubernetes_secrets_client import (
+            KubernetesSecretsClient,
+        )
+
+        secrets_client = KubernetesSecretsClient(namespace=namespace)
+        secret = secrets_client.load_secret(name=name)
+        return secret
+
+    @classmethod
+    def builtin_providers(cls, as_str: bool = False) -> List:
+        """Return list of all Kubetorch providers (as class objects) supported out of the box.
+
+        Args:
+            as_str (bool, optional): Whether to return the providers as a string or as a class.
+                (Default: ``False``)
+        """
+        from .provider_secrets.providers import _str_to_provider_class
+
+        if as_str:
+            return list(_str_to_provider_class.keys())
+        return list(_str_to_provider_class.values())
+
+    @classmethod
+    def from_provider(
+        cls, provider: str, name: str = None, path: str = None, override: bool = False
+    ):
+        """Return kubetorch provider secret object
+
+        Args:
+            provider (str): Provider's name
+            name (str, Optional): Secret name
+            path (str, optional): Path where the secret values are held.
+            override (Bool, optional): If True, override the secret's values in Kubernetes if a secret with the same name already exists.
+        """
+        from .provider_secrets.providers import _get_provider_class
+
+        secret_class = _get_provider_class(provider)
+        if not secret_class:
+            raise ValueError(
+                f"{provider} is not a supported provider: {Secret.builtin_providers(as_str=True)}"
+            )
+        return secret_class(name=name, provider=provider, path=path, override=override)
+
+    @classmethod
+    def from_path(cls, path: str, name: str = None, override: bool = False):
+        """Return kubetorch provider secret object
+
+        Args:
+            path (str): Local path to the secret values file
+            name (str, Optional): Secret name
+            override (Bool, optional): If True, override the secret's values in Kubernetes if a secret with the same name already exists.
+        """
+        from .provider_secrets.providers import _get_provider_class
+
+        secret_class = _get_provider_class(path) or Secret
+        if not secret_class._PROVIDER and not name:
+            raise ValueError("secret name must be provided.")
+
+        return secret_class(name=name, path=path, override=override)
+
+    @classmethod
+    def from_env(cls, env_vars: dict, name: str = None, override: bool = False):
+        """Return kubetorch provider secret object
+
+        Args:
+            env_vars (dict): Dictionary mapping secret keys to the corresponding
+            environment variable key.
+            name (str, Optional): Secret name
+            override (Bool, optional): If True, override the secret's values in Kubernetes if a secret with the same name already exists.
+        """
+        from .provider_secrets.providers import _get_provider_class
+
+        secret_class = _get_provider_class(env_vars) or Secret
+        return secret_class(name=name, env_vars=env_vars, override=override)

kubetorch/resources/secrets/secret_factory.py

@@ -0,0 +1,70 @@
+from typing import Dict, Optional
+
+from kubetorch.globals import config
+
+from .secret import Secret
+
+
+def secret(
+    name: Optional[str] = None,
+    provider: Optional[str] = None,
+    path: Optional[str] = None,
+    env_vars: Optional[Dict] = None,
+    namespace: Optional[str] = config.namespace,
+    override: Optional[bool] = False,
+) -> Secret:
+    """
+    Builds an instance of :class:`Secret`. At most one of `provider`, `path`, or `env_vars` can be provided, to maintain
+    one source of truth. For a provider, the values are inferred from the default path or environment variables for that
+    provider. To load a secret by name, provide its name and namespace.
+
+    Args:
+        namespace (str, optional): Namespace to load the secret from, if we create a secret from name. Default: "default".
+        name (str, optional): Name to assign the resource. If none is provided, resource name defaults to the
+            provider name.
+        provider (str, optional): Provider corresponding to the secret (e.g. "aws", "gcp"). To see all supported provider
+            types, run ``kt.Secret.builtin_providers(as_str=True)``.
+        path (str, optional): Path where the secret values are held.
+        env_vars (Dict, optional): Dictionary mapping secret keys to the corresponding
+            environment variable key.
+        override (Bool, optional): If True, override the secret's values in Kubernetes if a secret with the same name already exists.
+
+    Returns:
+        Secret: The resulting secret object.
+
+    Examples:
+
+    .. code-block:: python
+
+        import kubetorch as kt
+
+        local_secret = kt.secret(name="in_memory_secret", values={"secret_key": "secret_val"})
+        aws_secret = kt.secret(provider="aws")
+        gcp_secret = kt.secret(name="my-gcp-secret", path="~/.gcp/credentials")
+        lambda_secret = kt.secret(name= "my-lambda-secret", env_vars={"api_key": "LAMBDA_API_KEY"})
+    """
+
+    # env_vars or path or provider are provided
+    valid_input = sum([bool(x) for x in [provider, path, env_vars]]) == 1 or (
+        provider and path
+    )
+    valid_from_name_input = (
+        sum([bool(x) for x in [provider, path, env_vars]]) == 0 and name
+    )
+
+    if not (valid_from_name_input or valid_input):
+        raise ValueError(
+            "You must provide exactly one of: `provider`, `path`, or `env_vars`. Alternatively, you may provide `name` to load a secret from name."
+        )
+
+    if valid_input:
+        if provider:
+            return Secret.from_provider(
+                provider=provider, name=name, path=path, override=override
+            )
+        elif path and not provider:  # the case where provider + path are provided are
+            return Secret.from_path(path=path, name=name, override=override)
+        elif env_vars:
+            return Secret.from_env(env_vars=env_vars, name=name, override=override)
+    else:
+        return Secret.from_name(name=name, namespace=namespace)

kubetorch/resources/secrets/utils.py

@@ -0,0 +1,209 @@
+import os
+import re
+import time
+from typing import List, Optional
+
+from kubernetes import client, config
+from kubernetes.client import V1Pod, V1TokenReview, V1TokenReviewSpec
+from kubernetes.stream import stream
+
+from kubetorch.globals import config as kt_config
+
+from kubetorch.logger import get_logger
+from kubetorch.servers.http.utils import is_running_in_kubernetes
+
+logger = get_logger(__name__)
+
+
+def get_k8s_identity_name() -> Optional[str]:
+    try:
+        if is_running_in_kubernetes():
+            config.load_incluster_config()
+        else:
+            config.load_kube_config()
+        configuration = client.Configuration.get_default_copy()
+
+        token = configuration.api_key.get("authorization", "")
+        token = re.sub(r"^Bearer\s", "", token)
+
+        api = client.AuthenticationV1Api()
+        token_review = V1TokenReview(spec=V1TokenReviewSpec(token=token))
+
+        result = api.create_token_review(token_review)
+
+        if result.status.authenticated:
+            user = result.status.user
+            # For EKS IAM users/roles
+            if hasattr(user, "username") and user.username.endswith("amazonaws.com"):
+                # ARN format: arn:aws:iam::ACCOUNT_ID:user/USERNAME or
+                # arn:aws:sts::ACCOUNT_ID:assumed-role/ROLE_NAME/SESSION_NAME
+                arn_parts = user.username.split("/")
+                if "assumed-role" in user.username:
+                    return "role-" + arn_parts[-2].lower()  # Returns ROLE_NAME
+                return "user-" + arn_parts[-1].lower()  # Returns USERNAME for IAM users
+            # For Kubernetes service accounts (works for both GKE and EKS)
+            elif (
+                hasattr(user, "username") and "system:serviceaccount:" in user.username
+            ):
+                return (
+                    "sa-" + user.username.split(":")[-1].lower()
+                )  # Returns service account name
+
+    except Exception as e:
+        logger.info(f"Failed to get identity name: {e}")
+
+    return None
+
+
+def read_files_as_secrets_dict(path: str, filenames: List[str]):
+    values = {}
+    cred_path = os.path.expanduser(path)
+
+    for filename in filenames:
+        file_path = os.path.join(cred_path, filename)
+        # Read the files
+        content = _read_file_if_exists(file_path)
+        if content:
+            values[filename] = content
+        # # Base64 encode the content
+        # encoded = base64.b64encode(content).decode("utf-8")
+        # values[filename] = encoded
+
+    return values
+
+
+def _read_file_if_exists(file_path: str) -> Optional[str]:
+    try:
+        with open(file_path, "r") as f:  # "rb" if you encode above.
+            return f.read()
+    except FileNotFoundError:
+        logger.error(f"Warning: {file_path} not found, using empty content")
+        return None
+
+
+# ------------------------------------------------------------------------------------------------
+# Secret testing utils
+# ------------------------------------------------------------------------------------------------
+
+
+def check_path_on_kubernetes_pods(
+    path: str, service_name: str, namespace: str = None
+) -> bool:
+    """
+    Check if a path exists on a specific Knative service's pods
+    """
+    namespace = namespace or kt_config.namespace
+    # Load Kubernetes configuration
+    config.load_kube_config()
+    # Initialize API clients
+    core_v1_api = client.CoreV1Api()
+
+    pods = _fetch_pods_for_kubernetes_service(service_name, namespace, core_v1_api)
+    if not pods:
+        logger.error(
+            f"No pods found for service {service_name} in namespace {namespace}"
+        )
+        return False
+
+    path_found = True
+    for pod in pods:
+        pod_name = pod.metadata.name
+        command = ["/bin/bash", "-c", f"[ -f {path} ] && echo yes || echo no"]
+        try:
+            resp = stream(
+                core_v1_api.connect_get_namespaced_pod_exec,
+                name=pod_name,
+                namespace=namespace,
+                command=command,
+                container="kubetorch",
+                stderr=True,
+                stdout=True,
+            )
+            if "yes" in resp:
+                continue
+        except client.exceptions.ApiException as e:
+            logger.error(f"Error executing command on pod {pod_name}: {e}")
+
+        path_found = False
+
+    return path_found
+
+
+def check_env_vars_on_kubernetes_pods(
+    env_vars: list, service_name: str, namespace: str = None
+) -> dict:
+    """
+    Check if an AWS role is assumed on a specific Knative service's pods
+
+    :param namespace: Kubernetes namespace
+    :param service_name: Name of the Knative service
+    :return: Dictionary with role assumption details
+    """
+    namespace = namespace or kt_config.namespace
+    # Load Kubernetes configuration
+    config.load_kube_config()
+    # Initialize API clients
+    core_v1_api = client.CoreV1Api()
+
+    pods = _fetch_pods_for_kubernetes_service(service_name, namespace, core_v1_api)
+    if not pods:
+        logger.error(
+            f"No pods found for service {service_name} in namespace {namespace}"
+        )
+        return {}
+
+    found_env_vars = {}
+
+    for pod in pods:
+        for env_var in env_vars:
+            if found_env_vars.get(env_var):
+                # Skip if already found on another pod
+                continue
+            pod_name = pod.metadata.name
+            command = ["/bin/bash", "-c", f"echo ${env_var}"]
+            try:
+                resp = stream(
+                    core_v1_api.connect_get_namespaced_pod_exec,
+                    name=pod_name,
+                    namespace=namespace,
+                    command=command,
+                    container="kubetorch",
+                    stderr=True,
+                    stdout=True,
+                )
+                if len(resp.strip()) > 0:
+                    found_env_vars[env_var] = resp.strip()
+            except client.exceptions.ApiException as e:
+                logger.error(f"Error executing command: {e}")
+
+        if set(found_env_vars.keys()) == set(env_vars):
+            # Found all env vars: skip the remaining pods
+            break
+
+    return found_env_vars
+
+
+def _fetch_pods_for_kubernetes_service(
+    service_name: str, namespace: str, client_api: client.CoreV1Api
+) -> List[V1Pod]:
+    """
+    Fetch pods for a specific Knative service with timeout
+    """
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        try:
+            # List pods matching the service
+            pods = client_api.list_namespaced_pod(
+                namespace=namespace,
+                label_selector=f"kubetorch.com/service={service_name}",
+            )
+            ready_pods = [pod for pod in pods.items if pod.status.phase == "Running"]
+            if ready_pods:
+                return ready_pods
+        except Exception as e:
+            logger.error(
+                f"Error fetching pods for service {service_name} in namespace {namespace}: {e}"
+            )
+        time.sleep(1)
+
+    return []