worker.hpp

Go to the documentation of this file.

/*
 *   Copyright 2025 Huawei Technologies Co., Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <thread>
#include <memory>
#include <utility>
#include <vector>
#include <set>
#include <unistd.h>
#include <hicr/core/definitions.hpp>
#include <hicr/core/exceptions.hpp>
#include <hicr/core/processingUnit.hpp>
#include <hicr/core/computeManager.hpp>
#include "task.hpp"
 
namespace HiCR::tasking
{
 
constexpr size_t _DEFAULT_SUSPEND_INTERVAL_MS = 1000;
 
constexpr size_t _MILISECONDS_PER_SECOND = 1000;
 
using pullFunction_t = std::function<HiCR::tasking::Task *()>;
 
class Worker
{
  public:
 
  enum callback_t
  {
    onWorkerStart,
 
    onWorkerTaskPulled,
 
    onWorkerSuspend,
 
    onWorkerResume,
 
    onWorkerTerminate,
  };
 
  using workerCallbackMap_t = HiCR::tasking::CallbackMap<Worker *, callback_t>;
 
  enum state_t
  {
    uninitialized,
 
    ready,
 
    running,
 
    suspending,
 
    suspended,
 
    resuming,
 
    terminating,
 
    terminated
  };
 
  Worker(HiCR::ComputeManager *executionStateComputeManager,
         HiCR::ComputeManager *processingUnitComputeManager,
         pullFunction_t        pullFunction,
         workerCallbackMap_t  *callbackMap = nullptr)
    : _executionStateComputeManager(executionStateComputeManager),
      _processingUnitComputeManager(processingUnitComputeManager),
      _pullFunction(std::move(pullFunction)),
      _callbackMap(callbackMap)
  {
    _state = state_t::uninitialized;
  }
 
  virtual ~Worker() = default;
 
  __INLINE__ const state_t getState() { return _state; }
 
  __INLINE__ void setCallbackMap(workerCallbackMap_t *callbackMap) { _callbackMap = callbackMap; }
 
  __INLINE__ workerCallbackMap_t *getCallbackMap() { return _callbackMap; }
 
  __INLINE__ HiCR::tasking::Task *getCurrentTask() { return _currentTask; }
 
  __INLINE__ void initialize()
  {
    // Grabbing state value
    auto prevState = _state.load();
 
    // Checking we have at least one assigned resource
    if (_processingUnits.empty()) HICR_THROW_LOGIC("Attempting to initialize worker without any assigned resources");
 
    // Checking state
    if (prevState != state_t::uninitialized && prevState != state_t::terminated) HICR_THROW_RUNTIME("Attempting to initialize already initialized worker");
 
    // Initializing all resources
    for (auto &r : _processingUnits) _processingUnitComputeManager->initialize(r);
 
    // Transitioning state
    _state = state_t::ready;
  }
 
  __INLINE__ void start()
  {
    // Grabbing state value
    auto prevState = _state.load();
 
    // Checking state
    if (prevState != state_t::ready) HICR_THROW_RUNTIME("Attempting to start worker that is not in the 'initialized' state");
 
    // Setting state
    _state = state_t::running;
 
    // Creating new execution unit (the processing unit must support an execution unit of 'host' type)
    auto executionUnit = _executionStateComputeManager->createExecutionUnit([](void *worker) { static_cast<HiCR::tasking::Worker *>(worker)->mainLoop(); });
 
    // Creating worker's execution state
    auto executionState = _executionStateComputeManager->createExecutionState(executionUnit, this);
 
    // Launching worker in the lead resource (first one to be added)
    _processingUnitComputeManager->start(_processingUnits[0], executionState);
  }
 
  __INLINE__ bool suspend()
  {
    // Doing an atomic exchange
    state_t expected  = state_t::running;
    bool    succeeded = _state.compare_exchange_weak(expected, state_t::suspending);
 
    // Checking exchange
    return succeeded;
  }
 
  __INLINE__ bool resume()
  {
    // Doing an atomic exchange
    state_t expected  = state_t::suspended;
    bool    succeeded = _state.compare_exchange_weak(expected, state_t::resuming);
 
    // Checking exchange
    return succeeded;
  }
 
  __INLINE__ void terminate()
  {
    // Transitioning state
    auto prevState = _state.exchange(state_t::terminating);
 
    // Checking state
    if (prevState != state_t::running && prevState != state_t::suspending) HICR_THROW_RUNTIME("Attempting to stop worker that is not in a terminate-able state");
  }
 
  __INLINE__ void await()
  {
    // Getting state
    auto prevState = _state.load();
 
    if (prevState != state_t::terminating && prevState != state_t::running && prevState != state_t::suspended && prevState != state_t::suspending && prevState != state_t::resuming)
      HICR_THROW_RUNTIME("Attempting to wait for a worker that has not yet started or has already terminated");
 
    // Wait for the resources to free up
    for (auto &p : _processingUnits) _processingUnitComputeManager->await(p);
 
    // Transitioning state
    _state = state_t::terminated;
  }
 
  __INLINE__ void addProcessingUnit(std::unique_ptr<HiCR::ProcessingUnit> pu) { _processingUnits.push_back(std::move(pu)); }
 
  __INLINE__ std::vector<std::unique_ptr<HiCR::ProcessingUnit>> &getProcessingUnits() { return _processingUnits; }
 
  __INLINE__ void setSuspendInterval(size_t suspendIntervalMs) { _suspendIntervalMs = suspendIntervalMs; }
 
  protected:
 
  __INLINE__ virtual bool checkResumeConditions() { return _state == state_t::resuming; }
 
  private:
 
  HiCR::ComputeManager *const _executionStateComputeManager;
 
  HiCR::ComputeManager *const _processingUnitComputeManager;
 
  HiCR::tasking::Task *_currentTask = nullptr;
 
  const pullFunction_t _pullFunction;
 
  size_t _suspendIntervalMs = _DEFAULT_SUSPEND_INTERVAL_MS;
 
  std::atomic<state_t> _state;
 
  std::vector<std::unique_ptr<HiCR::ProcessingUnit>> _processingUnits;
 
  workerCallbackMap_t *_callbackMap = nullptr;
 
  __INLINE__ void mainLoop()
  {
    // Calling appropriate callback
    if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onWorkerStart);
 
    // Start main worker loop (run until terminated)
    while (true)
    {
      // Attempt to get a task by executing the pull function
      _currentTask = _pullFunction();
 
      // Calling appropriate callback
      if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onWorkerTaskPulled);
 
      // If a task was returned, then start or execute it
      if (_currentTask != nullptr) 
      {
        // If the task hasn't been initialized yet, we need to do it now
        if (_currentTask->getState() == HiCR::ExecutionState::state_t::uninitialized)
        {
          // First, create new execution state for the processing unit
          auto executionState = _executionStateComputeManager->createExecutionState(_currentTask->getExecutionUnit(), _currentTask);
 
          // Then initialize the task with the new execution state
          _currentTask->initialize(std::move(executionState));
        }
 
        // Now actually run the task
        _currentTask->run();
      }
 
      // Requesting processing units to terminate as soon as possible
      if (_state == state_t::suspending) 
      {
        // Setting state as suspended
        _state = state_t::suspended;
 
        // Calling appropriate callback
        if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onWorkerSuspend);
 
        // Suspending other processing units
        for (size_t i = 1; i < _processingUnits.size(); i++) _processingUnitComputeManager->suspend(_processingUnits[i]);
 
        // Putting current processing unit to check every so often
        while (checkResumeConditions() == false) usleep(_suspendIntervalMs * _MILISECONDS_PER_SECOND);
 
        // Calling appropriate callback
        if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onWorkerResume);
 
        // Resuming other processing units
        for (size_t i = 1; i < _processingUnits.size(); i++) _processingUnitComputeManager->resume(_processingUnits[i]);
 
        // Setting worker as running
        _state = state_t::running;
      }
 
      // Requesting processing units to terminate as soon as possible
      if (_state == state_t::terminating)
      {
        // Calling appropriate callback
        if (_callbackMap != nullptr) _callbackMap->trigger(this, callback_t::onWorkerTerminate);
 
        // Terminate secondary processing units first
        for (size_t i = 1; i < _processingUnits.size(); i++) _processingUnitComputeManager->terminate(_processingUnits[i]);
 
        // Then terminate current processing unit
        _processingUnitComputeManager->terminate(_processingUnits[0]);
 
        // Return immediately
        return;
      }
    }
  }
}; // class Worker
 
} // namespace HiCR::tasking