],
)

tf_cc_test(
    name = "parallel_device_remote_test",
    srcs = ["parallel_device_remote_test.cc"],
    # TODO(b/136478427): Enable global heap checking when servers shut down
    # cleanly.
    args = ["--heap_check="],
    deps = [
        ":parallel_device",
        ":parallel_device_testlib",
        "//tensorflow/c:c_api",
        "//tensorflow/c:c_api_experimental",

    echo "# Size of $TF_WHEEL is $WHEEL_MEGABYTES / $LARGEST_OK_SIZE megabytes." >&3
    test "$WHEEL_MEGABYTES" -le "$LARGEST_OK_SIZE"
}

# Note: this runs before the tests further down the file, so TF is installed in
# the venv and the venv is active when those tests run. The venv gets cleaned
# up in teardown_file() above.
@test "Wheel is installable" {
    python3 -m venv /tf/venv

    ],
)

tf_cuda_cc_test(
    name = "c_api_remote_test",
    size = "medium",
    srcs = [
        "c_api_remote_test.cc",
    ],
    # TODO(b/136478427): Figure out how to correctly shut the server down
    args = ["--heap_check="],
    tags = [
        "no_windows",
    ],
    deps = [
        ":c_api",
        ":c_api_experimental",
        ":c_api_internal",
        ":c_api_remote_test_util",

  TFE_ContextRemoveFunction(ctx, "AddVariablesFunction", status);
  TFE_DeleteContext(ctx);

  TF_DeleteStatus(status);

  // TODO(b/136478427): Figure out how to correctly shut the server down.
  worker_server1.release();
  worker_server2.release();
}

TEST(CAPI, TestLocalFunctionWithPackedInput) {
  TestFunctionWithPackedInput(/*remote=*/false);
}

    ```

    **Note:** You can add `--test_sharding_strategy=disabled` to the `flags` to
    disable the sharding so that all the test outputs are in one file. However,
    it may slow down the tests for not running in parallel and may cause the
    test to timeout but it could be useful when you need to execute a single
    test or more in general your filtered/selected tests have a very low

      // One of the inputs we're trying to pack is on a custom device. We'll let
      // the first custom device we see handle all of the packing.
      auto* custom_device_handle =
          tensorflow::down_cast<tensorflow::CustomDeviceTensorHandle*>(
              unwrapped_handle);
      tensorflow::ImmediateExecutionTensorHandle* result;
      status->status = custom_device_handle->device()->Pack(

    if (!t) {
      return tensorflow::errors::InvalidArgument(
          "Unable to cast input to GraphTensor");
    }
    inputs_.push_back(t->output_);
    *output = tensorflow::down_cast<TracingTensorHandle*>(outputs[0]);
    return absl::OkStatus();
  }

  Status Finalize(OutputList* outputs, AbstractFunction** f) override {
    std::vector<TF_Output> graph_outputs;

                              &Deallocator, &deallocator_called);

  TF_Tensor* o = TF_TensorMaybeMove(t);
  ASSERT_TRUE(o == nullptr);  // It is unsafe to move memory TF might not own.
  TF_DeleteTensor(t);
  EXPECT_TRUE(deallocator_called);
}

TEST(CAPI, LibraryLoadFunctions) {
  // TODO(b/73318067): Fix linking for the GPU test generated by the

// The config key must be set before getting its value. Getting value of
// non-existing config keys will result in errors.
// If `timeout_in_ms=0`, this call will block until the key-value is set or the
// worker shuts down.
TF_CAPI_EXPORT extern void TFE_GetConfigKeyValue(TFE_Context* ctx,
                                                 const char* key,
                                                 int64_t timeout_in_ms,

  tensorflow::TensorProto from_tensor_proto;
  status->status = BufferToMessage(from, &from_tensor_proto);
  if (!status->status.ok()) {
    return;
  }
  status->status =
      tensorflow::down_cast<tensorflow::TensorInterface*>(to->tensor)
          ->FromProto(from_tensor_proto);
}
// --------------------------------------------------------------------------