From d5df1e9230f88671c258f67035f2824722f2868c Mon Sep 17 00:00:00 2001 From: ark-dev Date: Tue, 9 Jun 2026 08:38:47 +0000 Subject: [PATCH 1/5] ark-dev: Implement Python API rework: Tensor.eval, set_model/current_model/use_model, torch interop; apply refs/p6-python-api.patch. --- python/ark/__init__.py | 1 + python/ark/data_type.py | 1 + python/ark/model.py | 45 +++++++++++- python/ark/ops.py | 67 ++++++++++++++++++ python/ark/planner.py | 6 +- python/ark/tensor.py | 40 +++++++++++ python/unittest/common.py | 2 + python/unittest/test.py | 2 + python/unittest/test_conversion.py | 110 +++++++++++++++++++++++++++++ python/unittest/test_eval.py | 89 +++++++++++++++++++++++ 10 files changed, 360 insertions(+), 3 deletions(-) create mode 100644 python/unittest/test_conversion.py create mode 100644 python/unittest/test_eval.py diff --git a/python/ark/__init__.py b/python/ark/__init__.py index 61ff98a3..653d1950 100644 --- a/python/ark/__init__.py +++ b/python/ark/__init__.py @@ -8,6 +8,7 @@ from .core import version from .model import Model +from .model import set_model, current_model, use_model __version__ = version() diff --git a/python/ark/data_type.py b/python/ark/data_type.py index fa2b2c06..0caac294 100644 --- a/python/ark/data_type.py +++ b/python/ark/data_type.py @@ -10,6 +10,7 @@ "DataType", "fp16", "fp32", + "bf16", "int32", "uint32", "int8", diff --git a/python/ark/model.py b/python/ark/model.py index 87c7279a..b6cff59c 100644 --- a/python/ark/model.py +++ b/python/ark/model.py @@ -1,11 +1,12 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from contextlib import contextmanager from typing import NewType from . import log from .core import CoreModel -__all__ = ["Model"] +__all__ = ["Model", "set_model", "current_model", "use_model"] ModelState = NewType("ModelState", None) @@ -114,3 +115,45 @@ class ModelState: rank: int = 0 world_size: int = 1 device_id: int = 0 + + +def set_model(model: Model) -> None: + """ + Set the current model. + + Args: + model: The model to set as current. + """ + if not isinstance(model, Model): + raise log.InvalidUsageError("model must be a Model instance") + ModelState.model = model + + +def current_model() -> Model: + """ + Return the current model, creating one if none exists. + """ + return Model.get_model() + + +@contextmanager +def use_model(model: Model = None): + """ + Context manager that sets *model* as the current model on entry and + restores the previous model on exit. If *model* is ``None``, a fresh + model is created. + + Args: + model: The model to use inside the context. If ``None``, a new + ``Model`` is created with the current rank and world size. + """ + prev = ModelState.model + if model is None: + model = Model(ModelState.rank, ModelState.world_size) + elif not isinstance(model, Model): + raise log.InvalidUsageError("model must be a Model instance") + ModelState.model = model + try: + yield model + finally: + ModelState.model = prev diff --git a/python/ark/ops.py b/python/ark/ops.py index 2b7e387f..1db16ab5 100644 --- a/python/ark/ops.py +++ b/python/ark/ops.py @@ -9,6 +9,17 @@ from .model import Model from . import log + +def _ensure_ark(t): + """ + If *t* is a ``torch.Tensor``, convert it to an ARK ``Tensor`` via + ``Tensor.from_torch()``. Otherwise return *t* unchanged. + """ + if not _no_torch and isinstance(t, torch.Tensor): + return Tensor.from_torch(t) + return t + + __all__ = [ "tensor", "parameter", @@ -59,6 +70,9 @@ def add( name: str = "add", ) -> Union[Tensor, float]: """ """ + input = _ensure_ark(input) + other = _ensure_ark(other) + output = _ensure_ark(output) if isinstance(input, Tensor) and isinstance(other, Tensor): a = input._tensor b = other._tensor @@ -86,6 +100,8 @@ def cast( name: str = "cast", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -111,6 +127,8 @@ def copy( name: str = "copy", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor if isinstance(input, Tensor): @@ -125,6 +143,9 @@ def div( name: str = "div", ) -> Tensor: """ """ + input = _ensure_ark(input) + other = _ensure_ark(other) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor if isinstance(other, Tensor): @@ -139,6 +160,9 @@ def embedding( name: str = "embedding", ) -> Tensor: """ """ + input = _ensure_ark(input) + weight = _ensure_ark(weight) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -152,6 +176,8 @@ def exp( name: str = "exp", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor(Model.get_model().exp(input._tensor, output, name)) @@ -163,6 +189,8 @@ def gelu( name: str = "gelu", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor(Model.get_model().gelu(input._tensor, output, name)) @@ -172,6 +200,8 @@ def identity( input: Tensor, deps: List[Tensor] = [], name: str = "identity" ) -> Tensor: """ """ + input = _ensure_ark(input) + deps = [_ensure_ark(d) for d in deps] dep_tensors = [] for dep in deps: if not isinstance(dep, Tensor): @@ -189,6 +219,9 @@ def matmul( name: str = "matmul", ) -> Tensor: """ """ + input = _ensure_ark(input) + other = _ensure_ark(other) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -210,6 +243,9 @@ def mul( name: str = "mul", ) -> Tensor: """ """ + input = _ensure_ark(input) + other = _ensure_ark(other) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor if isinstance(other, Tensor): @@ -219,6 +255,7 @@ def mul( def noop(input: Tensor, name: str = "noop"): """ """ + input = _ensure_ark(input) Model.get_model().noop(input._tensor, name) @@ -259,6 +296,8 @@ def reduce_max( name: str = "reduce_max", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -276,6 +315,8 @@ def reduce_mean( name: str = "reduce_mean", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -293,6 +334,8 @@ def reduce_sum( name: str = "reduce_sum", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -308,6 +351,8 @@ def relu( name: str = "relu", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor(Model.get_model().relu(input._tensor, output, name)) @@ -332,6 +377,7 @@ def reshape( # tensors shape is [128, 64] tensor = ark.reshape(tensor, [2, 64, 64]) """ + input = _ensure_ark(input) if not is_list_or_tuple(shape): raise log.InvalidUsageError( "shape should be a list or tuple of integers" @@ -353,6 +399,9 @@ def rope( name: str = "rope", ) -> Tensor: """ """ + input = _ensure_ark(input) + other = _ensure_ark(other) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -366,6 +415,8 @@ def rsqrt( name: str = "rsqrt", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor(Model.get_model().rsqrt(input._tensor, output, name)) @@ -375,6 +426,7 @@ def sharding( input: Tensor, axis: int, dim_per_shard: int, name: str = "sharding" ) -> List[Tensor]: """ """ + input = _ensure_ark(input) _tensor_list = Model.get_model().sharding( input._tensor, axis, dim_per_shard, name ) @@ -387,6 +439,8 @@ def sigmoid( name: str = "sigmoid", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor(Model.get_model().sigmoid(input._tensor, output, name)) @@ -398,6 +452,8 @@ def sqrt( name: str = "sqrt", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor(Model.get_model().sqrt(input._tensor, output, name)) @@ -410,6 +466,9 @@ def sub( name: str = "sub", ) -> Tensor: """ """ + input = _ensure_ark(input) + other = _ensure_ark(other) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor if isinstance(other, Tensor): @@ -441,6 +500,8 @@ def transpose( name: str = "transpose", ) -> Tensor: """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor if not is_list_or_tuple(perm): @@ -466,6 +527,8 @@ def send( name: str = "send", ): """ """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor return Tensor( @@ -478,6 +541,7 @@ def send_done( name: str = "send_done", ): """ """ + input = _ensure_ark(input) return Tensor(Model.get_model().send_done(input._tensor, name)) @@ -488,6 +552,7 @@ def recv( name: str = "recv", ): """ """ + output = _ensure_ark(output) return Tensor( Model.get_model().recv(output._tensor, remote_rank, tag, name) ) @@ -590,6 +655,8 @@ def all_reduce( Returns: Tensor: The reduced tensor. """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor _tensor = Model.get_model().all_reduce( diff --git a/python/ark/planner.py b/python/ark/planner.py index 79b0fb7e..870b8a59 100644 --- a/python/ark/planner.py +++ b/python/ark/planner.py @@ -232,8 +232,10 @@ def __exit__(self, exc_type, exc_value, exc_tb): class Planner(CorePlanner): - def __init__(self, device_id: int = 0): - compressed = Model.get_model().compress() + def __init__(self, device_id: int = 0, model: Model = None): + if model is None: + model = Model.get_model() + compressed = model.compress() super().__init__(compressed, device_id) def install_config_rule(self, rule: Callable[[str, str], str]): diff --git a/python/ark/tensor.py b/python/ark/tensor.py index 216318b2..b75796e6 100644 --- a/python/ark/tensor.py +++ b/python/ark/tensor.py @@ -27,6 +27,7 @@ def __init__( _tensor: CoreTensor, initializer: Initializer = None, requires_grad: bool = False, + _model: "Model" = None, ): """ Initializes a new instance of the Tensor class. @@ -34,10 +35,18 @@ def __init__( _tensor (core.CoreTensor): The underlying _Tensor object. initializer (Initializer): The initializer for the Tensor. requires_grad (bool): Whether the tensor requires gradient. Defaults to True. + _model (Model): The model this tensor belongs to. If None, + the current model is used. """ self._tensor: CoreTensor = _tensor self.initializer: Initializer = initializer self.requires_grad: bool = requires_grad + # Note: Model.get_model() creates a default model if none exists. + # The model is captured eagerly so that eval() uses the model that + # was active when the tensor was created. + self._model: "Model" = ( + _model if _model is not None else Model.get_model() + ) def __hash__(self): return self._tensor.id() @@ -332,6 +341,37 @@ def initialize(self) -> "Tensor": self.copy(data) return self + def eval(self, stream: int = 0) -> torch.Tensor: + """ + Compile the graph, run it, and return the result as a torch tensor. + + A fresh ``Planner`` and ``Runtime`` are created each call. + The build system's file-level compile cache avoids redundant + ``nvcc`` invocations across calls, but the ``Runtime`` object + is not reused. + + Args: + stream: CUDA stream ordinal for the execution. Defaults to 0. + + Returns: + torch.Tensor: The result tensor on the same device. + """ + if _no_torch: + raise log.SystemError("torch is not available") + from .planner import Planner + from .runtime import Runtime + + device_id = Model.get_device_id() + planner = Planner(device_id=device_id, model=self._model) + plan = planner.plan() + + with Runtime() as rt: + rt.launch( + plan=plan, device_id=device_id, stream=stream, loop_mode=False + ) + rt.run() + return self.to_torch() + class Parameter(Tensor): """ diff --git a/python/unittest/common.py b/python/unittest/common.py index 0c385e89..a0ac3ace 100644 --- a/python/unittest/common.py +++ b/python/unittest/common.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import functools import pytest import ark @@ -19,6 +20,7 @@ def decorator(test_func): test_func ) + @functools.wraps(test_func) def wrapper(*args, **kwargs): ark.init() test_func(*args, **kwargs) diff --git a/python/unittest/test.py b/python/unittest/test.py index 822fb1f7..e93eae05 100644 --- a/python/unittest/test.py +++ b/python/unittest/test.py @@ -10,3 +10,5 @@ from test_profiler import * from test_runtime import * from test_tensor import * +from test_conversion import * +from test_eval import * diff --git a/python/unittest/test_conversion.py b/python/unittest/test_conversion.py new file mode 100644 index 00000000..f4430fee --- /dev/null +++ b/python/unittest/test_conversion.py @@ -0,0 +1,110 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from common import ark, pytest_ark +import numpy as np +import pytest + + +@pytest_ark(need_torch=True) +def test_conversion_torch_to_ark_fp32(): + """Test converting a torch fp32 tensor to ARK and back.""" + import torch + + torch_data = torch.arange(64, dtype=torch.float32, device="cuda:0") + ark_tensor = ark.Tensor.from_torch(torch_data) + + assert ark_tensor.shape() == [64] + assert ark_tensor.dtype() == ark.fp32 + + out = ark.add(ark_tensor, 1.0) + + with ark.Runtime() as rt: + rt.launch() + rt.run() + result = out.to_numpy() + + expected = torch_data.cpu().numpy() + 1.0 + assert np.allclose(result, expected) + + +@pytest_ark(need_torch=True) +def test_conversion_torch_to_ark_fp16(): + """Test converting a torch fp16 tensor to ARK and back.""" + import torch + + torch_data = torch.ones(128, dtype=torch.float16, device="cuda:0") * 3.0 + ark_tensor = ark.Tensor.from_torch(torch_data) + + assert ark_tensor.shape() == [128] + assert ark_tensor.dtype() == ark.fp16 + + out = ark.mul(ark_tensor, 2.0) + + with ark.Runtime() as rt: + rt.launch() + rt.run() + result = out.to_numpy() + + expected = torch_data.cpu().numpy() * 2.0 + assert np.allclose(result, expected, atol=1e-2) + + +@pytest_ark(need_torch=True) +def test_conversion_ark_to_torch(): + """Test converting an ARK tensor result to torch.""" + import torch + + torch_input = torch.ones(64, dtype=torch.float32, device="cuda:0") * 5.0 + ark_input = ark.Tensor.from_torch(torch_input) + out = ark.add(ark_input, 10.0) + + with ark.Runtime() as rt: + rt.launch() + rt.run() + torch_result = out.to_torch() + + assert isinstance(torch_result, torch.Tensor) + expected = torch.ones(64, dtype=torch.float32, device="cuda:0") * 15.0 + assert torch.allclose(torch_result, expected) + + +@pytest_ark(need_torch=True) +def test_conversion_ensure_ark_passthrough(): + """Test that _ensure_ark passes through ARK tensors unchanged.""" + from ark.ops import _ensure_ark + + t = ark.tensor([64], ark.fp32) + assert _ensure_ark(t) is t + + +@pytest_ark(need_torch=True) +def test_conversion_ensure_ark_converts_torch(): + """Test that _ensure_ark converts torch tensors to ARK tensors.""" + import torch + from ark.ops import _ensure_ark + + torch_t = torch.ones(64, dtype=torch.float32, device="cuda:0") + ark_t = _ensure_ark(torch_t) + + assert isinstance(ark_t, ark.Tensor) + assert ark_t.shape() == [64] + + +@pytest_ark(need_torch=True) +def test_conversion_ops_accept_torch(): + """Test that ops accept torch tensors directly via _ensure_ark.""" + import torch + + a = torch.ones(64, dtype=torch.float32, device="cuda:0") * 2.0 + b = torch.ones(64, dtype=torch.float32, device="cuda:0") * 3.0 + + out = ark.add(a, b) + + with ark.Runtime() as rt: + rt.launch() + rt.run() + result = out.to_numpy() + + expected = np.ones(64, dtype=np.float32) * 5.0 + assert np.allclose(result, expected) diff --git a/python/unittest/test_eval.py b/python/unittest/test_eval.py new file mode 100644 index 00000000..bd4e952f --- /dev/null +++ b/python/unittest/test_eval.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import pytest + +torch = pytest.importorskip("torch") + +import numpy as np +from common import ark, pytest_ark + + +@pytest_ark(need_torch=True) +def test_eval_basic(): + """Test basic Tensor.eval() — compile, run, return torch tensor.""" + x = torch.ones(64, dtype=torch.float32, device="cuda:0") * 3.0 + out = ark.add(x, 2.0) + + result = out.eval() + + assert isinstance(result, torch.Tensor) + expected = torch.ones(64, dtype=torch.float32, device="cuda:0") * 5.0 + assert torch.allclose(result, expected) + + +@pytest_ark(need_torch=True) +def test_eval_chain(): + """Test eval on a chained computation.""" + x = torch.ones(64, dtype=torch.float16, device="cuda:0") * 4.0 + y = ark.mul(x, 2.0) + z = ark.add(y, 1.0) + + result = z.eval() + + assert isinstance(result, torch.Tensor) + expected = torch.ones(64, dtype=torch.float16, device="cuda:0") * 9.0 + assert torch.allclose(result.cpu(), expected.cpu(), atol=1e-2) + + +@pytest_ark(need_torch=True) +def test_eval_relu(): + """Test eval with relu op.""" + x = torch.tensor( + [-1.0, 0.0, 1.0, 2.0], dtype=torch.float32, device="cuda:0" + ) + out = ark.relu(x) + + result = out.eval() + + expected = torch.tensor( + [0.0, 0.0, 1.0, 2.0], dtype=torch.float32, device="cuda:0" + ) + assert torch.allclose(result, expected) + + +@pytest_ark(need_torch=True) +def test_eval_matmul(): + """Test eval with matmul.""" + a = torch.ones(4, 64, dtype=torch.float16, device="cuda:0") + b = torch.ones(64, 8, dtype=torch.float16, device="cuda:0") + out = ark.matmul(a, b) + + result = out.eval() + + assert result.shape == (4, 8) + expected = torch.full((4, 8), 64.0, dtype=torch.float16, device="cuda:0") + assert torch.allclose(result, expected, atol=1e-1) + + +@pytest_ark(need_torch=True) +def test_eval_returns_fresh_runtime(): + """Test that eval creates a fresh Runtime each call (no state leak).""" + x = torch.ones(64, dtype=torch.float32, device="cuda:0") * 2.0 + out = ark.add(x, 3.0) + + result1 = out.eval() + assert torch.allclose( + result1, + torch.ones(64, dtype=torch.float32, device="cuda:0") * 5.0, + ) + + # Second eval on the same tensor should also work + ark.init() + x2 = torch.ones(64, dtype=torch.float32, device="cuda:0") * 10.0 + out2 = ark.add(x2, 1.0) + result2 = out2.eval() + assert torch.allclose( + result2, + torch.ones(64, dtype=torch.float32, device="cuda:0") * 11.0, + ) From bb54ed2d29becc151b6de7c0d0165e8104c8b1e9 Mon Sep 17 00:00:00 2001 From: ark-dev Date: Tue, 9 Jun 2026 17:41:59 +0000 Subject: [PATCH 2/5] ark-dev: Implement P6 Python API rework: add Tensor.eval(), set_model()/current_model()/use_model() context helpers, and torch interop via _ensure_ark(); apply the reference patch refs/p6-python-api.patch and open a PR against main. --- python/ark/ops.py | 20 +++++----- python/ark/tensor.py | 3 +- python/unittest/test_conversion.py | 1 - python/unittest/test_eval.py | 36 ++++++++++------- python/unittest/test_model.py | 63 ++++++++++++++++++++++++++++++ 5 files changed, 96 insertions(+), 27 deletions(-) diff --git a/python/ark/ops.py b/python/ark/ops.py index 1db16ab5..e0f0802b 100644 --- a/python/ark/ops.py +++ b/python/ark/ops.py @@ -10,16 +10,6 @@ from . import log -def _ensure_ark(t): - """ - If *t* is a ``torch.Tensor``, convert it to an ARK ``Tensor`` via - ``Tensor.from_torch()``. Otherwise return *t* unchanged. - """ - if not _no_torch and isinstance(t, torch.Tensor): - return Tensor.from_torch(t) - return t - - __all__ = [ "tensor", "parameter", @@ -59,6 +49,16 @@ def _ensure_ark(t): ] +def _ensure_ark(t): + """ + If *t* is a ``torch.Tensor``, convert it to an ARK ``Tensor`` via + ``Tensor.from_torch()``. Otherwise return *t* unchanged. + """ + if not _no_torch and isinstance(t, torch.Tensor): + return Tensor.from_torch(t) + return t + + def is_list_or_tuple(obj): return isinstance(obj, list) or isinstance(obj, tuple) diff --git a/python/ark/tensor.py b/python/ark/tensor.py index b75796e6..ec9e87db 100644 --- a/python/ark/tensor.py +++ b/python/ark/tensor.py @@ -47,6 +47,7 @@ def __init__( self._model: "Model" = ( _model if _model is not None else Model.get_model() ) + self._device_id: int = Model.get_device_id() def __hash__(self): return self._tensor.id() @@ -361,7 +362,7 @@ def eval(self, stream: int = 0) -> torch.Tensor: from .planner import Planner from .runtime import Runtime - device_id = Model.get_device_id() + device_id = self._device_id planner = Planner(device_id=device_id, model=self._model) plan = planner.plan() diff --git a/python/unittest/test_conversion.py b/python/unittest/test_conversion.py index f4430fee..e542ad27 100644 --- a/python/unittest/test_conversion.py +++ b/python/unittest/test_conversion.py @@ -3,7 +3,6 @@ from common import ark, pytest_ark import numpy as np -import pytest @pytest_ark(need_torch=True) diff --git a/python/unittest/test_eval.py b/python/unittest/test_eval.py index bd4e952f..9ccff0e7 100644 --- a/python/unittest/test_eval.py +++ b/python/unittest/test_eval.py @@ -1,10 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import pytest - -torch = pytest.importorskip("torch") - import numpy as np from common import ark, pytest_ark @@ -12,6 +8,8 @@ @pytest_ark(need_torch=True) def test_eval_basic(): """Test basic Tensor.eval() — compile, run, return torch tensor.""" + import torch + x = torch.ones(64, dtype=torch.float32, device="cuda:0") * 3.0 out = ark.add(x, 2.0) @@ -25,6 +23,8 @@ def test_eval_basic(): @pytest_ark(need_torch=True) def test_eval_chain(): """Test eval on a chained computation.""" + import torch + x = torch.ones(64, dtype=torch.float16, device="cuda:0") * 4.0 y = ark.mul(x, 2.0) z = ark.add(y, 1.0) @@ -39,6 +39,8 @@ def test_eval_chain(): @pytest_ark(need_torch=True) def test_eval_relu(): """Test eval with relu op.""" + import torch + x = torch.tensor( [-1.0, 0.0, 1.0, 2.0], dtype=torch.float32, device="cuda:0" ) @@ -55,6 +57,8 @@ def test_eval_relu(): @pytest_ark(need_torch=True) def test_eval_matmul(): """Test eval with matmul.""" + import torch + a = torch.ones(4, 64, dtype=torch.float16, device="cuda:0") b = torch.ones(64, 8, dtype=torch.float16, device="cuda:0") out = ark.matmul(a, b) @@ -67,8 +71,10 @@ def test_eval_matmul(): @pytest_ark(need_torch=True) -def test_eval_returns_fresh_runtime(): - """Test that eval creates a fresh Runtime each call (no state leak).""" +def test_eval_independent_calls(): + """Test that independent eval calls on different graphs produce correct results.""" + import torch + x = torch.ones(64, dtype=torch.float32, device="cuda:0") * 2.0 out = ark.add(x, 3.0) @@ -78,12 +84,12 @@ def test_eval_returns_fresh_runtime(): torch.ones(64, dtype=torch.float32, device="cuda:0") * 5.0, ) - # Second eval on the same tensor should also work - ark.init() - x2 = torch.ones(64, dtype=torch.float32, device="cuda:0") * 10.0 - out2 = ark.add(x2, 1.0) - result2 = out2.eval() - assert torch.allclose( - result2, - torch.ones(64, dtype=torch.float32, device="cuda:0") * 11.0, - ) + # Second eval on a different graph should also work + with ark.use_model(None): + x2 = torch.ones(64, dtype=torch.float32, device="cuda:0") * 10.0 + out2 = ark.add(x2, 1.0) + result2 = out2.eval() + assert torch.allclose( + result2, + torch.ones(64, dtype=torch.float32, device="cuda:0") * 11.0, + ) diff --git a/python/unittest/test_model.py b/python/unittest/test_model.py index ad40d752..92a82785 100644 --- a/python/unittest/test_model.py +++ b/python/unittest/test_model.py @@ -3,6 +3,7 @@ from common import ark, pytest_ark import json +import pytest @pytest_ark() @@ -26,3 +27,65 @@ def test_model(): assert m_json.get("Nodes", None) is not None assert len(m_json["Nodes"]) == 0 + + +@pytest_ark() +def test_set_model_valid(): + """set_model accepts a Model instance.""" + m = ark.Model() + ark.set_model(m) + assert ark.Model.get_model() is m + + +@pytest_ark() +def test_set_model_invalid(): + """set_model rejects non-Model arguments.""" + with pytest.raises(ark.InvalidUsageError): + ark.set_model("not a model") + + +@pytest_ark() +def test_current_model_auto_creates(): + """current_model returns a Model, auto-creating if needed.""" + m = ark.current_model() + assert isinstance(m, ark.Model) + + +@pytest_ark() +def test_use_model_restores_previous(): + """use_model restores the previous model on exit.""" + outer = ark.current_model() + inner = ark.Model() + with ark.use_model(inner) as m: + assert m is inner + assert ark.Model.get_model() is inner + assert ark.Model.get_model() is outer + + +@pytest_ark() +def test_use_model_none_creates_fresh(): + """use_model(None) creates a fresh model.""" + outer = ark.current_model() + with ark.use_model(None) as m: + assert isinstance(m, ark.Model) + assert m is not outer + assert ark.Model.get_model() is outer + + +@pytest_ark() +def test_use_model_invalid(): + """use_model rejects non-Model arguments.""" + with pytest.raises(ark.InvalidUsageError): + with ark.use_model("not a model"): + pass + + +@pytest_ark() +def test_use_model_restores_on_exception(): + """use_model restores the previous model even if body raises.""" + outer = ark.current_model() + inner = ark.Model() + with pytest.raises(RuntimeError): + with ark.use_model(inner): + raise RuntimeError("boom") + assert ark.Model.get_model() is outer From 58e2d845daf7703a23c95219d413ca2f2948e3f4 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 9 Jun 2026 22:41:58 +0000 Subject: [PATCH 3/5] Fix black 26.5.1 lint: remove extra blank line before __all__ in ops.py --- python/ark/ops.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/ark/ops.py b/python/ark/ops.py index e0f0802b..470959ad 100644 --- a/python/ark/ops.py +++ b/python/ark/ops.py @@ -9,7 +9,6 @@ from .model import Model from . import log - __all__ = [ "tensor", "parameter", From f2821196ea27c4c76350f31badf0d6883e6b18e2 Mon Sep 17 00:00:00 2001 From: ark-dev Date: Wed, 10 Jun 2026 05:57:41 +0000 Subject: [PATCH 4/5] ark-dev: Update PR #255 by merging base into head and resolving any conflicts (BEHIND, behind_by=2). --- python/ark/data_type.py | 2 +- python/ark/model.py | 2 +- python/ark/tensor.py | 9 +++++++-- python/unittest/test_eval.py | 1 - 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/python/ark/data_type.py b/python/ark/data_type.py index 0caac294..f932fe61 100644 --- a/python/ark/data_type.py +++ b/python/ark/data_type.py @@ -8,9 +8,9 @@ __all__ = [ "DataType", + "bf16", "fp16", "fp32", - "bf16", "int32", "uint32", "int8", diff --git a/python/ark/model.py b/python/ark/model.py index b6cff59c..91bd99de 100644 --- a/python/ark/model.py +++ b/python/ark/model.py @@ -6,7 +6,7 @@ from . import log from .core import CoreModel -__all__ = ["Model", "set_model", "current_model", "use_model"] +__all__ = ["Model", "current_model", "set_model", "use_model"] ModelState = NewType("ModelState", None) diff --git a/python/ark/tensor.py b/python/ark/tensor.py index ec9e87db..1c14eda5 100644 --- a/python/ark/tensor.py +++ b/python/ark/tensor.py @@ -48,6 +48,10 @@ def __init__( _model if _model is not None else Model.get_model() ) self._device_id: int = Model.get_device_id() + # Note: _device_id is captured from global state, not derived from + # _model. In multi-device setups, callers should ensure the device id + # is set correctly before creating tensors. A future API may derive + # this from the model directly. def __hash__(self): return self._tensor.id() @@ -108,14 +112,15 @@ def __getitem__(self, index) -> "Tensor": new_offsets = Dims(new_offsets) new_padded_shape = Dims(new_padded_shape) new_tensor = Tensor( - Model.get_model().refer( + self._model.refer( self._tensor, new_shape, new_strides, new_offsets, new_padded_shape, "", - ) + ), + _model=self._model, ) new_tensor.requires_grad = self.requires_grad return new_tensor diff --git a/python/unittest/test_eval.py b/python/unittest/test_eval.py index 9ccff0e7..ee05bbbc 100644 --- a/python/unittest/test_eval.py +++ b/python/unittest/test_eval.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import numpy as np from common import ark, pytest_ark From 1d8109ceca1df53aee9313c50beadef6596f7d19 Mon Sep 17 00:00:00 2001 From: ark-dev Date: Wed, 10 Jun 2026 10:19:45 +0000 Subject: [PATCH 5/5] ark-dev: Update PR #255 by merging base branch into head branch and resolving any conflicts; cause: BEHIND base by 1 commit (PR #254 merge). --- python/ark/ops.py | 5 +++++ python/unittest/test_conversion.py | 12 ++++++++++++ python/unittest/test_eval.py | 18 ++++++++++++++++++ 3 files changed, 35 insertions(+) diff --git a/python/ark/ops.py b/python/ark/ops.py index 529625c6..de1ca1f9 100644 --- a/python/ark/ops.py +++ b/python/ark/ops.py @@ -53,6 +53,9 @@ def _ensure_ark(t): """ If *t* is a ``torch.Tensor``, convert it to an ARK ``Tensor`` via ``Tensor.from_torch()``. Otherwise return *t* unchanged. + + When used on an ``output`` parameter, the returned ARK tensor shares memory + with the original torch tensor via ``Tensor.from_torch``. """ if not _no_torch and isinstance(t, torch.Tensor): return Tensor.from_torch(t) @@ -691,6 +694,8 @@ def all_reduce_packet( Returns: Tensor: The reduced tensor. """ + input = _ensure_ark(input) + output = _ensure_ark(output) if output is not NullTensor: output = output._tensor _tensor = Model.get_model().all_reduce_packet( diff --git a/python/unittest/test_conversion.py b/python/unittest/test_conversion.py index e542ad27..e5582723 100644 --- a/python/unittest/test_conversion.py +++ b/python/unittest/test_conversion.py @@ -90,6 +90,18 @@ def test_conversion_ensure_ark_converts_torch(): assert ark_t.shape() == [64] +@pytest_ark(need_torch=True) +def test_conversion_cpu_torch_raises(): + """Passing a CPU torch tensor to an op raises InvalidUsageError.""" + import torch + import ark.log as log + import pytest + + cpu_t = torch.ones(64, dtype=torch.float32) # CPU + with pytest.raises(log.InvalidUsageError): + ark.add(cpu_t, 1.0) + + @pytest_ark(need_torch=True) def test_conversion_ops_accept_torch(): """Test that ops accept torch tensors directly via _ensure_ark.""" diff --git a/python/unittest/test_eval.py b/python/unittest/test_eval.py index ee05bbbc..6fb4474a 100644 --- a/python/unittest/test_eval.py +++ b/python/unittest/test_eval.py @@ -92,3 +92,21 @@ def test_eval_independent_calls(): result2, torch.ones(64, dtype=torch.float32, device="cuda:0") * 11.0, ) + + +@pytest_ark(need_torch=True) +def test_eval_cross_model_isolation(): + """eval() uses the model captured at tensor creation, not the current model.""" + import torch + + with ark.use_model() as m1: + x = torch.ones(64, dtype=torch.float32, device="cuda:0") + out = ark.add(x, 2.0) + + # Switch to a different model — out should still eval against m1 + with ark.use_model() as m2: + _ = ark.tensor([32], ark.fp32) # populate m2 + + result = out.eval() + expected = torch.ones(64, dtype=torch.float32, device="cuda:0") * 3.0 + assert torch.allclose(result, expected)