secretflow.device package#

Subpackages#

Submodules#

secretflow.device.driver module#

Functions:

with_device(dev, *[, num_returns, ...])

Set up a wrapper for scheduling function to this device.

to(device, data[, spu_vis])

Device object conversion.

reveal(func_or_object)

Get plaintext data from device.

wait(objects)

Wait for device objects until all are ready or error occurrency.

init([parties, address, num_cpus, ...])

Connect to an existing Ray cluster or start one and connect to it.

shutdown()

Disconnect the worker, and terminate processes started by secretflow.init().

secretflow.device.driver.with_device(dev: Device, *, num_returns: Optional[int] = None, static_argnames: Optional[Union[str, Iterable[str]]] = None)[源代码]#

Set up a wrapper for scheduling function to this device.

Agrs:

dev (Device): Target device. num_returns (int): Number of returned DeviceObject. static_argnames (Union[str, Iterable[str], None]): See jax.jit() docstring.

示例

>>> p1, spu = PYU(), SPU()
>>> # dynamic decorator
>>> x = with_device(p1)(load_data)('alice.csv')
>>> # static decorator
>>> @with_device(spu)
>>> def selu(x, alpha=1.67, lmbda=1.05):
>>>     return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
>>> x_ = x.to(spu)
>>> y = selu(x_)
secretflow.device.driver.to(device: Device, data: Any, spu_vis: str = 'secret')[源代码]#

Device object conversion.

参数
  • device (Device) – Target device.

  • data (Any) – DeviceObject or plaintext data.

  • spu_vis (str) – Deivce object visibility, SPU device only. secret: Secret sharing with protocol spdz-2k, aby3, etc. public: Public sharing, which means data will be replicated to each node.

返回

Target device object.

返回类型

DeviceObject

secretflow.device.driver.reveal(func_or_object)[源代码]#

Get plaintext data from device.

NOTE: Use this function with extreme caution, as it may cause privacy leaks. In SecretFlow, we recommend that data should flow between different devices and rarely revealed to driver. Only use this function when data dependency control flow occurs.

参数

func_or_object – May be callable or any Python objects which contains Device objects.

secretflow.device.driver.wait(objects: Any)[源代码]#

Wait for device objects until all are ready or error occurrency.

参数

objects – struct of device objects.

secretflow.device.driver.init(parties: Optional[Union[List[str], str]] = None, address: Optional[str] = None, num_cpus: Optional[int] = None, log_to_driver=False, omp_num_threads: Optional[int] = None, **kwargs)[源代码]#

Connect to an existing Ray cluster or start one and connect to it.

参数
  • parties – parties this node represents, e.g: ‘alice’, [‘alice’, ‘bob’, ‘carol’].

  • address – The address of the Ray cluster to connect to. If this address is not provided, then a raylet, a plasma store, a plasma manager, and some workers will be started.

  • num_cpus – Number of CPUs the user wishes to assign to each raylet.

  • log_to_driver – Whether direct output of worker processes on all nodes to driver.

  • omp_num_threads – set environment variable OMP_NUM_THREADS. It works only when address is None.

  • **kwargs – see ray.init() parameters.

secretflow.device.driver.shutdown()[源代码]#

Disconnect the worker, and terminate processes started by secretflow.init().

This will automatically run at the end when a Python process that uses Ray exits. It is ok to run this twice in a row. The primary use case for this function is to cleanup state between tests.

secretflow.device.proxy module#

Functions:

proxy(device_object_type[, max_concurrency])

Define a device class which should accept DeviceObject as method parameters and return DeviceObject.

secretflow.device.proxy.proxy(device_object_type: Type[DeviceObject], max_concurrency=None)[源代码]#

Define a device class which should accept DeviceObject as method parameters and return DeviceObject.

This proxy function mainly does the following work: 1. Add an additional parameter device: Device to init method __init__. 2. Wrap class methods, allow passing DeviceObject as parameters, which must be on the same device as the class instance. 3. According to the return annotation of class methods, return the corresponding number of DeviceObject.

@proxy(PYUObject)
class Model:
    def __init__(self, builder):
        self.weights = builder()

    def build_dataset(self, x, y):
        self.dataset_x = x
        self.dataset_y = y

    def get_weights(self) -> np.ndarray:
        return self.weights

    def train_step(self, step) -> Tuple[np.ndarray, int]:
        return self.weights, 100

alice = PYU('alice')
model = Model(builder, device=alice)
x, y = alice(load_data)()
model.build_dataset(x, y)
w = model.get_weights()
w, n = model.train_step(10)
参数
  • device_object_type (Type[DeviceObject]) – DeviceObject type, eg. PYUObject.

  • max_concurrency (int) – Actor threadpool size.

返回

Wrapper function.

返回类型

Callable

Module contents#