kubetorch 0.2.0__py3-none-any.whl

kubetorch/resources/compute/decorators.py

@@ -0,0 +1,139 @@
+import os
+from functools import update_wrapper
+from typing import List, Union
+
+from kubetorch.resources.callables.cls.cls import cls
+from kubetorch.resources.callables.fn.fn import fn
+from kubetorch.resources.compute.compute import Compute
+
+# Helper class which allows us to chain decorators in a way that allows us to reverse the order of the decorators
+# The `compute` decorator ultimately unwinds the calls to properly construct the Module.
+class PartialModule:
+    def __init__(
+        self,
+        fn_or_cls=None,
+        distribute_args=None,
+        autoscale_args=None,
+        async_=False,
+    ):
+        self.fn_or_cls = fn_or_cls
+        self.distribute_args = distribute_args
+        self.autoscale_args = autoscale_args
+        self.async_ = async_
+
+
+# @kubetorch.compute decorator that the user can use to wrap a function they want to deploy to a cluster,
+# and then deploy it with `kt deploy my_app.py` (we collect all the decorated functions imported in the file
+# to deploy them).
+def compute(
+    get_if_exists: bool = False, reload_prefixes: Union[str, List[str]] = [], **kwargs
+):
+    def decorator(func_or_cls):
+        from kubetorch.globals import disable_decorators
+
+        if disable_decorators:
+            return func_or_cls
+
+        if isinstance(func_or_cls, PartialModule):
+            distribute_args = func_or_cls.distribute_args
+            autoscale_args = func_or_cls.autoscale_args
+            async_ = func_or_cls.async_
+            func_or_cls = func_or_cls.fn_or_cls
+        else:
+            distribute_args = None
+            autoscale_args = None
+            async_ = False
+
+        # If we're on the server attempting to load this function or class, just return it as is
+        if (
+            os.environ.get("KT_CLS_OR_FN_NAME") == func_or_cls.__name__
+            and os.environ.get("KT_MODULE_NAME") == func_or_cls.__module__
+        ):
+            return func_or_cls
+
+        module_name = kwargs.pop("name", None)
+
+        if isinstance(func_or_cls, type):
+            new_module = cls(
+                func_or_cls,
+                name=module_name,
+                get_if_exists=get_if_exists,
+                reload_prefixes=reload_prefixes,
+            )
+        else:
+            new_module = fn(
+                func_or_cls,
+                name=module_name,
+                get_if_exists=get_if_exists,
+                reload_prefixes=reload_prefixes,
+            )
+
+        # All remaining kwargs are passed to the Compute constructor
+        new_module.compute = Compute(**kwargs)
+        new_module.compute.service_name = new_module.service_name
+        if distribute_args:
+            distribute_args, distribute_kwargs = distribute_args
+            new_module.compute.distribute(*distribute_args, **distribute_kwargs)
+        if autoscale_args:
+            autoscale_args, autoscale_kwargs = autoscale_args
+            new_module.compute.autoscale(*autoscale_args, **autoscale_kwargs)
+        if async_:
+            new_module.async_ = async_
+
+        # update_wrapper(new_module, func_or_cls)
+        return new_module
+
+    return decorator
+
+
+def async_(func_or_cls):
+    from kubetorch.globals import disable_decorators
+
+    if disable_decorators:
+        return func_or_cls
+
+    # If it's already a PartialModule (from other decorators), update it
+    if isinstance(func_or_cls, PartialModule):
+        func_or_cls.async_ = True
+        return func_or_cls
+
+    # Otherwise, create a new PartialModule
+    partial_module = PartialModule(fn_or_cls=func_or_cls, async_=True)
+    update_wrapper(partial_module, func_or_cls)
+    return partial_module
+
+
+def distribute(*args, **kwargs):
+    def decorator(func_or_cls):
+        from kubetorch.globals import disable_decorators
+
+        if disable_decorators:
+            return func_or_cls
+
+        # This is a partial so the order of decorator chaining can be reversed for best aesthetics
+        # the deploy method will actually call .distribute on the function or class after it's been deployed
+        partial_module = PartialModule(
+            fn_or_cls=func_or_cls, distribute_args=(args, kwargs)
+        )
+        update_wrapper(partial_module, func_or_cls)
+        return partial_module
+
+    return decorator
+
+
+def autoscale(*args, **kwargs):
+    def decorator(func_or_cls):
+        from kubetorch.globals import disable_decorators
+
+        if disable_decorators:
+            return func_or_cls
+
+        # This is a partial so the order of decorator chaining can be reversed for best aesthetics
+        # the deploy method will actually call .distribute on the function or class after it's been deployed
+        partial_module = PartialModule(
+            fn_or_cls=func_or_cls, autoscale_args=(args, kwargs)
+        )
+        update_wrapper(partial_module, func_or_cls)
+        return partial_module
+
+    return decorator

kubetorch/resources/compute/rbac.py

@@ -0,0 +1,74 @@
+from typing import Optional
+
+from kubernetes import client, config
+
+from kubetorch import globals
+from kubetorch.logger import get_logger
+
+from kubetorch.serving.constants import DEFAULT_SERVICE_ACCOUNT_NAME
+
+logger = get_logger(__name__)
+
+
+def find_service_account_in_namespace(
+    namespace: str, service_account_name: str, core_v1: client.CoreV1Api
+) -> Optional[client.V1ServiceAccount]:
+    """Find a ServiceAccount in a given namespace."""
+    try:
+        return core_v1.read_namespaced_service_account(
+            name=service_account_name, namespace=namespace
+        )
+    except client.exceptions.ApiException as e:
+        if e.status != 404:
+            raise e
+        return None
+
+
+def cleanup_namespace_and_service_account(namespace: str) -> None:
+    """
+    Cleanup a ServiceAccount and its associated ClusterRoleBindings in a given namespace.
+    """
+    if namespace == globals.config.install_namespace:
+        raise ValueError(
+            f"Cannot cleanup ${globals.config.install_namespace} ServiceAccount"
+        )
+
+    # Load kube config (works for both local and in-cluster)
+    try:
+        config.load_incluster_config()
+    except config.ConfigException:
+        config.load_kube_config()
+
+    core_v1 = client.CoreV1Api()
+    rbac_v1 = client.RbacAuthorizationV1Api()
+
+    service_account = find_service_account_in_namespace(
+        namespace, DEFAULT_SERVICE_ACCOUNT_NAME, core_v1
+    )
+    if not service_account:
+        return
+
+    # Delete the ServiceAccount
+    core_v1.delete_namespaced_service_account(
+        name=DEFAULT_SERVICE_ACCOUNT_NAME, namespace=namespace
+    )
+    logger.info(f"Deleted ServiceAccount {namespace}/{DEFAULT_SERVICE_ACCOUNT_NAME}")
+
+    # Delete all ClusterRoleBindings that reference the ServiceAccount
+    bindings = rbac_v1.list_namespaced_role_binding(
+        namespace=namespace,
+        label_selector=f"kubetorch.com/service-account={DEFAULT_SERVICE_ACCOUNT_NAME},kubetorch.com/namespace={namespace}",
+    )
+    for binding in bindings.items:
+        rbac_v1.delete_namespaced_role_binding(
+            name=binding.metadata.name, namespace=namespace
+        )
+        logger.info(f"Deleted RoleBinding {binding.metadata.name}")
+
+    logger.info(
+        f"Cleanup completed for ServiceAccount {namespace}/{DEFAULT_SERVICE_ACCOUNT_NAME}"
+    )
+
+    # Delete the namespace
+    core_v1.delete_namespace(name=namespace)
+    logger.info(f"Deleted namespace {namespace}")