[Colossal-Inference] (v0.1.0) Merge pull request #5739 from hpcaitech…

…/feature/colossal-infer [Inference] Merge feature/colossal-infer
hpcaitech · May 22, 2024 · df67476 · df67476
2 parents 22ce873 + 498f42c
commit df67476
Show file tree

Hide file tree

Showing 226 changed files with 20,482 additions and 9,722 deletions.
diff --git a/.github/workflows/build_on_pr.yml b/.github/workflows/build_on_pr.yml
@@ -91,7 +91,7 @@ jobs:
     container:
       image: hpcaitech/pytorch-cuda:2.1.0-12.1.0
       options: --gpus all --rm -v /dev/shm -v /data/scratch/llama-tiny:/data/scratch/llama-tiny
-    timeout-minutes: 60
+    timeout-minutes: 90
     defaults:
       run:
         shell: bash
@@ -140,7 +140,7 @@ jobs:
 
       - name: Install Colossal-AI
         run: |
-          pip install -v -e .
+          BUILD_EXT=1 pip install -v -e .
           pip install -r requirements/requirements-test.txt
 
       - name: Store Colossal-AI Cache

diff --git a/.github/workflows/doc_test_on_pr.yml b/.github/workflows/doc_test_on_pr.yml
@@ -58,7 +58,7 @@ jobs:
     container:
       image: hpcaitech/pytorch-cuda:2.1.0-12.1.0
       options: --gpus all --rm
-    timeout-minutes: 20
+    timeout-minutes: 30
     defaults:
       run:
         shell: bash

diff --git a/.github/workflows/example_check_on_pr.yml b/.github/workflows/example_check_on_pr.yml
@@ -8,6 +8,7 @@ on:
     # any change in the examples folder will trigger check for the corresponding example.
     paths:
       - "examples/**"
+      - "!examples/**.md"
 
 jobs:
   # This is for changed example files detect and output a matrix containing all the corresponding directory name.
@@ -19,6 +20,7 @@ jobs:
     outputs:
       matrix: ${{ steps.setup-matrix.outputs.matrix }}
       anyChanged: ${{ steps.setup-matrix.outputs.anyChanged }}
+      anyExtensionFileChanged: ${{ steps.find-extension-change.outputs.any_changed }}
     name: Detect changed example files
     concurrency:
       group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-detect-change
@@ -37,6 +39,16 @@ jobs:
           echo $commonCommit
           echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT
 
+      - name: Find the changed extension-related files
+        id: find-extension-change
+        uses: tj-actions/changed-files@v35
+        with:
+          base_sha: ${{ steps.locate-base-sha.outputs.baseSHA }}
+          files: |
+            op_builder/**
+            colossalai/kernel/**
+            setup.py
+
       - name: Get all changed example files
         id: changed-files
         uses: tj-actions/changed-files@v35
@@ -79,17 +91,28 @@ jobs:
     container:
       image: hpcaitech/pytorch-cuda:2.1.0-12.1.0
       options: --gpus all --rm -v /data/scratch/examples-data:/data/ -v /dev/shm
-    timeout-minutes: 20
+    timeout-minutes: 30
     concurrency:
       group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-run-example-${{ matrix.directory }}
       cancel-in-progress: true
     steps:
       - uses: actions/checkout@v3
 
+      - name: Restore Colossal-AI Cache
+        if: needs.detect.outputs.anyExtensionFileChanged != 'true'
+        run: |
+          if [ -d /github/home/cuda_ext_cache ] && [ ! -z "$(ls -A /github/home/cuda_ext_cache/)" ]; then
+            cp -p -r /github/home/cuda_ext_cache/* /__w/ColossalAI/ColossalAI/
+          fi
+
       - name: Install Colossal-AI
         run: |
           BUILD_EXT=1 pip install -v .
 
+      - name: Store Colossal-AI Cache
+        run: |
+          cp -p -r /__w/ColossalAI/ColossalAI/build /github/home/cuda_ext_cache/
+
       - name: Test the example
         run: |
           example_dir=${{ matrix.directory }}

diff --git a/.github/workflows/example_check_on_schedule.yml b/.github/workflows/example_check_on_schedule.yml
@@ -36,7 +36,7 @@ jobs:
     container:
       image: hpcaitech/pytorch-cuda:2.1.0-12.1.0
       options: --gpus all --rm -v /data/scratch/examples-data:/data/ -v /dev/shm
-    timeout-minutes: 10
+    timeout-minutes: 30
     steps:
       - name: 📚 Checkout
         uses: actions/checkout@v3

diff --git a/colossalai/inference/README.md b/colossalai/inference/README.md
diff --git a/colossalai/inference/__init__.py b/colossalai/inference/__init__.py
@@ -1,4 +1,4 @@
-from .engine import InferenceEngine
-from .engine.policies import BloomModelInferPolicy, ChatGLM2InferPolicy, LlamaModelInferPolicy
+from .config import InferenceConfig
+from .core import InferenceEngine
 
-__all__ = ["InferenceEngine", "LlamaModelInferPolicy", "BloomModelInferPolicy", "ChatGLM2InferPolicy"]
+__all__ = ["InferenceConfig", "InferenceEngine"]