| 408 | } |
| 409 | |
| 410 | Status LookupSliceInfo( |
| 411 | BundleReader* reader, const string& table_name, |
| 412 | const std::vector<string>& tensibles_name, |
| 413 | std::vector<TensorSliceProto>* out_table_slices, |
| 414 | std::vector<std::vector<TensorSliceProto>>* out_tensible_slices) { |
| 415 | std::vector<TensorSliceProto> table_slices; |
| 416 | TF_RETURN_IF_ERROR(reader->LookupTensorSliceProtos( |
| 417 | table_name, &table_slices)); |
| 418 | |
| 419 | std::sort(table_slices.begin(), table_slices.end(), |
| 420 | [](const TensorSliceProto& lhs, const TensorSliceProto& rhs) -> bool { |
| 421 | return lhs.hash_slice_begin() < rhs.hash_slice_begin(); |
| 422 | }); |
| 423 | |
| 424 | std::vector<std::vector<TensorSliceProto>> tensible_slices( |
| 425 | tensibles_name.size()); |
| 426 | for (size_t i = 0; i < tensibles_name.size(); i++) { |
| 427 | Status st = reader->LookupTensorSliceProtos( |
| 428 | tensibles_name[i], &tensible_slices[i]); |
| 429 | if (!st.ok()) { |
| 430 | LOG(ERROR) << "Slice " << tensibles_name[i] << " not found, ignored!"; |
| 431 | tensible_slices[i].clear(); |
| 432 | continue; |
| 433 | } |
| 434 | std::sort(tensible_slices[i].begin(), tensible_slices[i].end(), |
| 435 | [](const TensorSliceProto& lhs, const TensorSliceProto& rhs) -> bool { |
| 436 | return lhs.hash_slice_begin() < rhs.hash_slice_begin(); |
| 437 | }); |
| 438 | } |
| 439 | for (auto&& slice : tensible_slices) { |
| 440 | if (slice.empty()) { |
| 441 | continue; |
| 442 | } |
| 443 | if (slice.size() != table_slices.size()) { |
| 444 | return errors::FailedPrecondition("Checkpoint's Slice is mismatched"); |
| 445 | } |
| 446 | for (size_t i = 0; i < table_slices.size(); ++i) { |
| 447 | if (slice[i].hash_slice_begin() != table_slices[i].hash_slice_begin()) { |
| 448 | return errors::FailedPrecondition("Checkpoint's Slice is mismatched"); |
| 449 | } |
| 450 | if (slice[i].hash_slice_length() != table_slices[i].hash_slice_length()) { |
| 451 | return errors::FailedPrecondition("Checkpoint's Slice is mismatched"); |
| 452 | } |
| 453 | if (slice[i].extent(0).start() != table_slices[i].extent(0).start()) { |
| 454 | return errors::FailedPrecondition("Checkpoint's Slice is mismatched"); |
| 455 | } |
| 456 | if (slice[i].extent(0).length() != table_slices[i].extent(0).length()) { |
| 457 | return errors::FailedPrecondition("Checkpoint's Slice is mismatched"); |
| 458 | } |
| 459 | } |
| 460 | } |
| 461 | *out_table_slices = std::move(table_slices); |
| 462 | *out_tensible_slices = std::move(tensible_slices); |
| 463 | return Status::OK(); |
| 464 | } |
| 465 | |
| 466 | void BuildRestoreSlice( |
| 467 | const std::vector<TensorSliceProto>& table_slices, |