1 /* 2 * Copyright 2022 Google LLC 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #include "tensorflow/core/framework/op_kernel.h" 17 #include "tensorflow/core/framework/op_requires.h" 18 #include "tensorflow/core/framework/tensor_shape.h" 19 #include "tensorflow/core/framework/versions.pb.h" 20 #include "tensorflow/core/platform/errors.h" 21 #include "tensorflow/core/platform/stringpiece.h" 22 #include "tensorflow/core/util/saved_tensor_slice.pb.h" 23 24 namespace fcp { 25 namespace { 26 27 using ::tensorflow::OpKernel; 28 using ::tensorflow::OpKernelConstruction; 29 using ::tensorflow::OpKernelContext; 30 31 } // namespace 32 33 class DeleteDirOp : public OpKernel { 34 public: DeleteDirOp(OpKernelConstruction * context)35 explicit DeleteDirOp(OpKernelConstruction* context) : OpKernel(context) {} 36 Compute(OpKernelContext * context)37 void Compute(OpKernelContext* context) override { 38 const tensorflow::Tensor& dirname_t = context->input(0); 39 const tensorflow::Tensor& recursively_t = context->input(1); 40 { 41 const int64_t size = dirname_t.NumElements(); 42 OP_REQUIRES( 43 context, size == 1, 44 tensorflow::errors::InvalidArgument( 45 "Input 0 (dirname) must be a string scalar; got a tensor of ", 46 size, "elements")); 47 } 48 { 49 const int64_t size = recursively_t.NumElements(); 50 OP_REQUIRES( 51 context, size == 1, 52 tensorflow::errors::InvalidArgument( 53 "Input 1 (recursively) must be a string scalar; got a tensor of ", 54 size, "elements")); 55 } 56 const tensorflow::tstring& dirname = 57 dirname_t.scalar<tensorflow::tstring>()(); 58 const bool recursively = recursively_t.scalar<bool>()(); 59 if (context->env()->IsDirectory(dirname).ok()) { 60 if (recursively) { 61 int64_t undeleted_files = 0; 62 int64_t undeleted_dirs = 0; 63 tensorflow::Status delete_status = context->env()->DeleteRecursively( 64 dirname, &undeleted_files, &undeleted_dirs); 65 if (!delete_status.ok()) { 66 // The directory could be already deleted by another op. Let's not 67 // propagate this error. 68 LOG(WARNING) << "Failed to recursively delete the directory '" 69 << dirname << "' (remaining files: " << undeleted_files 70 << ", remaining dirs: " << undeleted_dirs << "). " 71 << delete_status; 72 } 73 } else { 74 tensorflow::Status delete_status = context->env()->DeleteDir(dirname); 75 if (!delete_status.ok()) { 76 // The directory could be already deleted by another op. Let's not 77 // propagate this error. 78 LOG(WARNING) << "Failed to delete the directory '" << dirname << "'. " 79 << delete_status; 80 } 81 } 82 } 83 } 84 }; 85 86 REGISTER_OP("DeleteDir") 87 .Input("dirname: string") 88 .Input("recursively: bool") 89 .SetIsStateful(); 90 REGISTER_KERNEL_BUILDER(Name("DeleteDir").Device(tensorflow::DEVICE_CPU), 91 DeleteDirOp); 92 93 class DeleteFileOp : public OpKernel { 94 public: DeleteFileOp(OpKernelConstruction * context)95 explicit DeleteFileOp(OpKernelConstruction* context) 96 : OpKernel(context) {} 97 Compute(OpKernelContext * context)98 void Compute(OpKernelContext* context) override { 99 const tensorflow::Tensor& filename_t = context->input(0); 100 { 101 const int64_t size = filename_t.NumElements(); 102 OP_REQUIRES( 103 context, size == 1, 104 tensorflow::errors::InvalidArgument( 105 "Input 0 (filename) must be a string scalar; got a tensor of ", 106 size, "elements")); 107 } 108 const tensorflow::tstring& filename = 109 filename_t.scalar<tensorflow::tstring>()(); 110 if (context->env()->FileExists(filename).ok()) { 111 tensorflow::Status delete_status = context->env()->DeleteFile(filename); 112 if (!delete_status.ok()) { 113 // The file could be already deleted by another op. Let's not propagate 114 // this error. 115 LOG(WARNING) << "Failed to delete the file '" << filename << "'. " 116 << delete_status; 117 } 118 } 119 } 120 }; 121 122 REGISTER_OP("DeleteFile").Input("filename: string").SetIsStateful(); 123 REGISTER_KERNEL_BUILDER( 124 Name("DeleteFile").Device(tensorflow::DEVICE_CPU), 125 DeleteFileOp); 126 127 } // namespace fcp 128