Merge pull request #18 from DocsaidLab/feat/ort_io_binding

[A] add io_binding mode for onnxengine with cuda graph on

.github/workflows/.pylintrc

@@ -535,5 +535,5 @@ preferred-modules=
 
 # Exceptions that will emit a warning when being caught. Defaults to
 # "BaseException, Exception".
-overgeneral-exceptions=BaseException,
-                       Exception
+overgeneral-exceptions=builtins.BaseException,
+                       builtins.Exception

.github/workflows/pull_request.yml

@@ -1,57 +1,96 @@
 name: Pull Request
 
 on:
-  workflow_dispatch:
   pull_request:
-    branches:
-      - main
+    branches: [main]
+    paths-ignore:
+      - "**/version.h"
+      - "doc/**"
+      - "**.md"
+
+env:
+  LOCAL_REGISTRY: localhost:5000
+  DOCKERFILE: docker/pr.dockerfile
 
 jobs:
-  test:
-    name: Run Tests
+  get_runner_and_uid:
+    name: Get Runner
     runs-on: [self-hosted, unicorn]
+    steps:
+      - name: Get UID and GID
+        id: uid_gid
+        run: |
+          echo "uid_gid=$(id -u):$(id -g)" >> "$GITHUB_OUTPUT"
+    outputs:
+      runner: ${{ runner.name }}
+      uid: ${{ steps.uid_gid.outputs.uid_gid }}
+
+  build_docker_image:
+    name: Build Docker Image
+    needs: [get_runner_and_uid]
+    runs-on: ${{ needs.get_runner_and_uid.outputs.runner }}
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v4
+
+      - name: Build Docker Image
+        id: docker_build
+        uses: docker/build-push-action@v3
+        with:
+          file: ${{ env.DOCKERFILE }}
+          no-cache: false
+          push: true
+          tags: ${{ env.LOCAL_REGISTRY }}/capybara-container:ci
+
+    outputs:
+      image: ${{ env.LOCAL_REGISTRY }}/capybara-container:ci
+
+  ci:
+    name: CI
+    needs: [get_runner_and_uid, build_docker_image]
+    runs-on: ${{ needs.get_runner_and_uid.outputs.runner }}
     strategy:
       matrix:
         python-version:
           - "3.10"
+    container:
+      image: ${{ needs.build_docker_image.outputs.image }}
+      options: --user ${{ needs.get_runner_and_uid.outputs.uid }} --gpus all
+
     steps:
-      - name: Checkout code
+      - name: Checkout Repository
         uses: actions/checkout@v4
 
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install packages
+      - name: Install Dependencies
         run: |
