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)[source]#

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.

Examples

>>> 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')[source]#

Device object conversion.

Parameters

  • 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.

Returns

Target device object.

Return type

DeviceObject

secretflow.device.driver.reveal(func_or_object)[source]#

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.

Parameters

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

secretflow.device.driver.wait(objects: Any)[source]#

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

Parameters

objects – struct of device objects.

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

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

Parameters

  • 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()[source]#

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)[source]#

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)
Parameters

  • device_object_type (Type[DeviceObject]) – DeviceObject type, eg. PYUObject.

  • max_concurrency (int) – Actor threadpool size.

Returns

Wrapper function.

Return type

Callable

Module contents#