DPL analysis: support timeslice rate limiting in DPL resource manager

ktf · ktf · commit 7044ca33caf2 · 2025-10-12T10:49:27.000+02:00
Use DPL resource manager rather than the ad-hoc solution for reconstruction.
diff --git a/Framework/Core/include/Framework/CommonDataProcessors.h b/Framework/Core/include/Framework/CommonDataProcessors.h
@@ -37,6 +37,9 @@ struct CommonDataProcessors {
   /// and simply discards them. @a rateLimitingChannelConfig is the configuration
   /// for the rate limiting channel, if any required.
   static DataProcessorSpec getDummySink(std::vector<InputSpec> const& danglingInputs, std::string rateLimitingChannelConfig);
+  /// @return a dummy DataProcessorSpec which requires all the passed @a InputSpec
+  /// and simply discards them. Rate limiting goes through the DPL driver
+  static DataProcessorSpec getScheduledDummySink(std::vector<InputSpec> const& danglingInputs);
   static AlgorithmSpec wrapWithRateLimiting(AlgorithmSpec spec);
 };
 
diff --git a/Framework/Core/include/Framework/ComputingQuotaOffer.h b/Framework/Core/include/Framework/ComputingQuotaOffer.h
@@ -44,6 +44,8 @@ struct ComputingQuotaOffer {
   int64_t memory = 0;
   /// How much shared memory it can allocate
   int64_t sharedMemory = 0;
+  /// How many timeslices it can process without giving back control
+  int64_t timeslices = 0;
   /// How much runtime it can use before giving back the resource
   /// in milliseconds.
   int64_t runtime = 0;
diff --git a/Framework/Core/include/Framework/ResourcePolicyHelpers.h b/Framework/Core/include/Framework/ResourcePolicyHelpers.h
@@ -22,6 +22,7 @@ namespace o2::framework
 struct ResourcePolicyHelpers {
   static ResourcePolicy trivialTask(char const* taskMatcher);
   static ResourcePolicy cpuBoundTask(char const* taskMatcher, int maxCPUs = 1);
+  static ResourcePolicy rateLimitedSharedMemoryBoundTask(char const* taskMatcher, int maxMemory, int maxTimeslices);
   static ResourcePolicy sharedMemoryBoundTask(char const* taskMatcher, int maxMemory);
 };
 
diff --git a/Framework/Core/src/ArrowSupport.cxx b/Framework/Core/src/ArrowSupport.cxx
@@ -77,23 +77,24 @@ struct MetricIndices {
   size_t shmOfferBytesConsumed = -1;
   size_t timeframesRead = -1;
   size_t timeframesConsumed = -1;
+  size_t timeframesExpired = -1;
 };
 
 std::vector<MetricIndices> createDefaultIndices(std::vector<DeviceMetricsInfo>& allDevicesMetrics)
 {
   std::vector<MetricIndices> results;
 
   for (auto& info : allDevicesMetrics) {
-    MetricIndices indices;
-    indices.arrowBytesCreated = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-bytes-created");
-    indices.arrowBytesDestroyed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-bytes-destroyed");
-    indices.arrowMessagesCreated = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-messages-created");
-    indices.arrowMessagesDestroyed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-messages-destroyed");
-    indices.arrowBytesExpired = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-bytes-expired");
-    indices.shmOfferBytesConsumed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "shm-offer-bytes-consumed");
-    indices.timeframesRead = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "df-sent");
-    indices.timeframesConsumed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "consumed-timeframes");
-    results.push_back(indices);
+    results.emplace_back(MetricIndices{
+      .arrowBytesCreated = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-bytes-created"),
+      .arrowBytesDestroyed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-bytes-destroyed"),
+      .arrowMessagesCreated = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-messages-created"),
+      .arrowMessagesDestroyed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-messages-destroyed"),
+      .arrowBytesExpired = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "arrow-bytes-expired"),
+      .shmOfferBytesConsumed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "shm-offer-bytes-consumed"),
+      .timeframesRead = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "df-sent"),
+      .timeframesConsumed = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "consumed-timeframes"),
+      .timeframesExpired = DeviceMetricsHelper::bookNumericMetric<uint64_t>(info, "expired-timeframes")});
   }
   return results;
 }
