def testPackedVariable(self, distribution): with distribution.scope(): v0 = variables_lib.Variable(0.) self.assertIsNone(v0._packed_var) distribution._enable_packed_variable_in_eager_mode = True with distribution.scope(): v1 = variables_lib.Variable(0) self.assertIsInstance(v1._packed_var, packed.PackedDistributedVariable) devices = v1._devices for i in range(1, len(devices)): with distribute_lib.ReplicaContext(distribution, i): v1.assign(i) val = v1._get() self.assertIsInstance(val, packed.PackedVarAndDevice) self.assertEqual(val.device, devices[0]) self.assertEqual(self.evaluate(val.read_value()), 0) for i in range(0, len(devices)): with distribute_lib.ReplicaContext(distribution, i): val = v1._get() self.assertIsInstance(val, packed.PackedVarAndDevice) self.assertEqual(val.device, devices[i]) self.assertEqual(self.evaluate(val.read_value()), i)
def _call_for_each_replica(self, fn, args, kwargs): with distribute_lib.ReplicaContext( self._container_strategy(), replica_id_in_sync_group=constant_op.constant(0, dtypes.int32)): # TODO(rchao): Support multi-replica per worker or sync-group. return distribute_utils.regroup((fn(*args, **kwargs), ))
def schedule(self, fn, args=None, kwargs=None): """Schedules `fn` to be dispatched to a worker for execution asynchronously. When calling `schedule` with a function `fn`, `fn` will be executed on a remote worker at some later time. The process is asynchronous, meaning `schedule` returns immediately, possibly without having the result ready yet. `schedule` returns a structure of `RemoteValue` object, which wraps the output of the function. Call `fetch()` on `RemoteValue` to wait for the function execution to finish and retrieve its output from the remote worker. `schedule` guarantees that `fn` will be executed on a worker at least once; it could be more than once if its corresponding worker fails in the middle of its execution. Note that since worker can fail at any point when executing the function, it is possible that the function is partially executed, but `Client` guarantees that in those events, the function will eventually be fully executed, possibly on a different worker that is available. If any previously scheduled function raises an error, `schedule` will fail by raising any one of those errors, and clear the errors collected so far. There are two implications when this happens: 1) user should call `schedule` with `fn` again to re-schedule, and 2) some of the previously scheduled functions may have not been executed. User can call `fetch` on the returned `RemoteValue` to inspect if they have executed, failed, or cancelled, and reschedule the corresponding function if needed. When `schedule` raises, it guarantees that there is no function that is still being executed. At this time, there is no support of worker assignment for function execution, or priority of the workers. `args` and `kwargs` are the arguments passed into `fn`, when `fn` is executed on a worker. They can be `PerWorkerValues`, which is a collection of values, each of which represents a component specific to a worker; in this case, the argument will be substituted with the corresponding component on the target worker. Arguments that are not `PerWorkerValues` will be passed into `fn` as-is. Args: fn: A `tf.function`; the function to be dispatched to a worker for execution asynchronously. args: Positional arguments for `fn`. kwargs: Keyword arguments for `fn`. Returns: A structure of `RemoteValue` object. Raises: Exception: one of the exceptions caught by the client by any previously scheduled function since the last time an error was thrown or since the beginning of the program. """ # TODO(b/160702436): Invoke `strategy.run` for user's function so it enters # a `ReplicaContext` in a logically correct way. with distribute_lib.ReplicaContext( self._strategy, replica_id_in_sync_group=constant_op.constant(0, dtypes.int32)): with self._translate_parameter_server_failure(): return self.cluster.schedule(fn, args=args, kwargs=kwargs)
def _call_for_each_replica(self, fn, args, kwargs): with distribute_lib.ReplicaContext(self._container_strategy(), replica_id_in_sync_group=0), \ ops.device(self._ipu_device): # Make sure it is compiled as a single engine when called in graph mode. # This is similar to the mechanism used by xla.compile. xla_context = control_flow_ops.XLAControlFlowContext() try: xla_context.Enter() _validate_function_for_arguments(fn, args, kwargs) return fn(*args, **kwargs) finally: xla_context.Exit()
def _call_for_each_replica(self, fn, args, kwargs): with distribute_lib.ReplicaContext( self._container_strategy(), replica_id_in_sync_group=0), \ ops.device(self._ipu_device): return fn(*args, **kwargs)