#!/usr/bin/env python3
# coding=UTF-8
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""ObjectRef"""
import asyncio
import functools
import importlib
import json
from concurrent.futures import Future
import logging
import threading
from typing import Any, Union
from yr.exception import YRInvokeError, YRError, YRRuntimeError, GeneratorFinished, raise_yr_value_error
from yr.err_type import ErrorInfo, ErrorCode
from yr.libruntime_pb2 import FunctionMeta
import yr
from yr.common import constants
from yr.ds_tensor_client_manager import get_tensor_client
_logger = logging.getLogger(__name__)
_runtime_holder_module = None
_runtime_holder_lock = threading.Lock()
def _get_runtime_holder():
global _runtime_holder_module
if _runtime_holder_module is None:
with _runtime_holder_lock:
if _runtime_holder_module is None:
_runtime_holder_module = importlib.import_module("yr.runtime_holder")
return _runtime_holder_module
def _set_future_helper(
result: Any,
*,
f: Union[asyncio.Future, Future],
):
if f.done():
return
if isinstance(result, ErrorInfo):
if result.error_code in (ErrorCode.ERR_GENERATOR_FINISHED, ErrorCode.ERR_GENERATOR_FINISHED.value):
f.set_exception(GeneratorFinished(""))
return
if result.error_code not in (ErrorCode.ERR_OK, ErrorCode.ERR_OK.value):
f.set_exception(YRRuntimeError.from_error_info(result))
elif isinstance(result, YRInvokeError):
f.set_exception(result.origin_error())
elif isinstance(result, YRError):
f.set_exception(result)
elif isinstance(result, RuntimeError):
f.set_exception(result)
else:
f.set_result(result)
[docs]
class ObjectRef:
"""ObjectRef"""
_id = None
_task_id = None
[docs]
def __init__(self, object_id: str, task_id=None, need_incre=True, need_decre=True, exception=None,
enable_tensor_transport=False):
"""Initialize the ObjectRef."""
self._id = object_id
self._task_id = task_id
self._need_decre = need_decre
self._exception = exception
self._data = None
self._enable_tensor_transport = enable_tensor_transport
self.npu_obj_ids = []
self.instance_id = ""
runtime_holder = _get_runtime_holder()
global_runtime = runtime_holder.global_runtime.get_runtime()
if need_incre and global_runtime and exception is None:
global_runtime.increase_global_reference([self._id])
def __del__(self):
try:
runtime_holder = _get_runtime_holder()
global_runtime = runtime_holder.global_runtime.get_runtime()
except (AttributeError, ImportError, RuntimeError):
return
try:
if self._need_decre and global_runtime:
global_runtime.decrease_global_reference([self._id])
if self._enable_tensor_transport and len(self.npu_obj_ids) != 0:
ds_tensor_client = get_tensor_client()
ds_tensor_client.dev_delete(self.npu_obj_ids)
# In future versions, this step will no longer be required
if len(self.instance_id) != 0 and global_runtime:
opts = yr.InvokeOptions()
opts.is_delete_remote_tensor = True
global_runtime.invoke_instance(FunctionMeta(), self.instance_id, self.npu_obj_ids, opts, 1)
except (AttributeError, TypeError, NameError):
# 忽略模块清理时的属性访问错误
pass
def __copy__(self):
return self
def __deepcopy__(self, memo):
return self
def __str__(self):
return self.id
def __eq__(self, other):
return self.id == other.id
def __hash__(self):
return hash(self.id)
def __repr__(self):
return self.id
def __await__(self):
return self.as_future().__await__()
def __reduce__(self):
"""Support for cloudpickle serialization.
Returns a tuple (callable, args) for pickle to reconstruct this ObjectRef.
The _object_ref_deserializer will be called during unpickling.
"""
from yr.serialization.serializers import _object_ref_deserializer, global_thread_local
if not hasattr(global_thread_local, "object_refs"):
global_thread_local.object_refs = set()
global_thread_local.object_refs.add(self)
return _object_ref_deserializer, (self._id,)
@property
def task_id(self):
"""Task id."""
return self._task_id
@task_id.setter
def task_id(self, value):
"""Task id."""
if value is not None:
self._task_id = value
@property
def id(self):
"""ObjectRef id."""
return self._id
[docs]
def as_future(self) -> asyncio.Future:
"""
Wrap `ObjectRef` with an `asyncio.Future`.
Note that the future cancellation will not cancel the corresponding
task when the ObjectRef representing return object of a task.
Returns:
An `asyncio.Future` that wraps `ObjectRef`. Data type is `asyncio.Future`.
"""
return asyncio.wrap_future(self.get_future())
[docs]
def get_future(self):
"""
Get future.
Returns:
The future of ObjectRef. Data type is Future.
"""
f = Future()
if self._exception is not None:
_set_future_helper(self._exception, f=f)
return f
if self._data is not None:
f.set_result(self._data)
return f
self.on_complete(functools.partial(_set_future_helper, f=f))
f.object_ref = self
return f
[docs]
def wait(self, timeout=None):
"""
Wait stateless function done.
Args:
timeout (int, optional): The number of seconds to wait for the exception if the future isn't done.
The default value ``None`` indicates that there is no limit on the wait time.
"""
future = self.get_future()
if future is not None:
future.result(timeout=timeout)
[docs]
def is_exception(self) -> bool:
"""
Whether future exception.
Returns:
Whether future exception. Data type is bool.
"""
future = self.get_future()
if future is None:
return False
return future.exception() is not None
[docs]
def done(self):
"""
Return ``True`` if the obj future was cancelled or finished executing.
Returns:
Whether the obj future was cancelled or finished executing. Data type is bool.
"""
future = self.get_future()
if future:
return future.done()
return True
[docs]
def cancel(self):
"""
Cancel the obj future
Returns:
Returns ``True`` if the future was cancelled, ``False`` otherwise. A future cannot be cancelled if it is
running or has already completed. Data type is bool.
"""
future = self.get_future()
if future:
future.cancel()
[docs]
def on_complete(self, callback):
"""
Register callback.
Args:
callback (Callable): User callback.
"""
runtime_holder = _get_runtime_holder()
runtime_holder.global_runtime.get_runtime().set_get_callback(self.id, callback)
[docs]
def get(self, timeout: int = constants.DEFAULT_GET_TIMEOUT) -> Any:
"""This function is used to retrieve an object.
Args:
timeout (int, optional): The maximum time in seconds to wait for the
interface object to be retrieved. Defaults to ``300``.
Returns:
The retrieved object. Data type is Any.
Raises:
ValueError: If the timeout parameter is less than or equal to the 0 and is not -1.
YRInvokeError: If the retrieved result is an instance of YRInvokeError.
"""
self.exception()
if timeout <= constants.MIN_TIMEOUT_LIMIT and timeout != constants.NO_LIMIT:
raise_yr_value_error("Parameter 'timeout' should be greater than 0 or equal to -1 (no timeout)")
runtime_holder = _get_runtime_holder()
objects = runtime_holder.global_runtime.get_runtime().get([self.id], timeout, False)
result_str = objects[0]
try:
obj = json.loads(result_str)
except json.decoder.JSONDecodeError:
_logger.warning(f"Failed to decode the result with object ID [{self.id}] using 'json.loads'."
f"result string: {result_str}")
obj = result_str
return obj
[docs]
def exception(self):
"""Raise exception if exception is not none."""
if self._exception is not None:
raise self._exception
[docs]
def set_data(self, data):
"""
Set data.
Args:
data (ObjectRef): Data to be set.
"""
self._data = data
[docs]
def set_exception(self, e):
"""Set exception."""
self._exception = e
class ObjectRefDirect(ObjectRef):
"""ObjectRefDirect bypasses datasystem — no IncreaseRef/DecreaseRef.
Return values exceeding the truncation threshold (100MB) will be truncated.
"""
def __init__(self, object_id: str, task_id=None, exception=None):
super().__init__(object_id, task_id=task_id, need_incre=False, need_decre=False, exception=exception)
