TracccTritonInitializer.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/

 
#include "Python.h"
 
#include "TracccTritonInitializer.h"
 
#include "CxxUtils/checker_macros.h"
 
#include <dlfcn.h>
 
#include <atomic>
#include <cstddef>
#include <mutex>
#include <stdexcept>
#include <string>
 
// Embedded Athena application.
#include "GaudiKernel/Bootstrap.h"
#include "GaudiKernel/IAppMgrUI.h"
#include "GaudiKernel/IAlgManager.h"
#include "GaudiKernel/IAlgorithm.h"
#include "GaudiKernel/ISvcLocator.h"
#include "GaudiKernel/SmartIF.h"
#include "GaudiKernel/StateMachine.h"
 
 
namespace triton { namespace backend { namespace traccc {
 
// TracccTritonBootstrap.bootstrap() brings up the embedded Athena application
// with the device chain configured by TracccTritonDeviceRecoCfg, scheduled
// inside an AthSequencer
namespace {
    void promoteSelfToGlobal() {
        Dl_info info{};
        if (!dladdr(reinterpret_cast<void*>(&promoteSelfToGlobal), &info) ||
            !info.dli_fname) {
            throw std::runtime_error(
                "TracccTritonInitializer: dladdr() could not determine this "
                "library's own path");
        }
        if (!dlopen(info.dli_fname, RTLD_NOW | RTLD_GLOBAL | RTLD_NOLOAD)) {
            throw std::runtime_error(
                std::string("TracccTritonInitializer: failed to promote '") +
                info.dli_fname + "' to RTLD_GLOBAL: " + dlerror());
        }
    }
}  // namespace
 
struct TracccTritonInitializer::Impl {
    Config config;
    SmartIF<IAppMgrUI> app;
    SmartIF<ISvcLocator> svcLocator;
    bool ready = false;
    std::atomic<std::size_t> nextSlot{0};
    PyThreadState* pyMainThreadState = nullptr;

    SmartIF<IAlgorithm> algorithm(const std::string& name) const {
        IAlgManager* algMgr = svcLocator.as<IAlgManager>();
        if (!algMgr) {
            throw std::runtime_error(
                "TracccTritonInitializer: no IAlgManager available");
        }
        SmartIF<IAlgorithm> alg = algMgr->algorithm(name, /*createIf*/ false);
        if (!alg) {
            throw std::runtime_error(
                "TracccTritonInitializer: could not retrieve algorithm '" +
                name + "' from the embedded Athena application");
        }
        return alg;
    }

    static std::string keyFromDict(PyObject* dict, const char* name) {
        PyObject* item = PyDict_GetItemString(dict, name);  // borrowed
        if (!item || !PyUnicode_Check(item)) {
            throw std::runtime_error(
                std::string("TracccTritonInitializer: TracccTritonBootstrap."
                            "bootstrap() returned no string entry '") + name +
                "' in its StoreGate key dictionary");
        }
        const char* value = PyUnicode_AsUTF8(item);
        if (!value || !*value) {
            throw std::runtime_error(
                std::string("TracccTritonInitializer: empty StoreGate key '") +
                name + "' returned by TracccTritonBootstrap.bootstrap()");
        }
        return value;
    }

