jet-k8s 0.1.3

A CLI Job Execution Toolkit (Jet) for Kubernetes

Jet: A CLI Job Execution Toolkit (Jet) for Kubernetes

Skip the YAML. A lightweight command-line Job Execution Toolkit (Jet) for Kubernetes that simplifies batch job management with a focus on ML workloads.

Features

• 🚀 Simplified Job Submission - Define and submit Kubernetes jobs directly from the command line without writing YAML files manually.
• 📊 Easy Monitoring - Track and manage batch jobs with a fast and responsive Terminal User Interface (TUI).
• 📄 Work with Templates - Save custom job templates to standardize and simplify job configurations, making your experiments reproducible.
• 📓 Jupyter Integration - Launch Jupyter notebooks on Kubernetes with automatic port forwarding.
• 🐛 Debug Sessions - Spin up interactive debug pods for quick troubleshooting.
• 🤖 ML Focused - Designed with Python machine learning workloads and data processing tasks in mind.

Overview

Jet eliminates the complexity of Kubernetes YAML configuration files, providing a streamlined CLI experience for:

• Defining and submitting batch jobs
• Monitoring job status and logs with a lightweight and fast Terminal User Interface (TUI) inspired by k9s.
• Running interactive Jupyter notebook sessions on Kubernetes with automatic port forwarding.
• Creating interactive shell debug environments for troubleshooting and debugging.
• Automatic job cleanup for Jupyter and debug sessions.

Perfect for ML engineers and researchers who want to leverage Kubernetes for ML training, inference and experimentation jobs without the YAML overhead.

Demos

[!TIP] Click the GIFs to view the demos on Asciinema player with media controls.

Submitting Jobs

Monitoring Jobs with TUI

Starting Jupyter Notebook Sessions and Auto Port-forwarding

Starting Interactive Debug Sessions

Saving and Using Job Templates

Installation

Dependencies

1. Python 3.8.1 or higher.

2. kubectl installed and configured on your local machine. Refer to the official Kubernetes documentation for installation instructions.

3. A running Kubernetes cluster, with kubeconfig properly set up to access the cluster from your local machine.

Install Jet

Jet can be installed using pip from PyPI:

pip install jet-k8s

Usage

After installation, you can use the jet command in your terminal. Here are some basic commands:

Please refer to the following sections for detailed user guides.

• Submitting Jobs
• Starting Jupyter Notebook Sessions
• Starting Debug Sessions
• Using Job Templates
• Monitoring Jobs
• Other Commands

Why Jobs?

Some key reasons for using Kubernetes Jobs for ML workloads:

1. Batch Workloads: Jobs are designed for batch processing tasks, which aligns well with ML training and data processing workloads that are typically non-interactive and run to completion.
2. Automatic Retry: Jobs have built-in retry mechanisms for failed tasks, which is beneficial for long-running ML jobs that may encounter transient failures.
3. Resource Management: Jobs can be scheduled and managed more effectively with schedulers such as KAI-scheduler. For example, pods within jobs can be prempted and automatically rescheduled on different nodes if a high priority job needs resources or to organize pods to optimize cluster resource utilization.
4. Completion Tracking: Jobs provide a clear way to track the completion status of tasks, making it easier to manage and monitor ML workloads.

Notes

1. Jet currently supports some of the frequently used job and pod spec configurations through command-line arguments. If you have specific requirements not currently supported, please raise an issue or contribute a PR to add the necessary features.

2. Jet currently supports only Kubernetes clusters with NVIDIA GPU nodes.

3. Pod's restartPolicy is set to Never for all jobs types by default and job's themselves have backoffLimit set to None (so defaults to Kubernetes defaults of 6). This configuration is to ensure that when the containers in pods fail, they are not restarted indefinitely on the same resources, but instead rescheduled on different resources by the job controller. You can override this using the --restart-policy argument.

4. The argument --gpu-type is implemented using node selectors. Ensure that your cluster nodes are labeled appropriately for the GPU types you intend to use. For example, to label a node with an A100 GPU, you can use:

   kubectl label nodes <node-name> gpu-type=a100
   

5. The pod security context is set to run containers with the same user and group ID as the user executing the jet command. This is to ensure proper file permission handling when mounting host directories or volumes. If your use case requires running containers with different user/group IDs, please raise an issue or contribute a PR to make this configurable.

6. The --pyenv argument mounts a Python virtual environment from the host into the container at the same path and updates the container's PATH accordingly.

   • Requirements:

     • Shared storage: The venv directory must be accessible at the same path on the node where the pod runs. This works automatically with single-node clusters or shared storage (NFS, BeeGFS), but may not work on multi-node clusters without shared storage.

     • Python compatibility: The venv's Python executable (read from pyvenv.cfg) must be available inside the container. This works if:

       • The container image has Python installed at the same path (e.g., /usr/bin/python3.x for system Python venvs), or
       • The venv includes its own Python rather than system Python (e.g., venvs created with uv or conda using a specific Python version).

TODOs:

• Add support for fractional GPUs using HAMi plugin (In dev: KAI-scheduler #60).
• Add support for other accelerator types such as AMDs and TPUs.
• Evaluate support for other kubernetes schedulers such as Volcano.
• Ability to submit jobs with parallism and gang scheduling for usecases such as multi-node training jobs.
• Add support for job dependencies and chaining.
• Add TUI support for port forwarding.
• Add TUI support to change namespaces and contexts.