@@ -258,6 +259,7 @@ o2::framework::ServiceSpec ArrowSupport::arrowBackendSpec()
                        int64_t totalMessagesDestroyed = 0;
                        int64_t totalTimeframesRead = 0;
                        int64_t totalTimeframesConsumed = 0;
+                       int64_t totalTimeframesExpired = 0;
                        auto &driverMetrics = sm.driverMetricsInfo;
                        auto &allDeviceMetrics = sm.deviceMetricsInfos;
                        auto &specs = sm.deviceSpecs;
@@ -266,9 +268,14 @@ o2::framework::ServiceSpec ArrowSupport::arrowBackendSpec()
                        static auto stateMetric = DeviceMetricsHelper::createNumericMetric<uint64_t>(driverMetrics, "rate-limit-state");
                        static auto totalBytesCreatedMetric = DeviceMetricsHelper::createNumericMetric<uint64_t>(driverMetrics, "total-arrow-bytes-created");
                        static auto shmOfferConsumedMetric = DeviceMetricsHelper::createNumericMetric<uint64_t>(driverMetrics, "total-shm-offer-bytes-consumed");
+                       // These are really to monitor the rate limiting
                        static auto unusedOfferedSharedMemoryMetric = DeviceMetricsHelper::createNumericMetric<int>(driverMetrics, "total-unused-offered-shared-memory");
+                       static auto unusedOfferedTimeslicesMetric = DeviceMetricsHelper::createNumericMetric<int>(driverMetrics, "total-unused-offered-timeslices");
                        static auto availableSharedMemoryMetric = DeviceMetricsHelper::createNumericMetric<int>(driverMetrics, "total-available-shared-memory");
+                       static auto availableTimeslicesMetric = DeviceMetricsHelper::createNumericMetric<int>(driverMetrics, "total-available-timeslices");
                        static auto offeredSharedMemoryMetric = DeviceMetricsHelper::createNumericMetric<int>(driverMetrics, "total-offered-shared-memory");
+                       static auto offeredTimeslicesMetric = DeviceMetricsHelper::createNumericMetric<int>(driverMetrics, "total-offered-timeslices");
+
                        static auto totalBytesDestroyedMetric = DeviceMetricsHelper::createNumericMetric<uint64_t>(driverMetrics, "total-arrow-bytes-destroyed");
                        static auto totalBytesExpiredMetric = DeviceMetricsHelper::createNumericMetric<uint64_t>(driverMetrics, "total-arrow-bytes-expired");
                        static auto totalMessagesCreatedMetric = DeviceMetricsHelper::createNumericMetric<uint64_t>(driverMetrics, "total-arrow-messages-created");
@@ -390,6 +397,18 @@ o2::framework::ServiceSpec ArrowSupport::arrowBackendSpec()
                            auto const& timestamps = DeviceMetricsHelper::getTimestampsStore<uint64_t>(deviceMetrics)[info.storeIdx];
                            lastTimestamp = std::max(lastTimestamp, timestamps[(info.pos - 1) % data.size()]);
                          }
+                         {
+                           size_t index = indices.timeframesExpired;
+                           assert(index < deviceMetrics.metrics.size());
+                           changed |= deviceMetrics.changed[index];
+                           MetricInfo info = deviceMetrics.metrics[index];
+                           assert(info.storeIdx < deviceMetrics.uint64Metrics.size());
+                           auto& data = deviceMetrics.uint64Metrics[info.storeIdx];
+                           auto value = (int64_t)data[(info.pos - 1) % data.size()];
+                           totalTimeframesExpired += value;
+                           auto const& timestamps = DeviceMetricsHelper::getTimestampsStore<uint64_t>(deviceMetrics)[info.storeIdx];
+                           lastTimestamp = std::max(lastTimestamp, timestamps[(info.pos - 1) % data.size()]);
+                         }
                        }
                        static uint64_t unchangedCount = 0;
                        if (changed) {
@@ -407,26 +426,45 @@ o2::framework::ServiceSpec ArrowSupport::arrowBackendSpec()
                          unchangedCount++;
                        }
                        changedCountMetric(driverMetrics, unchangedCount, timestamp);
