Skip to content

Add batch processing logic to expressions #93

Description

@jerrydzhang

Currently each individual call to cas_simplify, expression.score, expression.fit, are not easily parallelize on the python side since it still holds the GIL. These are quite common operations that are call on large population of expressions at once.

It would be advantageous for the cpp side to expose batch processing here is an example I have been using that inserts into bingo/expressions/agraph/cppagraph/bindings/bind_expression.cpp

  m.def(
      "fit_batch",
      [](py::iterable expression_objects,
         const RowMatrixXd &X,
         const Eigen::VectorXd &y,
         const std::string &,
         int n_threads)
      {
          std::vector<py::object> owners;
          std::vector<AGraphExpression *> expressions;

          for (py::handle item : expression_objects)
          {
              if (!py::isinstance<AGraphExpression>(item))
              {
                  throw py::type_error(
                      "Every item in expressions must be an "
                      "AGraphExpression");
              }

              owners.emplace_back(
                  py::reinterpret_borrow<py::object>(item));

              expressions.push_back(
                  &owners.back().cast<AGraphExpression &>());
          }

          const size_t N = expressions.size();

          if (N == 0)
          {
              return py::list();
          }

          std::unordered_set<AGraphExpression *> unique_expressions;

          for (AGraphExpression *expression : expressions)
          {
              if (!unique_expressions.insert(expression).second)
              {
                  throw py::value_error(
                      "fit_batch does not accept duplicate expression "
                      "objects because fit() mutates them");
              }
          }

          if (n_threads <= 0)
          {
              const unsigned hw = std::thread::hardware_concurrency();
              n_threads = hw ? static_cast<int>(hw) : 1;
          }

          const size_t n_workers = std::min(
              static_cast<size_t>(n_threads), N);

          std::atomic<size_t> next{0};
          std::atomic<bool> stop{false};
          std::vector<std::exception_ptr> errors(N);

          {
              py::gil_scoped_release release;

              auto worker = [&]()
              {
                  while (!stop.load(std::memory_order_relaxed))
                  {
                      const size_t i = next.fetch_add(
                          1, std::memory_order_relaxed);

                      if (i >= N)
                      {
                          break;
                      }

                      try
                      {
                          expressions[i]->fit(X, y);
                      }
                      catch (...)
                      {
                          errors[i] = std::current_exception();
                          stop.store(true, std::memory_order_relaxed);
                      }
                  }
              };

              std::vector<std::thread> pool;
              pool.reserve(n_workers - 1);

              for (size_t t = 1; t < n_workers; ++t)
              {
                  pool.emplace_back(worker);
              }

              worker();

              for (auto &thread : pool)
              {
                  thread.join();
              }
          }

          for (const auto &error : errors)
          {
              if (error)
              {
                  std::rethrow_exception(error);
              }
          }

          py::list result;
          for (const py::object &owner : owners)
          {
              result.append(owner);
          }

          return result;
      },
      py::arg("expressions"),
      py::arg("X"),
      py::arg("y"),
      py::arg("metric") = "mse",
      py::arg("n_threads") = 0,
      R"doc(
  Fit a batch of expressions in parallel.

  Every expression is fitted against the same X and y. The expressions
  are modified in place, and the same objects are returned in input order.

  Parameters
  ----------
  expressions : iterable of AGraphExpression
      Expressions to fit. Duplicate object references are not allowed.
  X : numpy.ndarray
      Input data with shape (n_samples, n_features).
  y : numpy.ndarray
      Target values with shape (n_samples,).
  metric : str, default="mse"
      Present for compatibility with AGraphExpression.fit. Currently
      ignored by the native fit implementation.
  n_threads : int, default=0
      Number of worker threads. Zero selects hardware concurrency.

  Returns
  -------
  list of AGraphExpression
      The same expression objects, fitted and in the original order.

  Notes
  -----
  If one fit operation fails, some earlier expressions may already have
  been fitted. The operation is not transactional.
  )doc");

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions