docs: enhance common module code documentation

Author: Ygnas

src/codeflare_sdk/common/kubernetes_cluster/auth.py

@@ -165,8 +165,20 @@ def load_kube_config(self):
 
 def config_check() -> str:
     """
-    Function for loading the config file at the default config location ~/.kube/config if the user has not
-    specified their own config file or has logged in with their token and server.
+    Check and load the Kubernetes config from the default location.
+
+    This function checks if a Kubernetes config file exists at the default path
+    (`~/.kube/config`). If none is provided, it tries to load in-cluster config.
+    If the `config_path` global variable is set by an external module (e.g., `auth.py`),
+    this path will be used directly.
+
+    Returns:
+        str:
+            The loaded config path if successful.
+
+    Raises:
+        PermissionError:
+            If no valid credentials or config file is found.
     """
     global config_path
     global api_client
@@ -215,7 +227,16 @@ def _gen_ca_cert_path(ca_cert_path: Optional[str]):
 
 
 def get_api_client() -> client.ApiClient:
-    "This function should load the api client with defaults"
+    """
+    Retrieve the Kubernetes API client with the default configuration.
+
+    This function returns the current API client instance if already loaded,
+    or creates a new API client with the default configuration.
+
+    Returns:
+        client.ApiClient:
+            The Kubernetes API client object.
+    """
     if api_client != None:
         return api_client
     to_return = client.ApiClient()

src/codeflare_sdk/common/kueue/kueue.py

@@ -19,8 +19,23 @@
 from kubernetes.client.exceptions import ApiException
 
 
-def get_default_kueue_name(namespace: str):
-    # If the local queue is set, use it. Otherwise, try to use the default queue.
+def get_default_kueue_name(namespace: str) -> Optional[str]:
+    """
+    Retrieves the default Kueue name from the provided namespace.
+
+    This function attempts to fetch the local queues in the given namespace and checks if any of them is annotated
+    as the default queue. If found, the name of the default queue is returned.
+
+    The default queue is marked with the annotation "kueue.x-k8s.io/default-queue" set to "true."
+
+    Args:
+        namespace (str):
+            The Kubernetes namespace where the local queues are located.
+
+    Returns:
+        Optional[str]:
+            The name of the default queue if it exists, otherwise None.
+    """
     try:
         config_check()
         api_instance = client.CustomObjectsApi(get_api_client())