-                       auto maxTimeframes = registry.get<RateLimitConfig>().maxTimeframes;
-                       if (maxTimeframes && (totalTimeframesRead - totalTimeframesConsumed) > maxTimeframes) {
-                         return;
-                       }
+
                        static const ResourceSpec shmResourceSpec{
                          .name = "shared memory",
                          .unit = "MB",
                          .api = "/shm-offer {}",
-                         .maxAvailable = (int64_t)calculateAvailableSharedMemory(registry),
+                         .maxAvailable = (int64_t)registry.get<RateLimitConfig>().maxMemory,
                          .maxQuantum = 100,
                          .minQuantum = 50,
                          .metricOfferScaleFactor = 1000000,
                        };
+                       static const ResourceSpec timesliceResourceSpec{
+                         .name = "timeslice",
+                         .unit = "timeslices",
+                         .api = "/timeslice-offer {}",
+                         .maxAvailable = (int64_t)registry.get<RateLimitConfig>().maxTimeframes,
+                         .maxQuantum = 2,
+                         .minQuantum = 1,
+                         .metricOfferScaleFactor = 1,
+                       };
                        static ResourceState shmResourceState{
                          .available = shmResourceSpec.maxAvailable,
                        };
+                       static ResourceState timesliceResourceState{
+                         .available = timesliceResourceSpec.maxAvailable,
+                       };
                        static ResourceStats shmResourceStats{
                          .enoughCount = shmResourceState.available - shmResourceSpec.minQuantum > 0 ? 1 : 0,
                          .lowCount = shmResourceState.available - shmResourceSpec.minQuantum > 0 ? 0 : 1
                        };
+                       static ResourceStats timesliceResourceStats{
+                         .enoughCount = shmResourceState.available - shmResourceSpec.minQuantum > 0 ? 1 : 0,
+                         .lowCount = shmResourceState.available - shmResourceSpec.minQuantum > 0 ? 0 : 1
+                       };
+
+                       offerResources(timesliceResourceState, timesliceResourceSpec, timesliceResourceStats,
+                                      specs, infos, manager, totalTimeframesConsumed, totalTimeframesExpired,
+                                      totalTimeframesRead, totalTimeframesConsumed, timestamp, driverMetrics,
+                                      availableTimeslicesMetric, unusedOfferedTimeslicesMetric, offeredTimeslicesMetric,
+                                      (void*)&sm);
 
                        offerResources(shmResourceState, shmResourceSpec, shmResourceStats,
                                       specs, infos, manager, shmOfferBytesConsumed, totalBytesExpired,
@@ -497,8 +535,8 @@ o2::framework::ServiceSpec ArrowSupport::arrowBackendSpec()
                        if (!once) {
                          O2_SIGNPOST_ID_GENERATE(sid, rate_limiting);
                          O2_SIGNPOST_EVENT_EMIT_INFO(rate_limiting, sid, "setup",
-                                                     "Rate limiting set up at %{bytes}llu MB distributed over %d readers",
-                                                     config->maxMemory, readers);
+                                                     "Rate limiting set up at %{bytes}llu MB and %llu timeframes distributed over %d readers",
+                                                     config->maxMemory, config->maxTimeframes, readers);
                          registry.registerService(ServiceRegistryHelpers::handleForService<RateLimitConfig>(config));
                          once = true;
                        } },
diff --git a/Framework/Core/src/CommonDataProcessors.cxx b/Framework/Core/src/CommonDataProcessors.cxx
@@ -44,6 +44,7 @@ using namespace o2::framework::data_matcher;
 
 // Special log to track callbacks we know about
 O2_DECLARE_DYNAMIC_LOG(callbacks);
+O2_DECLARE_DYNAMIC_LOG(rate_limiting);
 
 namespace o2::framework
 {
@@ -224,6 +225,35 @@ DataProcessorSpec CommonDataProcessors::getDummySink(std::vector<InputSpec> cons
     .labels = {{"resilient"}}};
 }
 
+// For the cases were the driver is guaranteed to be there (e.g. in analysis) we can use a
+// more sophisticated controller which can get offers for timeslices so that we can rate limit
+// across multiple input devices and rate limit shared memory usage without race conditions
+DataProcessorSpec CommonDataProcessors::getScheduledDummySink(std::vector<InputSpec> const& danglingOutputInputs)
+{
+  return DataProcessorSpec{
+    .name = "internal-dpl-injected-dummy-sink",
+    .inputs = danglingOutputInputs,
+    .algorithm = AlgorithmSpec{adaptStateful([](CallbackService& callbacks, DeviceState& deviceState, InitContext& ic) {
+      // We update the number of consumed timeframes based on the oldestPossingTimeslice
+      // this information will be aggregated in the driver which will then decide wether or not a new offer for
+      // a timeslice should be done and to which device
+      auto domainInfoUpdated = [](ServiceRegistryRef services, size_t timeslice, ChannelIndex channelIndex) {
+        LOGP(debug, "Domain info updated with timeslice {}", timeslice);
+        auto& timesliceIndex = services.get<TimesliceIndex>();
+        auto oldestPossingTimeslice = timesliceIndex.getOldestPossibleOutput().timeslice.value;
+        auto& stats = services.get<DataProcessingStats>();
+        O2_SIGNPOST_ID_GENERATE(sid, rate_limiting);
+        O2_SIGNPOST_EVENT_EMIT(rate_limiting, sid, "run", "Consumed timeframes to be set to %zu.", oldestPossingTimeslice);
+        stats.updateStats({(int)ProcessingStatsId::CONSUMED_TIMEFRAMES, DataProcessingStats::Op::Set, (int64_t)oldestPossingTimeslice});
+      };
+      callbacks.set<CallbackService::Id::DomainInfoUpdated>(domainInfoUpdated);
+
+      return adaptStateless([]() {
+      });
+    })},
+    .labels = {{"resilient"}}};
+}
+
 AlgorithmSpec CommonDataProcessors::wrapWithRateLimiting(AlgorithmSpec spec)
 {
   return PluginManager::wrapAlgorithm(spec, [](AlgorithmSpec::ProcessCallback& original, ProcessingContext& pcx) -> void {
diff --git a/Framework/Core/src/ComputingQuotaEvaluator.cxx b/Framework/Core/src/ComputingQuotaEvaluator.cxx
@@ -40,6 +40,7 @@ ComputingQuotaEvaluator::ComputingQuotaEvaluator(ServiceRegistryRef ref)
     0,
     0,
     0,
+    0,
     -1,
     -1,
     OfferScore::Unneeded,
@@ -97,7 +98,7 @@ bool ComputingQuotaEvaluator::selectOffer(int task, ComputingQuotaRequest const&
                         result.size(), totalOffer.cpu, totalOffer.memory, totalOffer.sharedMemory);
       for (auto& offer : result) {
         // We pretend each offer id is a pointer, to have a unique id.
-        O2_SIGNPOST_ID_FROM_POINTER(oid, quota, (void*)(int64_t)(offer*8));
+        O2_SIGNPOST_ID_FROM_POINTER(oid, quota, (void*)(int64_t)(offer * 8));
         O2_SIGNPOST_START(quota, oid, "offers", "Offer %d has been selected.", offer);
       }
       dpStats.updateStats({static_cast<short>(ProcessingStatsId::RESOURCES_SATISFACTORY), DataProcessingStats::Op::Add, 1});
@@ -168,6 +169,7 @@ bool ComputingQuotaEvaluator::selectOffer(int task, ComputingQuotaRequest const&
     tmp.cpu += offer.cpu;
     tmp.memory += offer.memory;
     tmp.sharedMemory += offer.sharedMemory;
+    tmp.timeslices += offer.timeslices;
     offer.score = selector(offer, tmp);
     switch (offer.score) {
       case OfferScore::Unneeded:
@@ -224,7 +226,7 @@ void ComputingQuotaEvaluator::dispose(int taskId)
       continue;
     }
     if (offer.sharedMemory <= 0) {
-      O2_SIGNPOST_ID_FROM_POINTER(oid, quota, (void*)(int64_t)(oi*8));
+      O2_SIGNPOST_ID_FROM_POINTER(oid, quota, (void*)(int64_t)(oi * 8));
       O2_SIGNPOST_END(quota, oid, "offers", "Offer %d back to not needed.", oi);
       offer.valid = false;
       offer.score = OfferScore::Unneeded;
@@ -269,7 +271,7 @@ void ComputingQuotaEvaluator::handleExpired(std::function<void(ComputingQuotaOff
   /// to the driver.
   for (auto& ref : mExpiredOffers) {
     auto& offer = mOffers[ref.index];
-    O2_SIGNPOST_ID_FROM_POINTER(oid, quota, (void*)(int64_t)(ref.index*8));
+    O2_SIGNPOST_ID_FROM_POINTER(oid, quota, (void*)(int64_t)(ref.index * 8));
     if (offer.sharedMemory < 0) {
       O2_SIGNPOST_END(quota, oid, "handleExpired", "Offer %d does not have any more memory. Marking it as invalid.", ref.index);
       offer.valid = false;
diff --git a/Framework/Core/src/ResourcePolicyHelpers.cxx b/Framework/Core/src/ResourcePolicyHelpers.cxx
@@ -11,9 +11,7 @@
 
 #include "Framework/ResourcePolicyHelpers.h"
 #include "Framework/DeviceSpec.h"
-#include "ResourcesMonitoringHelper.h"
 
-#include <string>
 #include <regex>
 
 namespace o2::framework
@@ -41,6 +39,36 @@ ResourcePolicy ResourcePolicyHelpers::cpuBoundTask(char const* s, int requestedC
     [requestedCPUs](ComputingQuotaOffer const& offer, ComputingQuotaOffer const& accumulated) -> OfferScore { return accumulated.cpu >= requestedCPUs ? OfferScore::Enough : OfferScore::More; }};
 }
 
+ResourcePolicy ResourcePolicyHelpers::rateLimitedSharedMemoryBoundTask(char const* s, int requestedSharedMemory, int requestedTimeslices)
+{
+  return ResourcePolicy{
+    "ratelimited-shm-bound",
+    [matcher = std::regex(s)](DeviceSpec const& spec) -> bool {
+      return std::regex_match(spec.name, matcher);
+    },
+    [requestedSharedMemory, requestedTimeslices](ComputingQuotaOffer const& offer, ComputingQuotaOffer const& accumulated) -> OfferScore { 
+      // If we have enough memory and not enough timeslices,
+      // ignore further shared memory.
+      if (accumulated.sharedMemory >= requestedSharedMemory && offer.timeslices == 0) {
+        return OfferScore::Unneeded;
+      }
+      // If we have enough timeslices and not enough shared memory
+      // ignore further timeslices.
+      if (accumulated.timeslices >= requestedTimeslices && offer.sharedMemory == 0) {
+        return OfferScore::Unneeded;
+      }
+      // If it does not offer neither shared memory nor timeslices, mark it as unneeded.
+      if (offer.sharedMemory == 0 && offer.timeslices == 0) {
+        return OfferScore::Unneeded;
+      }
+      // We have enough to process.
+      if (accumulated.sharedMemory >= requestedSharedMemory && accumulated.timeslices >= requestedTimeslices) {
+        return OfferScore::Enough;
+      }
+      // We need more resources
+      return OfferScore::More; }};
+}
+
 ResourcePolicy ResourcePolicyHelpers::sharedMemoryBoundTask(char const* s, int requestedSharedMemory)
 {
   return ResourcePolicy{
diff --git a/Framework/Core/src/WSDriverClient.cxx b/Framework/Core/src/WSDriverClient.cxx
diff --git a/Framework/Core/src/WorkflowHelpers.cxx b/Framework/Core/src/WorkflowHelpers.cxx
