Added Dockerfile for CI images #195

VeeraRajasekhar · 2025-05-28T22:36:25Z

Description

Added the dockerfile, which can be used to create the ci-artifactory images.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Added a new file docker/Dockerfile

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

docker/Dockerfile

wenchenvincent

Please address the comments.

ipanfilo · 2025-05-29T04:41:28Z

docker/Dockerfile

+RUN apt update \
+    && apt install -y nano wget ninja-build \
+    && apt install -y python3 python3-pip git \
+    && apt install -y sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev


What are those packages (sqlite3 and further) for?

I recall aotriton need sqlite3 long time ago, not sure about now

I have derived this dockerfile from this https://github.com/ROCm/DeepLearningModels/blob/main/docker/pyt_semianalysis_models.ubuntu.amd.Dockerfile,

Will remove these, if these are not required.

@VeeraRajasekhar Those are needed for the semianalysis models. We don't need them for the CI images. We only need to keep those packages necessary for building TE and running CI tests.

docker/Dockerfile

ipanfilo · 2025-05-29T04:43:51Z

docker/Dockerfile

+
+RUN python3 -m pip install --upgrade pip
+RUN pip install ninja cmake setuptools wheel
+RUN pip install uv tabulate


What are those for?

This is derived from https://github.com/ROCm/DeepLearningModels/blob/main/docker/pyt_semianalysis_models.ubuntu.amd.Dockerfile, I have removed these, will test the resultant docker after my changes

docker/Dockerfile

wangye805 · 2025-05-29T05:05:16Z

docker/Dockerfile

+RUN apt update \
+    && apt install -y nano wget ninja-build \
+    && apt install -y python3 python3-pip git \
+    && apt install -y sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev


I recall aotriton need sqlite3 long time ago, not sure about now

docker/Dockerfile

ipanfilo

Why conversations are marked as resolved w/o any actual action?

VeeraRajasekhar · 2025-05-29T15:50:54Z

Why conversations are marked as resolved w/o any actual action?

Some of them, I have resolved, some I have currently resolved in my local, just to keep track I will mark them resolved.

docker/Dockerfile.ci.deps

wenchenvincent · 2025-08-27T14:49:04Z

@VeeraRajasekhar Is this PR still needed?

wenchenvincent · 2026-01-14T03:41:43Z

@VeeraRajasekhar Could you remind me of what we had decided on this PR? It seemed that it is no longer relevant and we should close it.

VeeraRajasekhar · 2026-02-11T03:06:55Z

Hi @ipanfilo, @wangye805

I have updated this PR with latest 7.2 docker file and moved to .github/scripts.

Let me know if I need to add an action to automate docker build and upload to our artifactory?

Thanks.

ipanfilo · 2026-02-11T04:36:13Z

.github/scripts/Dockerfile.ci.deps

+RUN pip install ipython pytest fire pydantic pybind11 ninja pandas
+RUN apt-get update && apt-get install -y vim
+
+ARG PYTORCH_ROCM_ARCH=gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx950;gfx1151


Why is this arg needed?

Not needed, when I manually installed torch I mentioned this, I can remove this now.

ipanfilo · 2026-02-11T04:36:51Z

.github/scripts/Dockerfile.ci.deps

@@ -0,0 +1,27 @@
+# TE CI Dockerfile


Put copyright here

ipanfilo · 2026-02-11T04:37:14Z

.github/scripts/Dockerfile.ci.deps

+
+ARG PYTORCH_ROCM_ARCH=gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx950;gfx1151
+
+# Install flash-attention v2.8.1


Better make FA version arg

ipanfilo · 2026-02-11T04:54:07Z

.github/scripts/Dockerfile.ci.deps

+
+
+RUN pip install \
+    https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/jax_rocm7_pjrt-0.8.0%2Brocm7.2.0-py3-none-manylinux_2_28_x86_64.whl \


Since base image is ARG, JAX wheels should also be built based on ARG cause their ROCm and Python versions should match. So in general there should be ARGs for ROCm version, base image, JAX version, JAX wheels

VeeraRajasekhar · 2026-02-12T01:54:26Z

I had to force push to include new FA 2.8.3 support commit and my changes for 7.2 support to run the CI.

Thanks.

VeeraRajasekhar requested review from ipanfilo, wangye805 and wenchenvincent May 28, 2025 22:37

wenchenvincent reviewed May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

wenchenvincent reviewed May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

wenchenvincent reviewed May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

wenchenvincent requested changes May 29, 2025

View reviewed changes

ipanfilo reviewed May 29, 2025

View reviewed changes

wangye805 requested changes May 29, 2025

View reviewed changes

ipanfilo requested changes May 29, 2025

View reviewed changes

wangye805 reviewed May 29, 2025

View reviewed changes

docker/Dockerfile.ci.deps Outdated Show resolved Hide resolved

docker/Dockerfile.ci.deps Outdated Show resolved Hide resolved

ipanfilo requested changes May 29, 2025

View reviewed changes

docker/Dockerfile.ci.deps Outdated Show resolved Hide resolved

VeeraRajasekhar force-pushed the dockerfile branch from c4913b2 to e49b365 Compare February 11, 2026 03:04

VeeraRajasekhar self-assigned this Feb 11, 2026

VeeraRajasekhar requested review from ipanfilo and wangye805 February 11, 2026 03:07

ipanfilo reviewed Feb 11, 2026

View reviewed changes

VeeraRajasekhar added 7 commits February 12, 2026 01:35

Added Dockerfile for CI images

30833e7

Addressed reviews

342eb9a

Addressed reviews

7ba78c3

Addressed flash attention commit to tag

4fff8ff

Moved Dockerfile.ci.deps to .github/scripts, Updated dockerfile

41c9e42

Addressed comments

d1ecdd3

Adjusted tolerance for gfx942 for ROCm7.2 support

bdc75a2

VeeraRajasekhar force-pushed the dockerfile branch from e49b365 to bdc75a2 Compare February 12, 2026 01:53

Updated CI docker image

dde1a02


		ARG PYTORCH_ROCM_ARCH=gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx950;gfx1151

		# Install flash-attention v2.8.1



		RUN pip install \
		https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/jax_rocm7_pjrt-0.8.0%2Brocm7.2.0-py3-none-manylinux_2_28_x86_64.whl \

Added Dockerfile for CI images #195

Are you sure you want to change the base?

Added Dockerfile for CI images #195

Uh oh!

Conversation

VeeraRajasekhar commented May 28, 2025

Description

Type of change

Changes

Checklist:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wenchenvincent left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ipanfilo left a comment

Choose a reason for hiding this comment

Uh oh!

VeeraRajasekhar commented May 29, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wenchenvincent commented Aug 27, 2025

Uh oh!

wenchenvincent commented Jan 14, 2026

Uh oh!

VeeraRajasekhar commented Feb 11, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VeeraRajasekhar commented Feb 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants