kblaunch 0.1.0__py3-none-any.whl

kblaunch/cli.py

@@ -0,0 +1,195 @@
+import base64
+import os
+
+import typer
+from kubejobs.jobs import KubernetesJob, KueueQueue
+from kubernetes import client, config
+from loguru import logger
+
+app = typer.Typer()
+
+
+def check_if_completed(job_name: str, namespace: str = "informatics") -> bool:
+    # Load the kube config
+    config.load_kube_config()
+
+    # Create an instance of the API class
+    api = client.BatchV1Api()
+
+    job_exists = False
+    is_completed = True
+
+    # Check if the job exists in the specified namespace
+    jobs = api.list_namespaced_job(namespace)
+    if job_name in {job.metadata.name for job in jobs.items}:
+        job_exists = True
+
+    if job_exists is True:
+        job = api.read_namespaced_job(job_name, namespace)
+        is_completed = False
+
+        # Check the status conditions
+        if job.status.conditions:
+            for condition in job.status.conditions:
+                if condition.type == "Complete" and condition.status == "True":
+                    is_completed = True
+                elif condition.type == "Failed" and condition.status == "True":
+                    logger.error(f"Job {job_name} has failed.")
+        else:
+            logger.info(f"Job {job_name} still running or status is unknown.")
+
+        if is_completed:
+            api_res = api.delete_namespaced_job(
+                name=job_name,
+                namespace=namespace,
+                body=client.V1DeleteOptions(propagation_policy="Foreground"),
+            )
+            logger.info(f"Job '{job_name}' deleted. Status: {api_res.status}")
+    return is_completed
+
+
+def send_message_command(env_vars: dict) -> str:
+    if "SLACK_WEBHOOK" not in env_vars:
+        logger.debug("SLACK_WEBHOOK not found in env_vars.")
+        return ""
+    webhook = env_vars["SLACK_WEBHOOK"]
+    return (
+        """curl -X POST -H 'Content-type: application/json' --data '{"text":"Job started in '"$POD_NAME"'"}' """
+        + webhook
+        + " ; "
+    )
+
+
+def get_env_vars(
+    local_env_vars: list[str],
+    secrets_env_vars: list[str],
+    load_dotenv: bool = False,
+    namespace: str = "informatics",
+) -> dict:
+    """Get environment variables from local environment and secrets."""
+
+    if load_dotenv:
+        from dotenv import load_dotenv
+
+        load_dotenv()
+
+    env_vars = {}
+    for var_name in local_env_vars:
+        value = os.getenv(var_name)
+        if value is not None:
+            env_vars[var_name] = value
+        else:
+            logger.warning(
+                f"Environment variable {var_name} not found in local environment"
+            )
+
+    for secret_name in secrets_env_vars:
+        v1 = client.CoreV1Api()
+        try:
+            secret = v1.read_namespaced_secret(name=secret_name, namespace=namespace)
+            for key, value in secret.data.items():
+                decoded_value = base64.b64decode(value).decode("utf-8")
+                if key in env_vars:
+                    logger.warning(f"Key {key} already set in env_vars.")
+                env_vars[key] = decoded_value
+        except Exception as e:
+            logger.error(f"Error reading secret {secret_name}: {e}")
+            raise e
+
+    return env_vars
+
+
+def export_env_vars(env_vars: dict) -> str:
+    """Export environment variables."""
+    cmd = ""
+    for key, value in env_vars.items():
+        cmd += f" export {key}='{value}' &&"
+    cmd = cmd.strip(" &&") + " ; "
+    return cmd
+
+
+@app.command()
+def launch(
+    email: str = typer.Option(..., help="User email"),
+    job_name: str = typer.Option(..., help="Name of the Kubernetes job"),
+    docker_image: str = typer.Option(
+        "nvcr.io/nvidia/cuda:12.0.0-devel-ubuntu22.04", help="Docker image"
+    ),
+    namespace: str = typer.Option("informatics", help="Kubernetes namespace"),
+    queue_name: str = typer.Option(KueueQueue.INFORMATICS, help="Kueue queue name"),
+    interactive: bool = typer.Option(False, help="Run in interactive mode"),
+    command: str = typer.Option(..., help="Command to run in the container"),
+    cpu_request: str = typer.Option("1", help="CPU request"),
+    ram_request: str = typer.Option("8Gi", help="RAM request"),
+    gpu_limit: int = typer.Option(1, help="GPU limit"),
+    gpu_product: str = typer.Option("NVIDIA-A100-80GB-PCIe", help="GPU product"),
+    secrets_env_vars: list[str] = typer.Option(
+        [],  # Use empty list as default instead of None
+        help="List of secret environment variables to export to the container",
+    ),
+    local_env_vars: list[str] = typer.Option(
+        [],  # Use empty list as default instead of None
+        help="List of local environment variables to export to the container",
+    ),
+    load_dotenv: bool = typer.Option(
+        True, help="Load environment variables from .env file"
+    ),
+):
+    """Launch a Kubernetes job with the specified configuration."""
+
+    is_completed = check_if_completed(job_name, namespace=namespace)
+
+    if is_completed is True:
+        logger.info(f"Job '{job_name}' is completed. Launching a new job.")
+
+        if interactive:
+            cmd = "while true; do sleep 60; done;"
+        else:
+            cmd = command
+            logger.info(f"Command: {cmd}")
+
+        # Get local environment variables
+        env_vars = get_env_vars(
+            local_env_vars=local_env_vars,
+            secrets_env_vars=secrets_env_vars,
+            load_dotenv=load_dotenv,
+            namespace=namespace,
+        )
+
+        logger.info(f"Creating job for: {cmd}")
+        job = KubernetesJob(
+            name=job_name,
+            cpu_request=cpu_request,
+            ram_request=ram_request,
+            image=docker_image,
+            gpu_type="nvidia.com/gpu",
+            gpu_limit=gpu_limit,
+            gpu_product=gpu_product,
+            backoff_limit=0,
+            command=["/bin/bash", "-c", "--"],
+            args=[export_env_vars(env_vars) + send_message_command(env_vars) + cmd],
+            user_email=email,
+            namespace=namespace,
+            kueue_queue_name=queue_name,
+            secret_env_vars=env_vars,
+            volume_mounts={
+                "nfs": {"mountPath": "/nfs", "server": "10.24.1.255", "path": "/"}
+            },
+        )
+
+        job_yaml = job.generate_yaml()
+        logger.info(job_yaml)
+
+        # Run the Job on the Kubernetes cluster
+        job.run()
+    else:
+        logger.info(f"Job '{job_name}' is still running.")
+
+
+def cli():
+    """Entry point for the application"""
+    app()
+
+
+if __name__ == "__main__":
+    cli()

