fix imports & pytest

.dockerignore

@@ -6,6 +6,7 @@ build/
 *.pyc
 *.pyo
 *.pyd
+.pytest_cache/
 
 # Git
 **/.git

.github/workflows/codeql.yml

@@ -49,7 +49,7 @@ jobs:
     - name: Build
       run: |
         mkdir build && cd build
-        cmake -DCMAKE_BUILD_TYPE=Debug -DARK_BUILD_PYTHON=ON -DARK_BYPASS_GPU_CHECK=ON -DARK_USE_CUDA=ON -DARK_BUILD_TESTS=OFF ..
+        cmake -DCMAKE_BUILD_TYPE=Debug -DARK_BYPASS_GPU_CHECK=ON -DARK_USE_CUDA=ON -DARK_BUILD_TESTS=OFF ..
         make build ark_py
 
     - name: Perform CodeQL Analysis
@@ -95,7 +95,7 @@ jobs:
     - name: Build
       run: |
         mkdir build && cd build
-        CXX=/opt/rocm/bin/hipcc cmake -DCMAKE_BUILD_TYPE=Debug -DARK_BUILD_PYTHON=ON -DARK_BYPASS_GPU_CHECK=ON -DARK_USE_ROCM=ON -DARK_BUILD_TESTS=OFF ..
+        CXX=/opt/rocm/bin/hipcc cmake -DCMAKE_BUILD_TYPE=Debug -DARK_BYPASS_GPU_CHECK=ON -DARK_USE_ROCM=ON -DARK_BUILD_TESTS=OFF ..
         make -j build ark_py
 
     - name: Perform CodeQL Analysis

.github/workflows/ut-cuda.yml

@@ -44,7 +44,7 @@ jobs:
       - name: Build
         run: |
           mkdir build && cd build
-          cmake -DCMAKE_BUILD_TYPE=Debug -DARK_BUILD_PYTHON=ON ..
+          cmake -DCMAKE_BUILD_TYPE=Debug ..
           make -j ut ark_py
 
       - name: Run C++ UT
@@ -71,7 +71,11 @@ jobs:
       - name: Run Python UT
         run: |
           cd build
-          ARK_ROOT=$PWD pytest --cov=../python/ark --cov-report lcov:py_coverage.info --verbose ../python/unittest/test.py
+          PYTHONPATH=$PWD/python ARK_ROOT=$PWD python3 -m pytest \
+              --cov=../python/ark \
+              --cov-report lcov:py_coverage.info \
+              --verbose \
+              ../python/unittest/test.py
 
       - name: Report Coverage
         env:

python/CMakeLists.txt

@@ -18,7 +18,10 @@ FetchContent_MakeAvailable(pybind11)
 
 file(GLOB_RECURSE BIND_SOURCES CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)
 pybind11_add_module(ark_py ${BIND_SOURCES})
-set_target_properties(ark_py PROPERTIES OUTPUT_NAME _ark_core)
+set_target_properties(ark_py PROPERTIES OUTPUT_NAME _ark_core LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/ark)
+add_custom_command(TARGET ark_py POST_BUILD
+    COMMAND ${CMAKE_COMMAND} -E copy_directory ${CMAKE_CURRENT_SOURCE_DIR}/ark ${CMAKE_CURRENT_BINARY_DIR}/ark
+)
 target_link_libraries(ark_py PRIVATE ark_static)
 target_include_directories(ark_py SYSTEM PRIVATE ${DLPACK_INCLUDE_DIRS})
 target_include_directories(ark_py PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../ark)

python/ark/__init__.py

@@ -1,15 +1,12 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import sys
 import os
 
 if os.environ.get("ARK_ROOT", None) is None:
     os.environ["ARK_ROOT"] = os.path.abspath(os.path.dirname(__file__))
 
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-import _ark_core
+from . import _ark_core
 from .model import Model
 
 
@@ -21,11 +18,6 @@ def version():
     return __version__
 
 
-def srand(seed):
-    """Sets the seed for random number generation."""
-    _ark_core.srand(seed)
-
-
 def set_rank(rank):
     """Sets the rank of the current process."""
     Model.set_rank(rank)

python/ark/data_type.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 import numpy
-import _ark_core
+from . import _ark_core
 
 try:
     import torch

