Merge pull request #99 from deepray-AI/1022

fuhailin · web-flow · commit ebca9eebeb75 · 2025-11-05T19:28:43.000+08:00
[OPS] TFOkStatus compatibility code for C++
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
@@ -29,8 +29,10 @@ jobs:
     name: Notebook lint
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/setup-python@v1
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
+    - uses: actions/setup-python@v6
+      with:
+        python-version: '3.13' 
     - name: Install tensorflow-docs
       run: python3 -m pip install -U git+https://github.com/tensorflow/docs
     - name: Lint notebooks
diff --git a/build_deps/tf_dependency/build_defs.bzl.tpl b/build_deps/tf_dependency/build_defs.bzl.tpl
@@ -2,4 +2,4 @@
 
 D_GLIBCXX_USE_CXX11_ABI = "%{tf_cx11_abi}"
 CPLUSPLUS_VERSION = "%{tf_cplusplus_ver}"
-DTF_VERSION_INTEGER = "%{tf_version_integer}"
+TF_VERSION_INTEGER = "%{tf_version_integer}"
diff --git a/build_deps/tf_dependency/tf_configure.bzl b/build_deps/tf_dependency/tf_configure.bzl
@@ -213,7 +213,7 @@ def _tf_pip_impl(repository_ctx):
     tf_shared_cc_library_path = "%s/%s" % (tf_shared_library_dir, tf_shared_cc_library_name)
     tf_cx11_abi = "-D_GLIBCXX_USE_CXX11_ABI=%s" % (repository_ctx.os.environ[_TF_CXX11_ABI_FLAG])
     tf_cplusplus_ver = "-std=%s" % repository_ctx.os.environ[_TF_CPLUSPLUS_VER]