    void bootstrapPython() {
        const std::size_t nSlots = config.nSlots;
        const std::string& sequenceName = config.sequenceName;
        if (!Py_IsInitialized()) {
            if (!dlopen(TRACCC_TRITON_PYTHON_LIBRARY, RTLD_NOW | RTLD_GLOBAL)) {
                throw std::runtime_error(
                    std::string(
                        "TracccTritonInitializer: failed to re-load '"
                        TRACCC_TRITON_PYTHON_LIBRARY "' with RTLD_GLOBAL: ") +
                    dlerror());
            }
 
            Py_Initialize();
            if (!Py_IsInitialized()) {
                throw std::runtime_error(
                    "TracccTritonInitializer: Py_Initialize() failed");
            }
        }
 
        PyObject* module =
            PyImport_ImportModule("TracccTritonBackend.TracccTritonBootstrap");
        if (!module) {
            PyErr_Print();
            throw std::runtime_error(
                "TracccTritonInitializer: failed to import "
                "TracccTritonBackend.TracccTritonBootstrap -- check "
                "PYTHONPATH");
        }
 
        PyObject* func = PyObject_GetAttrString(module, "bootstrap");
        Py_DECREF(module);
        if (!func || !PyCallable_Check(func)) {
            Py_XDECREF(func);
            PyErr_Print();
            throw std::runtime_error(
                "TracccTritonInitializer: TracccTritonBootstrap.bootstrap "
                "not found or not callable");
        }
 
        PyObject* result =
            PyObject_CallFunction(func, "ns", static_cast<Py_ssize_t>(nSlots),
                                  sequenceName.c_str());
        Py_DECREF(func);
        if (!result) {
            PyErr_Print();
            throw std::runtime_error(
                "TracccTritonInitializer: TracccTritonBootstrap.bootstrap() "
                "raised a Python exception (see stderr above)");
        }
        if (!PyDict_Check(result)) {
            Py_DECREF(result);
            throw std::runtime_error(
                "TracccTritonInitializer: TracccTritonBootstrap.bootstrap() "
                "did not return a dictionary of StoreGate keys");
        }
        try {
            config.cellsKey = keyFromDict(result, "cells");
            config.measurementsKey = keyFromDict(result, "measurements");
            config.clustersKey = keyFromDict(result, "clusters");
            config.tracksKey = keyFromDict(result, "tracks");
        } catch (...) {
            Py_DECREF(result);
            throw;
        }
        Py_DECREF(result);
 
        pyMainThreadState = PyEval_SaveThread();
    }
};
 
TracccTritonInitializer::TracccTritonInitializer()
    : m_impl(std::make_unique<Impl>()) {}
 
TracccTritonInitializer::~TracccTritonInitializer() {
    try {
        finalize();
    } catch (...) {
    }
}
 
TracccTritonInitializer& TracccTritonInitializer::instance() {
    // The singleton is mutable, but every method that touches its state
    // serialises on its own mutex, so sharing it between Triton's worker
    // threads is safe.
    static TracccTritonInitializer theInstance ATLAS_THREAD_SAFE;
    return theInstance;
}
 
void TracccTritonInitializer::initialize(const Config& config) {
    static std::mutex initMutex;
    std::lock_guard<std::mutex> lock(initMutex);
 
    if (m_impl->ready) {
        // Gaudi is a process singleton; a second initialize() with a
        // different device id cannot be honored.
        if (config.deviceId != m_impl->config.deviceId) {
            throw std::runtime_error(
                "TracccTritonInitializer: already initialized on device " +
                std::to_string(m_impl->config.deviceId) +
                ", cannot re-initialize on device " +
                std::to_string(config.deviceId) +
                ". Pin the model to a single GPU in instance_group.");
        }
        if (config.nSlots != m_impl->config.nSlots) {
            throw std::runtime_error(
                "TracccTritonInitializer: already initialized with " +
                std::to_string(m_impl->config.nSlots) +
                " event slots, cannot re-initialize with " +
                std::to_string(config.nSlots));
        }
        return;
    }
 
    m_impl->config = config;
 
    promoteSelfToGlobal();
 
    // Fills in the StoreGate keys on m_impl->config from the python configuration
    m_impl->bootstrapPython();
 
    // Fetch the ApplicationMgr
    m_impl->app = Gaudi::createApplicationMgr();
    if (!m_impl->app) {
        throw std::runtime_error(
            "TracccTritonInitializer: Gaudi::createApplicationMgr() "
            "returned null after the Python bootstrap");
    }
    // bootstrap() takes the application all the way to RUNNING
    if (m_impl->app->FSMState() != Gaudi::StateMachine::RUNNING) {
        throw std::runtime_error(
            "TracccTritonInitializer: application manager is not "
            "RUNNING after the Python bootstrap (FSM state " +
            std::to_string(static_cast<int>(m_impl->app->FSMState())) + ")");
    }
 
    m_impl->svcLocator = SmartIF<ISvcLocator>(m_impl->app);
    if (!m_impl->svcLocator) {
        throw std::runtime_error(
            "TracccTritonInitializer: could not obtain ISvcLocator");
    }
 
    // Sanity-check that the device chain the Runner will execute exists.
    m_impl->algorithm(m_impl->config.sequenceName);
 
    m_impl->ready = true;
}
 
void TracccTritonInitializer::finalize() {
    if (!m_impl->ready) return;
 
    if (m_impl->app) {
        // finalize() expects the INITIALIZED state, so the RUNNING state the
        // bootstrap left the application in has to be unwound first.
        m_impl->app->stop().ignore();
        m_impl->app->finalize().ignore();
        m_impl->app->terminate().ignore();
    }
 
    m_impl->svcLocator = nullptr;
    m_impl->app = nullptr;
    m_impl->ready = false;
}
 
bool TracccTritonInitializer::isReady() const {
    return m_impl->ready;
}
 
const TracccTritonInitializer::Config&
TracccTritonInitializer::config() const {
    return m_impl->config;
}
 
std::size_t
TracccTritonInitializer::acquireSlot() {
    if (!m_impl->ready) {
        throw std::runtime_error(
            "TracccTritonInitializer::acquireSlot: not initialized");
    }
    const std::size_t slot = m_impl->nextSlot.fetch_add(1);
    if (slot >= m_impl->config.nSlots) {
        throw std::runtime_error(
            "TracccTritonInitializer: Triton created more model instances "
            "than the " + std::to_string(m_impl->config.nSlots) +
            " event slot(s) the application was configured with");
    }
    return slot;
}
 
int TracccTritonInitializer::deviceId() const {
    return m_impl->config.deviceId;
}
 
ISvcLocator& TracccTritonInitializer::serviceLocator() const {
    if (!m_impl->svcLocator) {
        throw std::runtime_error(
            "TracccTritonInitializer::serviceLocator: not initialized");
    }
    return *m_impl->svcLocator;
}
 
}}}  // namespace triton::backend::traccc