RF2 Integration

.buildkite/bin/setup

@@ -8,10 +8,10 @@ export TORCH_EXTENSIONS_DIR=~/$BUILDKITE_BUILD_ID/$BUILDKITE_AGENT_NAME/torch_ex
 
 # create env
 conda create -n $BUILDKITE_BUILD_ID -y
-conda env update -n $BUILDKITE_BUILD_ID  --file environments/env-dev.yml --prune
+conda env update -n $BUILDKITE_BUILD_ID  --file environments/linux-cuda/env.yml --prune
 conda activate $BUILDKITE_BUILD_ID
 
 set -x
 
 # install pkg
-pip install .
+pip install .[dev]

Dockerfile

@@ -0,0 +1,24 @@
+# syntax=docker/dockerfile:1
+FROM mambaorg/micromamba:jammy-cuda-12.1.1
+
+USER root
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+    unzip \
+    git \
+    vim \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+USER $MAMBA_USER
+
+ENV PATH "$MAMBA_ROOT_PREFIX/bin:$PATH"
+
+WORKDIR /code
+
+COPY --chown=$MAMBA_USER:$MAMBA_USER environments/linux-cuda /code/environments/linux-cuda
+RUN micromamba install -y -n base -f /code/environments/linux-cuda/env.yml && micromamba clean --all --yes
+RUN pip install -r /code/environments/linux-cuda/requirements-dev-linux-cuda.txt
+COPY --chown=$MAMBA_USER:$MAMBA_USER . /code
+RUN pip install .[dev]

environments/linux-cuda/build_requirements.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+
+pip-compile --verbose --no-emit-index-url --resolver=backtracking -o requirements-linux-cuda.txt ../../requirements.in
+pip-compile --verbose --no-emit-index-url --resolver=backtracking -c requirements-linux-cuda.txt -o requirements-dev-linux-cuda.txt ../../requirements-dev.in

environments/linux-cuda/env.yml

@@ -0,0 +1,10 @@
+name: tmol
+channels:
+  - nvidia/label/cuda-12.1.1
+  - conda-forge
+dependencies:
+  - python=3.11
+  - cuda
+  - pip
+  - pip:
+    - -r requirements-linux-cuda.txt