-    tf_version_integer = "-DTF_VERSION_INTEGER=%s" % (repository_ctx.os.environ[_TF_VERSION_INTEGER])
+    tf_version_integer = "TF_VERSION_INTEGER=%s" % (repository_ctx.os.environ[_TF_VERSION_INTEGER])
 
     tf_shared_library_rule = _symlink_genrule_for_dir(
         repository_ctx,
diff --git a/configure.py b/configure.py
@@ -25,7 +25,7 @@
 import shutil
 import subprocess
 import sys
-from typing import Optional
+from typing import List, Optional
 
 import tensorflow as tf
 from packaging.version import Version
@@ -873,7 +873,7 @@ def _find_executable_or_die(executable_name: str, executable_path: Optional[str]
   return resolved_path_to_exe
 
 
-def _get_cuda_compute_capabilities_or_die() -> list[str]:
+def _get_cuda_compute_capabilities_or_die() -> List[str]:
   """Finds compute capabilities via nvidia-smi or rasies exception.
 
   Returns:
diff --git a/deepray/custom_ops/parquet_dataset/BUILD b/deepray/custom_ops/parquet_dataset/BUILD
@@ -14,6 +14,7 @@ custom_op_library(
     deps = [
         ":arrow_util",
         ":parquet_batch_reader",
+        "//deepray/custom_ops/utils:ok_status_util",
     ],
 )
 
@@ -31,6 +32,7 @@ cc_library(
         "DEEPREC_ARROW_ZEROCOPY",
     ],
     deps = [
+        "//deepray/custom_ops/utils:ok_status_util",
         "@local_config_tf//:libtensorflow_framework",
         "@local_config_tf//:tf_header_lib",
         "@org_apache_arrow//:arrow",
@@ -47,6 +49,7 @@ cc_library(
     ],
     deps = [
         ":arrow_util",
+        "//deepray/custom_ops/utils:ok_status_util",
         "@local_config_tf//:libtensorflow_framework",
         "@local_config_tf//:tf_header_lib",
         "@org_apache_arrow//:arrow",
diff --git a/deepray/custom_ops/parquet_dataset/cc/kernels/arrow_util.cc b/deepray/custom_ops/parquet_dataset/cc/kernels/arrow_util.cc
@@ -25,6 +25,7 @@ limitations under the License.
 
 #include "arrow/array.h"
 #include "arrow/util/thread_pool.h"
+#include "deepray/custom_ops/utils/ok_status_util.h"
 #include "eigen.h"
 #include "tensorflow/core/framework/allocation_description.pb.h"
 
@@ -252,7 +253,7 @@ class RaggedTensorBuilder : public ::arrow::ArrayVisitor {
 #define CASE_ARROW_ENUM_SET_DTYPE(PTR, ENUM)                       \
   case ENUM: {                                                     \
     *PTR = DataTypeToEnum<ArrowEnumToDataType<ENUM>::Type>::value; \
-    return OkStatus();                                             \
+    return TFOkStatus;                                             \
   }
 
 Status MakeDataTypeAndRaggedRankFromArrowDataType(
@@ -280,7 +281,7 @@ Status MakeDataTypeAndRaggedRankFromArrowDataType(
       return errors::Unimplemented("Arrow data type ", arrow_dtype->ToString(),
                                    " not supported.");
   }
-  return OkStatus();
+  return TFOkStatus;
 }
 
 Status MakeTensorsFromArrowArray(
@@ -297,7 +298,7 @@ Status MakeTensorsFromArrowArray(
 
   RaggedTensorBuilder builder(dtype, ragged_rank);
   TF_RETURN_IF_ARROW_ERROR(builder.Build(arrow_array, output_tensors));
-  return OkStatus();
+  return TFOkStatus;
 }
 
 int UpdateArrowCpuThreadPoolCapacityFromEnv() {
diff --git a/deepray/custom_ops/parquet_dataset/cc/kernels/parquet_batch_reader.cc b/deepray/custom_ops/parquet_dataset/cc/kernels/parquet_batch_reader.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "absl/strings/match.h"
 #include "arrow_util.h"
+#include "deepray/custom_ops/utils/ok_status_util.h"
 
 namespace tensorflow {
 namespace data {
@@ -44,7 +45,7 @@ class ParquetBatchReader::Impl {
 
   Status Open() {
     if (TF_PREDICT_TRUE(batch_reader_)) {
-      return OkStatus();
+      return TFOkStatus;
     }
     if (TF_PREDICT_FALSE(partition_index_ >= partition_count_)) {
       return errors::InvalidArgument("Partition index ", partition_index_,
@@ -101,7 +102,7 @@ class ParquetBatchReader::Impl {
 
     TF_RETURN_IF_ARROW_ERROR(reader_->GetRecordBatchReader(
         row_group_indices_, column_indices_, &batch_reader_));
-    return OkStatus();
+    return TFOkStatus;
   }
 
   Status Read(std::vector<Tensor>* output_tensors) {
@@ -123,7 +124,7 @@ class ParquetBatchReader::Impl {
           field_dtypes_[i], field_ragged_ranks_[i], arrays[i], output_tensors));
     }
 
-    return OkStatus();
+    return TFOkStatus;
   }
 
  private:
diff --git a/deepray/custom_ops/parquet_dataset/cc/kernels/parquet_dataset_ops.cc b/deepray/custom_ops/parquet_dataset/cc/kernels/parquet_dataset_ops.cc
@@ -16,6 +16,7 @@ limitations under the License.
 
 #include <unordered_set>
 
+#include "deepray/custom_ops/utils/ok_status_util.h"
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_def_builder.h"
@@ -83,14 +84,14 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
     return output_shapes_;
   }
 
-  Status CheckExternalState() const override { return OkStatus(); }
+  Status CheckExternalState() const override { return TFOkStatus; }
 
   string DebugString() const override {
     return "ParquetTabularDatasetOp::Dataset";
   }
 
   Status InputDatasets(std::vector<const DatasetBase*>* inputs) const override {
-    return OkStatus();
+    return TFOkStatus;
   }
 
  protected:
@@ -122,7 +123,7 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {
                        {"partition_index", partition_index},
                        {"drop_remainder", drop_remainder}},
                       output));
-    return OkStatus();
+    return TFOkStatus;
   }
 
  private:
@@ -159,7 +160,7 @@ class ParquetTabularDatasetOp::Dataset::Iterator
       return s;
     }
     *end_of_sequence = true;
-    return OkStatus();
+    return TFOkStatus;
   }
 
  protected:
diff --git a/deepray/custom_ops/utils/BUILD b/deepray/custom_ops/utils/BUILD
@@ -1,4 +1,5 @@
 load("@local_config_cuda//cuda:build_defs.bzl", "cuda_library")
+load("@local_config_tf//:build_defs.bzl", "TF_VERSION_INTEGER")
 load("@org_tensorflow//tensorflow:tensorflow.bzl", "tf_copts")
 
 package(default_visibility = ["//visibility:public"])
@@ -24,6 +25,9 @@ cc_library(
     srcs = [
         "ok_status_util.h",
     ],
+    defines = [
+        TF_VERSION_INTEGER,
+    ],
     visibility = ["//visibility:public"],
 )
 
diff --git a/deepray/custom_ops/utils/ok_status_util.h b/deepray/custom_ops/utils/ok_status_util.h
@@ -33,7 +33,7 @@ This code is for compatibility.*/
 #else
 // #pragma message(PRINT_MACRO(TF_VERSION_INTEGER))
 // #define TFOkStatus Status::OK()
-#define TFOkStatus absl::OkStatus()
+#define TFOkStatus Status::OK()
 #endif
 }  // namespace deepray
 }  // namespace tensorflow
diff --git a/deepray/optimizers/adagrad.py b/deepray/optimizers/adagrad.py
@@ -18,10 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import sys
-
 import tensorflow as tf
-from absl import flags
 
 from deepray.custom_ops.embedding_variable import gen_kv_variable_ops
 from deepray.custom_ops.embedding_variable import kv_variable_ops
diff --git a/deepray/optimizers/ftrl.py b/deepray/optimizers/ftrl.py
@@ -1,7 +1,4 @@
-import sys
-
 import tensorflow as tf
-from absl import flags
 
 from deepray.custom_ops.embedding_variable import gen_kv_variable_ops
 from deepray.custom_ops.embedding_variable import kv_variable_ops
diff --git a/tools/docker/dev_container.Dockerfile b/tools/docker/dev_container.Dockerfile
@@ -78,6 +78,8 @@ ENV PATH=${PATH}:/opt/cmake/bin \
     LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/opt/cmake/lib
 
 RUN HOROVOD_WITH_MPI=1 HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_WITH_TENSORFLOW=1 pip install --no-cache-dir horovod -U
+COPY tools/releases/horovod_runner_patch.sh /tmp/
+RUN bash /tmp/horovod_runner_patch.sh
 
 COPY tools/install_deps/devel.requirements.txt /install_deps/devel.requirements.txt
 RUN pip install --no-cache-dir -r /install_deps/devel.requirements.txt -U
@@ -106,6 +108,8 @@ RUN bash /install_deps/install_openssh.sh
 
 # Setup openmpi
 COPY --from=openmpi_builder /opt/openmpi /opt/openmpi
+COPY --from=openmpi_builder /opt/ucx /opt/ucx
+COPY --from=openmpi_builder /opt/ucc /opt/ucc
 ENV PATH=${PATH}:/opt/openmpi/bin \
     LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/opt/openmpi/lib
 RUN mpirun --version

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ limitations under the License.`
`16`	`16`
`17`	`17`	`#include <unordered_set>`
`18`	`18`
	`19`	`+#include "deepray/custom_ops/utils/ok_status_util.h"`
`19`	`20`	`#include "tensorflow/core/framework/common_shape_fns.h"`
`20`	`21`	`#include "tensorflow/core/framework/op.h"`
`21`	`22`	`#include "tensorflow/core/framework/op_def_builder.h"`
`@@ -83,14 +84,14 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {`
`83`	`84`	`return output_shapes_;`
`84`	`85`	`}`
`85`	`86`
`86`		`- Status CheckExternalState() const override { return OkStatus(); }`
	`87`	`+ Status CheckExternalState() const override { return TFOkStatus; }`
`87`	`88`
`88`	`89`	`string DebugString() const override {`
`89`	`90`	`return "ParquetTabularDatasetOp::Dataset";`
`90`	`91`	`}`
`91`	`92`
`92`	`93`	`Status InputDatasets(std::vector<const DatasetBase> inputs) const override {`
`93`		`- return OkStatus();`
	`94`	`+ return TFOkStatus;`
`94`	`95`	`}`
`95`	`96`
`96`	`97`	`protected:`
`@@ -122,7 +123,7 @@ class ParquetTabularDatasetOp::Dataset : public DatasetBase {`
`122`	`123`	`{"partition_index", partition_index},`
`123`	`124`	`{"drop_remainder", drop_remainder}},`
`124`	`125`	`output));`
`125`		`- return OkStatus();`
	`126`	`+ return TFOkStatus;`
`126`	`127`	`}`
`127`	`128`
`128`	`129`	`private:`
`@@ -159,7 +160,7 @@ class ParquetTabularDatasetOp::Dataset::Iterator`
`159`	`160`	`return s;`
`160`	`161`	`}`
`161`	`162`	`*end_of_sequence = true;`
`162`		`- return OkStatus();`
	`163`	`+ return TFOkStatus;`
`163`	`164`	`}`
`164`	`165`
`165`	`166`	`protected:`