@@ -58,12 +73,14 @@ def list_local_queues(
         Depending on the version of the local queue API, the available flavors may not be present in the response.
 
     Args:
-        namespace (str, optional): The namespace to list local queues from. Defaults to None.
-        flavors (List[str], optional): The flavors to filter local queues by. Defaults to None.
+        namespace (str, optional):
+            The namespace to list local queues from. Defaults to None.
+        flavors (List[str], optional):
+            The flavors to filter local queues by. Defaults to None.
     Returns:
-        List[dict]: A list of dictionaries containing the name of the local queue and the available flavors
+        List[dict]:
+            A list of dictionaries containing the name of the local queue and the available flavors
     """
-
     from ...ray.cluster.cluster import get_current_namespace
 
     if namespace is None:  # pragma: no cover
@@ -92,8 +109,22 @@ def list_local_queues(
     return to_return
 
 
-def local_queue_exists(namespace: str, local_queue_name: str):
-    # get all local queues in the namespace
+def local_queue_exists(namespace: str, local_queue_name: str) -> bool:
+    """
+    Checks if a local queue with the provided name exists in the given namespace.
+
+    This function queries the local queues in the specified namespace and verifies if any queue matches the given name.
+
+    Args:
+        namespace (str):
+            The namespace where the local queues are located.
+        local_queue_name (str):
+            The name of the local queue to check for existence.
+
+    Returns:
+        bool:
+            True if the local queue exists, False otherwise.
+    """
     try:
         config_check()
         api_instance = client.CustomObjectsApi(get_api_client())
@@ -113,6 +144,24 @@ def local_queue_exists(namespace: str, local_queue_name: str):
 
 
 def add_queue_label(item: dict, namespace: str, local_queue: Optional[str]):
+    """
+    Adds a local queue name label to the provided item.
+
+    If the local queue is not provided, the default local queue for the namespace is used. The function validates if the
+    local queue exists, and if it does, the local queue name label is added to the resource metadata.
+
+    Args:
+        item (dict):
+            The resource where the label will be added.
+        namespace (str):
+            The namespace of the local queue.
+        local_queue (str, optional):
+            The name of the local queue to use. Defaults to None.
+
+    Raises:
+        ValueError:
+            If the provided or default local queue does not exist in the namespace.
+    """
     lq_name = local_queue or get_default_kueue_name(namespace)
     if lq_name == None:
         return

src/codeflare_sdk/common/utils/demos.py

@@ -13,10 +13,14 @@ def copy_demo_nbs(dir: str = "./demo-notebooks", overwrite: bool = False):
     Any files that exist in the directory that don't match these values will remain untouched.
 
     Args:
-        dir (str): The directory to copy the demo notebooks to. Defaults to "./demo-notebooks". overwrite (bool):
-        overwrite (bool): Whether to overwrite files in the directory if it already exists. Defaults to False.
+        dir (str):
+            The directory to copy the demo notebooks to. Defaults to "./demo-notebooks".
+        overwrite (bool):
+            Whether to overwrite files in the directory if it already exists. Defaults to False.
+
     Raises:
-        FileExistsError: If the directory already exists.
+        FileExistsError:
+            If the directory already exists.
     """
     # does dir exist already?
     if overwrite is False and pathlib.Path(dir).exists():

src/codeflare_sdk/common/utils/generate_cert.py

@@ -28,10 +28,20 @@
 
 
 def generate_ca_cert(days: int = 30):
-    # Generate base64 encoded ca.key and ca.cert
-    # Similar to:
-    # openssl req -x509 -nodes -newkey rsa:2048 -keyout ca.key -days 1826 -out ca.crt -subj '/CN=root-ca'
-    # base64 -i ca.crt -i ca.key
+    """
+    Generates a self-signed CA certificate and private key, encoded in base64 format.
+
+    Similar to:
+        openssl req -x509 -nodes -newkey rsa:2048 -keyout ca.key -days 1826 -out ca.crt -subj '/CN=root-ca'
+
+    Args:
+        days (int):
+            The number of days for which the CA certificate will be valid. Default is 30.
+
+    Returns:
+        Tuple[str, str]:
+            A tuple containing the base64-encoded private key and CA certificate.
+    """
 
     private_key = rsa.generate_private_key(
         public_exponent=65537,
@@ -79,6 +89,25 @@ def generate_ca_cert(days: int = 30):
 
 
 def get_secret_name(cluster_name, namespace, api_instance):
+    """
+    Retrieves the name of the Kubernetes secret containing the CA certificate for the given Ray cluster.
+
+    Args:
+        cluster_name (str):
+            The name of the Ray cluster.
+        namespace (str):
+            The Kubernetes namespace where the Ray cluster is located.
+        api_instance (client.CoreV1Api):
+            An instance of the Kubernetes CoreV1Api.
+
+    Returns:
+        str:
+            The name of the Kubernetes secret containing the CA certificate.
+
+    Raises:
+        KeyError:
+            If no secret matching the cluster name is found.
+    """
     label_selector = f"ray.openshift.ai/cluster-name={cluster_name}"
     try:
         secrets = api_instance.list_namespaced_secret(
@@ -97,7 +126,26 @@ def get_secret_name(cluster_name, namespace, api_instance):
 
 
 def generate_tls_cert(cluster_name, namespace, days=30):
-    # Create a folder tls-<cluster>-<namespace> and store three files: ca.crt, tls.crt, and tls.key
+    """
+    Generates a TLS certificate and key for a Ray cluster, saving them locally along with the CA certificate.
+
+    Args:
+        cluster_name (str):
+            The name of the Ray cluster.
+        namespace (str):
+            The Kubernetes namespace where the Ray cluster is located.
+        days (int):
+            The number of days for which the TLS certificate will be valid. Default is 30.
+
+    Files Created:
+        - ca.crt: The CA certificate.
+        - tls.crt: The TLS certificate signed by the CA.
+        - tls.key: The private key for the TLS certificate.
+
+    Raises:
+        Exception:
+            If an error occurs while retrieving the CA secret.
+    """
     tls_dir = os.path.join(os.getcwd(), f"tls-{cluster_name}-{namespace}")
     if not os.path.exists(tls_dir):
         os.makedirs(tls_dir)
@@ -181,6 +229,21 @@ def generate_tls_cert(cluster_name, namespace, days=30):
 
 
 def export_env(cluster_name, namespace):
+    """
+    Sets environment variables to configure TLS for a Ray cluster.
+
+    Args:
+        cluster_name (str):
+            The name of the Ray cluster.
+        namespace (str):
+            The Kubernetes namespace where the Ray cluster is located.
+
+    Environment Variables Set:
+        - RAY_USE_TLS: Enables TLS for Ray.
+        - RAY_TLS_SERVER_CERT: Path to the TLS server certificate.
+        - RAY_TLS_SERVER_KEY: Path to the TLS server private key.
+        - RAY_TLS_CA_CERT: Path to the CA certificate.
+    """
     tls_dir = os.path.join(os.getcwd(), f"tls-{cluster_name}-{namespace}")
     os.environ["RAY_USE_TLS"] = "1"
     os.environ["RAY_TLS_SERVER_CERT"] = os.path.join(tls_dir, "tls.crt")