python/ark/error.py

@@ -1,14 +1,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from _ark_core import _BaseError as BaseError
-from _ark_core import _InternalError as InternalError
-from _ark_core import _InvalidUsageError as InvalidUsageError
-from _ark_core import _ModelError as ModelError
-from _ark_core import _PlanError as PlanError
-from _ark_core import _UnsupportedError as UnsupportedError
-from _ark_core import _SystemError as SystemError
-from _ark_core import _GpuError as GpuError
+from ._ark_core import _BaseError as BaseError
+from ._ark_core import _InternalError as InternalError
+from ._ark_core import _InvalidUsageError as InvalidUsageError
+from ._ark_core import _ModelError as ModelError
+from ._ark_core import _PlanError as PlanError
+from ._ark_core import _UnsupportedError as UnsupportedError
+from ._ark_core import _SystemError as SystemError
+from ._ark_core import _GpuError as GpuError
 
 __all__ = [
     "BaseError",

python/ark/init.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import _ark_core
+from . import _ark_core
 from .model import Model
 from .runtime import _RuntimeState
 

python/ark/model.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 from typing import NewType
-from _ark_core import _Model
+from ._ark_core import _Model
 
 _ModelState = NewType("_ModelState", None)
 

python/ark/planner.py

@@ -5,7 +5,7 @@
 import json
 from typing import Callable, Dict, List, Any
 
-from _ark_core import _Planner, _PlannerContext
+from ._ark_core import _Planner, _PlannerContext
 from .model import Model
 
 

python/ark/runtime.py

@@ -4,7 +4,7 @@
 import logging
 from enum import Enum
 
-from _ark_core import _Executor
+from ._ark_core import _Executor
 from .planner import Planner, Plan
 
 

python/ark/tensor.py

@@ -4,7 +4,7 @@
 import numpy as np
 from typing import Callable, List, Union, Type
 
-from _ark_core import _Dims, _Tensor, _NullTensor
+from ._ark_core import _Dims, _Tensor, _NullTensor
 from .data_type import DataType
 from .runtime import Runtime
 from .model import Model

python/unittest/test.py

@@ -1,12 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import sys
-import os
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + "/..")
-sys.path.insert(0, os.environ.get("ARK_ROOT", ".") + "/python")
-
 from test_error import *
 from test_model import *
 from test_runtime import *

python/unittest/test_error.py

@@ -1,11 +1,11 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import ark
+from unittest_common import ark, pytest_ark
 
 
+@pytest_ark()
 def test_error():
-    ark.init()
     try:
         ark.tensor([0])
     except ark.BaseError as e:

python/unittest/test_model.py

@@ -1,13 +1,12 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import ark
+from unittest_common import ark, pytest_ark
 import json
 
 
+@pytest_ark()
 def test_model():
-    ark.init()
-
     input_tensor = ark.tensor([64, 64], ark.fp16)
     other_tensor = ark.tensor([64, 64], ark.fp16)
     ark.add(input_tensor, other_tensor)

python/unittest/test_runtime.py

@@ -1,12 +1,12 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import ark
+from unittest_common import ark, pytest_ark
 import numpy as np
 
 
+@pytest_ark()
 def test_runtime_relaunch():
-    ark.init()
     with ark.Runtime.get_runtime() as rt:
         assert rt.launched() == False
         rt.launch()
@@ -18,8 +18,8 @@ def test_runtime_relaunch():
         assert rt.launched() == True
 
 
+@pytest_ark()
 def test_runtime_init():
-    ark.init()
     M, N = 64, 64
     input_tensor = ark.tensor([M, N], ark.fp16)
     other_tensor = ark.tensor([M, N], ark.fp16)
@@ -51,8 +51,8 @@ def test_runtime_init():
     runtime.reset()
 
 
+@pytest_ark()
 def test_runtime_reuse_plans():
-    ark.init()
     M, N = 64, 64
     input_tensor = ark.tensor([M, N], ark.fp16)
     other_tensor = ark.tensor([M, N], ark.fp16)

python/unittest/test_tensor.py

@@ -1,8 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from unittest_common import pytest_ark
-import ark
+from unittest_common import ark, pytest_ark
 
 
 @pytest_ark(need_torch=True)