Skip to content

Commit

Permalink
docs fix: Updated the userguide in ImageSpec document, added two meth…
Browse files Browse the repository at this point in the history
…ods using parameters and requirements.txt to the subsections of Installing cuda
  • Loading branch information
sharsha315 committed Oct 26, 2024
1 parent 13b3d82 commit 4a1dd8d
Showing 1 changed file with 32 additions and 0 deletions.
32 changes: 32 additions & 0 deletions docs/user_guide/customizing_dependencies/imagespec.md
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,38 @@ def task2() -> int:
## Install CUDA in the image
There are few ways to install CUDA in the image.

### Use cuda and cudnn parameters in `ImageSpec`.
Specifying CUDA and cuDNN versions in ImageSpec.
```python
from flytekit import task, workflow
from flytekit.image_spec import ImageSpec

new_flytekit = "git+https://github.com/flyteorg/flytekit@8c7c43ad310ec74b4e8c06c2b4604b6cb9692f86"
image_spec = ImageSpec(
base_image="ubuntu20.04",
packages=[new_flytekit],
cuda="11.2.2",
cudnn="8",
registry="pingsutw",
apt_packages=["git"],
python_version="3.9",
)
```

### Use requirements.txt File in `ImageSpec`.
Including dependencies from a requirements.txt file using the requirements parameter.
```python
from flytekit import task, workflow, ImageSpec

new_flytekit = "git+https://github.com/flyteorg/flytekit@01f2d63221467face8d35ed0af5308bbb3e66eaa"
image = ImageSpec(
registry="pingsutw",
packages=["pandas", new_flytekit],
apt_packages=["git"],
requirements="requirement.txt"
)
```

### Use Nvidia docker image
CUDA is pre-installed in the Nvidia docker image. You can specify the base image in the `ImageSpec`.
```python
Expand Down

0 comments on commit 4a1dd8d

Please sign in to comment.