Image API

Kubetorch's Image class allows you to define the container image that will be used for your compute environment, and lightly customize it as well without rebuilding the Docker image in full. While in production, we think you probably should just use a single image that contains your code and all dependencies.

Arguments

The arguments are simple:

name (str, optional): Name to assign to the image.
image_id (str, optional): Base Docker image to build from (e.g., "nvcr.io/nvidia/pytorch:23.10-py3").

Methods

However, there are also methods to lightly interact with the base image that might be helpful during interactive deployment, to not need to alter your team Docker image or constantly rebuild for iteration.

.pip_install(reqs: List[str], force: bool = False)
Add a pip install step to the image.
```
image.pip_install(["transformers", "scikit-learn"])
```
.set_env_vars(env_vars: Dict[str, str])
Set environment variables inside the container.
```
image.set_env_vars({"CUDA_VISIBLE_DEVICES": "0"})
```
.sync_package(package: str, force: bool = False)
Sync a local editable package or folder into the container.
```
image.sync_package("my_local_module")
```
.run_bash(command: str, force: bool = False)
Run arbitrary shell commands as a build step.
```
image.run_bash("apt-get update && apt-get install -y libsndfile1")
```
.rsync(source: str, dest: str, contents: bool = False, filter_options: str = None, force: bool = False)
Copy files or directories from local to remote.
```
image.rsync("./data", "/workspace/data", contents=True)
```
.from_docker(image_id: str, docker_secret: Optional[Secret or str] = None)
Specify a prebuilt Docker image to use as the base image, optionally including a pull secret.
```
image.from_docker("myregistry/myimage:latest", docker_secret="my-docker-secret")
```

Example

As example, here is taking a base Docker image and minimally altering it with pip install and setting an env var based on my local machine's HuggingFace Token.

import kubetorch as kt

image = (
    kt.Image(image_id="base-image")
    .pip_install(["vllm", "datasets"])
    .set_env_vars({"HF_TOKEN": os.environ("HF_TOKEN")})
)