-          python -m pip install pytest wheel pylint pylint-flask
-          python setup.py bdist_wheel
-          wheel_file=$(echo dist/*.whl)
-          python -m pip install ${wheel_file} --force-reinstall
+          python3 -m pip install pytest wheel pylint pylint-flask pytest-cov typeguard
 
-      - name: Lint with pylint
+      - name: Build and Install Package
         run: |
-          python -m pylint ${{ github.workspace }}/capybara \
-            --rcfile=.github/workflows/.pylintrc \
-            --load-plugins pylint_flask \
+          python3 setup.py bdist_wheel && \
+          wheel_file=$(ls dist/*.whl 2>/dev/null || echo '') && \
+          if [ -z "$wheel_file" ]; then
+            echo 'Error: No wheel file found in dist directory.' && exit 1
+          fi && \
+          python3 -m pip install $wheel_file --force-reinstall
 
-      - name: Test with pytest
+      - name: Lint with Pylint
         run: |
-          python -m pip install pytest pytest-cov typeguard
+          python3 -m pylint capybara \
+          --rcfile=.github/workflows/.pylintrc \
+          --load-plugins pylint_flask
 
-          # Test all
-          python -m pytest tests
-
-          # Report all
-          mkdir -p tests/coverage
-          python -m pytest -x \
-            --junitxml=tests/coverage/cov-jumitxml.xml \
-            --cov=capybara tests | tee tests/coverage/cov.txt
+      - name: Run Tests with Pytest
+        run: |
+          mkdir -p tests/coverage && \
+          python3 -m pytest tests --junitxml=tests/coverage/cov-junitxml.xml \
+          --cov=capybara | tee tests/coverage/cov.txt
 
-      - name: Pytest coverage comment
+      - name: Pytest Coverage Comment
         id: coverageComment
         uses: MishaKav/pytest-coverage-comment@main
         with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
           pytest-coverage-path: tests/coverage/cov.txt
-          junitxml-path: tests/coverage/cov-jumitxml.xml
+          junitxml-path: tests/coverage/cov-junitxml.xml

capybara/onnxengine/__init__.py

@@ -1,4 +1,5 @@
 from .engine import Backend, ONNXEngine
+from .engine_io_binding import ONNXEngineIOBinding
 from .metadata import get_onnx_metadata, write_metadata_into_onnx
 from .tools import get_onnx_input_infos, get_onnx_output_infos
 

capybara/onnxengine/engine.py

@@ -118,43 +118,68 @@ def _get_provider_info(
         return providers, provider_option
 
     def __repr__(self) -> str:
+        import re
+
+        def strip_ansi_codes(text):
+            """Remove ANSI escape codes from a string."""
+            ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
+            return ansi_escape.sub('', text)
+
         def format_nested_dict(dict_data, indent=0):
-            info = ""
-            for k, v in dict_data.items():
-                prefix = "  " * indent
-                if isinstance(v, dict):
-                    info += f"{prefix}{k}:\n" + \
-                        format_nested_dict(v, indent + 1)
-                elif isinstance(v, str) and v.startswith('{') and v.endswith('}'):
+            """Recursively format nested dictionaries with indentation."""
+            info = []
+            prefix = "  " * indent
+            for key, value in dict_data.items():
+                if isinstance(value, dict):
+                    info.append(f"{prefix}{key}:")
+                    info.append(format_nested_dict(value, indent + 1))
+                elif isinstance(value, str) and value.startswith('{') and value.endswith('}'):
                     try:
-                        nested_dict = eval(v)
+                        nested_dict = eval(value)
                         if isinstance(nested_dict, dict):
-                            info += f"{prefix}{k}:\n" + \
-                                format_nested_dict(nested_dict, indent + 1)
+                            info.append(f"{prefix}{key}:")
+                            info.append(format_nested_dict(
+                                nested_dict, indent + 1))
                         else:
-                            info += f"{prefix}{k}: {v}\n"
-                    except:
-                        info += f"{prefix}{k}: {v}\n"
+                            info.append(f"{prefix}{key}: {value}")
+                    except Exception:
+                        info.append(f"{prefix}{key}: {value}")
                 else:
-                    info += f"{prefix}{k}: {v}\n"
-            return info
+                    info.append(f"{prefix}{key}: {value}")
+            return "\n".join(info)
 
         title = 'DOCSAID X ONNXRUNTIME'
+        divider_length = 50
+        divider = f"+{'-' * divider_length}+"
         styled_title = colored.stylize(
             title, [colored.fg('blue'), colored.attr('bold')])
-        divider_length = 50
-        title_length = len(title)
-        left_padding = (divider_length - title_length) // 2
-        right_padding = divider_length - title_length - left_padding
 
-        path = f'Model Path: {self.model_path}'
+        def center_text(text, width):
+            """Center text within a fixed width, handling ANSI escape codes."""
+            plain_text = strip_ansi_codes(text)
+            text_length = len(plain_text)
+            left_padding = (width - text_length) // 2
+            right_padding = width - text_length - left_padding
+            return f"|{' ' * left_padding}{text}{' ' * right_padding}|"
+
+        path = f"Model Path: {self.model_path}"
         input_info = format_nested_dict(self.input_infos)
         output_info = format_nested_dict(self.output_infos)
         metadata = format_nested_dict(self.metadata)
-        providers = f'Provider: {", ".join(self.providers)}'
+        providers = f"Provider: {', '.join(self.providers)}"
         provider_options = format_nested_dict(self.provider_options)
 
-        divider = colored.stylize(
-            f"+{'-' * divider_length}+", [colored.fg('blue'), colored.attr('bold')])
-        infos = f'\n\n{divider}\n|{" " * left_padding}{styled_title}{" " * right_padding}|\n{divider}\n\n{path}\n\n{input_info}\n{output_info}\n\n{metadata}\n\n{providers}\n\n{provider_options}\n{divider}'
-        return infos
+        sections = [
+            divider,
+            center_text(styled_title, divider_length),
+            divider,
+            path,
+            input_info,
+            output_info,
+            metadata,
+            providers,
+            provider_options,
+            divider,
+        ]
+
+        return "\n\n".join(sections)