kblaunch-0.1.0.dist-info/METADATA

@@ -0,0 +1,126 @@
+Metadata-Version: 2.4
+Name: kblaunch
+Version: 0.1.0
+Summary: A CLI tool for launching Kubernetes jobs in EIDF
+Author-email: Gautier Dagan <gautier.dagan@ed.ac.uk>
+License-Expression: MIT
+License-File: LICENSE
+Requires-Python: >=3.9
+Requires-Dist: kubejobs>=0.4.5
+Requires-Dist: kubernetes>=31.0.0
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: python-dotenv>=1.0.1
+Requires-Dist: typer>=0.15.1
+Provides-Extra: test
+Requires-Dist: pytest-cov>=4.1.0; extra == 'test'
+Requires-Dist: pytest-mock>=3.10.0; extra == 'test'
+Requires-Dist: pytest>=7.0; extra == 'test'
+Description-Content-Type: text/markdown
+
+# kblaunch
+
+[![Test](https://github.com/gautierdag/kblaunch/actions/workflows/test.yaml/badge.svg)](https://github.com/gautierdag/kblaunch/actions/workflows/test.yaml)
+![Python Version](https://img.shields.io/badge/python-3.9+-blue)
+![Ruff](https://img.shields.io/badge/linter-ruff-blue)
+[![PyPI Version](https://img.shields.io/pypi/v/kblaunch)](https://pypi.org/project/kblaunch/)
+
+A CLI tool for launching Kubernetes jobs with environment variable and secret management.
+
+## Installation
+
+```bash
+pip install git+https://github.com/gautierdag/kblaunch.git
+```
+
+## Usage
+
+### Basic Usage
+
+Launch a simple job:
+
+```bash
+kblaunch \
+    --email your.email@ed.ac.uk \
+    --job-name myjob \
+    --command "python script.py"
+```
+
+### With Environment Variables
+
+1. From local environment:
+
+    ```bash
+    kblaunch \
+        --job-name myjob \
+        --command "python script.py" \
+        --local-env-vars PATH \
+        --local-env-vars PYTHONPATH
+    ```
+
+2. From Kubernetes secrets:
+
+    ```bash
+    kblaunch \
+        --job-name myjob \
+        --command "python script.py" \
+        --secrets-env-vars mysecret1 \
+        --secrets-env-vars mysecret2
+    ```
+
+3. From .env file:
+
+    ```bash
+    kblaunch \
+        --job-name myjob \
+        --command "python script.py" \
+        --load-dotenv
+    ```
+
+### GPU Jobs
+
+Specify GPU requirements:
+
+```bash
+kblaunch \
+    --job-name gpu-job \
+    --command "python train.py" \
+    --gpu-limit 2 \
+    --gpu-product "NVIDIA-A100-80GB-PCIe"
+```
+
+### Interactive Mode
+
+Launch an interactive job:
+
+```bash
+kblaunch \
+    --job-name interactive \
+    --interactive
+```
+
+## Options
+
+- `--email`: User email [required]
+- `--job-name`: Name of the Kubernetes job [required]
+- `--docker-image`: Docker image (default: "nvcr.io/nvidia/cuda:12.0.0-devel-ubuntu22.04")
+- `--namespace`: Kubernetes namespace (default: "informatics")
+- `--queue-name`: Kueue queue name
+- `--interactive`: Run in interactive mode (default: False)
+- `--command`: Command to run in the container [required]
+- `--cpu-request`: CPU request (default: "1")
+- `--ram-request`: RAM request (default: "8Gi")
+- `--gpu-limit`: GPU limit (default: 1)
+- `--gpu-product`: GPU product (default: "NVIDIA-A100-80GB-PCIe")
+- `--secrets-env-vars`: List of secret environment variables
+- `--local-env-vars`: List of local environment variables
+- `--load-dotenv`: Load environment variables from .env file (default: True)
+
+## Features
+
+- Kubernetes job management
+- Environment variable handling from multiple sources
+- Kubernetes secrets integration
+- GPU job support
+- Interactive mode
+- Automatic job cleanup
+- Slack notifications (when configured)

kblaunch-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+kblaunch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kblaunch/cli.py,sha256=pV_YD5YeQq_ifJhre9iT-BcvY5ns1ZzUJNiX8uCRfdQ,6452
+kblaunch-0.1.0.dist-info/METADATA,sha256=bxOxNiD__aaFu_cNAktwN23Z3iKNeO7xdStmb1InTXw,3247
+kblaunch-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kblaunch-0.1.0.dist-info/entry_points.txt,sha256=XQxzjiVPA2ZKX8iIryH4_3VIrPHTREklu--gEnpWcB4,46
+kblaunch-0.1.0.dist-info/licenses/LICENSE,sha256=Jw_Wqf5swgrgy2K0Pgat5DOoHDct-nz_AsIQ-yWI1HI,1070
+kblaunch-0.1.0.dist-info/RECORD,,

kblaunch-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

kblaunch-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+kblaunch = kblaunch.cli:app

kblaunch-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Gautier Dagan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.