// Dynamic cluster API.

// Update an existing context with a new set of servers defined in a ServerDef
// proto. Servers can be added to and removed from the list of remote workers
// in the context. A New set of servers identified by the ServerDef must be up
// when the context is updated.
//
// This API is for experimental usage and may be subject to change.

// --------------------------------------------------------------------------
// The deprecated session API.  Please switch to the above instead of
// TF_ExtendGraph(). This deprecated API can be removed at any time without
// notice.

typedef struct TF_DeprecatedSession TF_DeprecatedSession;

TF_CAPI_EXPORT extern TF_DeprecatedSession* TF_NewDeprecatedSession(
    const TF_SessionOptions*, TF_Status* status);

class DeviceThread;

// Forwards operations to `devices`, maintaining ParallelTensor with components
// placed on each underlying device.
class ParallelDevice {
 public:
  // Eager async execution is only supported when remote eager is not in use
  // (b/157523095).
  explicit ParallelDevice(const std::vector<std::string>& devices,
                          bool is_async = false, int in_flight_nodes_limit = 0);

  ~ParallelDevice();

TF_CAPI_EXPORT extern void TFE_ContextAddFunction(TFE_Context* ctx,
                                                  TF_Function* function,
                                                  TF_Status* status);

// Removes a function from the context. Once removed, you can no longer
// TFE_Execute it or TFE_Execute any TFE_Op which has it as an attribute or any
// other function which calls it as an attribute.

// Frees the memory associated with the library handle.
// Does NOT unload the library.
TF_CAPI_EXPORT extern void TF_DeletePluggableDeviceLibraryHandle(
    TF_Library* lib_handle);

// Removes `func_name` from `g`. If `func_name` is not in `g`, an error will be
// returned.
TF_CAPI_EXPORT extern void TF_GraphRemoveFunction(TF_Graph* g,
                                                  const char* func_name,

  // "Jacobian-vector product"; `Watch`ing a new Tensor and immediately calling
  // FetchJVP for it would return `tangent`.
  void Watch(int64_t tensor_id, Gradient* tangent);

  // Removes the gradient associated with tensor_id. Should be called when the
  // Tensor associated with `tensor_id` is deleted.
  void DeleteGradient(int64_t tensor_id);

  // string is empty, no mutation has occurred. Otherwise the string is a
  // description of the mutation suitable for returning to the user.
  //
  // Sessions are added to this map in TF_NewSession, and removed in
  // TF_DeleteSession.
  // TF_Graph may only / must be deleted when
  //   sessions.size() == 0 && delete_requested == true
  //
  // TODO(b/74949947): mutations currently trigger a warning instead of a bad

  virtual bool IsCustomDevice(const string& device_name) = 0;

  // Register a custom device. It will return error is the device name is
  // already registered.
  // TODO(tfrt-devs): Remove this method. Let caller register it directly into
  // CustomDeviceOpHandler.
  virtual Status RegisterCustomDevice(const string& name,
                                      std::unique_ptr<CustomDevice> device) = 0;

// TODO(b/143949615): After all filesystems are converted, this file will be
// moved to core/platform, and this class can become a singleton and replace the
// need for `Env::Default()`. At that time, we might decide to remove the need
// for `Env::Default()` altogether, but that's a different project, not in
// scope for now. I'm just mentioning this here as that transition will mean
// removal of the registration part from `Env` and adding it here instead: we