hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/utils/extrack/operations_cpu.h

@@ -0,0 +1,588 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_UTILS_EXTRACK_OPERATIONS_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_UTILS_EXTRACK_OPERATIONS_CPU_H
+
+
+#include "extrack.h"
+
+#include <iostream>
+#include <chrono>
+#include <iomanip>
+#include <ctime>
+#include <sstream>
+#include <filesystem>
+#include <fstream>
+#include <string>
+#include <algorithm>
+#include <system_error>
+#include <vector>
+#include <iterator>
+
+#include <cstdlib>
+
+#if defined(_WIN32) || defined(_WIN64)
+#include <process.h>
+#include <stdlib.h>
+#else
+#include <unistd.h>
+#endif
+#ifdef __APPLE__
+#include <crt_externs.h>
+#endif
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+#ifdef RL_TOOLS_EXTRACK_GIT_DIFF
+    namespace utils::extrack::detail {
+        // Minimal HTML escape so that adversarial commit messages embedded
+        // inside a <pre> block cannot close the tag or inject markup in the
+        // TensorBoard text plugin. Covers the full escape surface for <pre>
+        // with no user attributes: &, <, >.
+        inline std::string html_escape(const std::string& in) {
+            std::string out;
+            out.reserve(in.size() + 16);
+            for (char c : in) {
+                switch (c) {
+                    case '&': out += "&amp;"; break;
+                    case '<': out += "&lt;"; break;
+                    case '>': out += "&gt;"; break;
+                    default:  out += c;
+                }
+            }
+            return out;
+        }
+    }
+#endif
+    namespace utils::extrack::detail {
+        inline std::vector<std::string> get_command_line_arguments(){
+            std::vector<std::string> arguments;
+#if defined(_WIN32) || defined(_WIN64)
+            if(__argv != nullptr){
+                for(int argument_i = 0; argument_i < __argc; argument_i++){
+                    arguments.emplace_back(__argv[argument_i]);
+                }
+            }
+#elif defined(__APPLE__)
+            int argc = *_NSGetArgc();
+            char** argv = *_NSGetArgv();
+            if(argv != nullptr){
+                for(int argument_i = 0; argument_i < argc; argument_i++){
+                    arguments.emplace_back(argv[argument_i]);
+                }
+            }
+#else
+            std::ifstream cmdline_file("/proc/self/cmdline", std::ios::binary);
+            std::string cmdline((std::istreambuf_iterator<char>(cmdline_file)), std::istreambuf_iterator<char>());
+            std::string::size_type start = 0;
+            while(start < cmdline.size()){
+                std::string::size_type end = cmdline.find('\0', start);
+                if(end == std::string::npos){
+                    end = cmdline.size();
+                }
+                arguments.push_back(cmdline.substr(start, end - start));
+                start = end + 1;
+            }
+#endif
+            return arguments;
+        }
+        inline std::string shell_quote(const std::string& argument){
+            bool needs_quoting = argument.empty();
+            for(char c : argument){
+                bool safe = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_' || c == '-' || c == '.' || c == '/' || c == '=' || c == ':' || c == ',' || c == '+' || c == '@' || c == '%';
+                if(!safe){
+                    needs_quoting = true;
+                    break;
+                }
+            }
+            if(!needs_quoting){
+                return argument;
+            }
+            std::string quoted = "'";
+            for(char c : argument){
+                if(c == '\''){
+                    quoted += "'\\''";
+                } else {
+                    quoted += c;
+                }
+            }
+            quoted += "'";
+            return quoted;
+        }
+    }
+    // note usage examples in ./extrack.h
+    template <typename DEVICE, typename TI>
+    void init(DEVICE& device, utils::extrack::Config<TI>& config, utils::extrack::Paths& paths, typename DEVICE::index_t seed = 0){
+        if(paths.experiment.empty()){
+            utils::assert_exit(device, !config.base_path.empty(), "Extrack base path (-e,--extrack) must be set if the Extrack experiment path (--ee,--extrack-experiment) is not set.");
+            auto environment_extrack_path = std::getenv("RL_TOOLS_EXTRACK_PATH");
+            if(environment_extrack_path != nullptr && config.base_path.is_relative()){
+                config.base_path = std::filesystem::path(environment_extrack_path) / config.base_path;
+            }
+            if(config.experiment.empty()){
+
+                auto environment_extrack_experiment = std::getenv("RL_TOOLS_EXTRACK_EXPERIMENT");
+                if(environment_extrack_experiment != nullptr){
+                    config.experiment = std::string(environment_extrack_experiment);
+                } else {
+                    config.experiment = utils::extrack::get_timestamp_string();
+                }
+            }
+            paths.experiment = config.base_path / config.experiment;
+        }
+        {
+#ifdef RL_TOOLS_EXTRACK_GIT_DIFF
+            std::string commit_hash = rl_tools::utils::extrack::git::project::available ? rl_tools::utils::extrack::git::project::commit : rl_tools::utils::extrack::git::rl_tools::commit;
+            if(commit_hash.empty() || commit_hash == "unknown" || commit_hash == "not_a_repository"){
+                commit_hash = "no-hash";
+            }
+            std::string short_hash = commit_hash.substr(0, std::min<std::size_t>(7, commit_hash.size()));
+#elif defined(RL_TOOLS_COMMIT_HASH_SET)
+            std::string commit_hash = RL_TOOLS_STRINGIFY(RL_TOOLS_COMMIT_HASH);
+            std::string short_hash = commit_hash.substr(0, 7);
+#else
+            std::string short_hash = "no-hash";
+#endif
+
+            std::string setup_name = short_hash + "_" + config.name;
+            if(!config.population_variates.empty()){
+                setup_name = setup_name + "_" + config.population_variates;
+            }
+            paths.setup = paths.experiment / setup_name;
+            paths.config = paths.setup / config.population_values;
+        }
+        {
+            std::stringstream padded_seed_ss;
+            padded_seed_ss << std::setw(4) << std::setfill('0') << seed;
+            paths.seed = paths.config / padded_seed_ss.str();
+        }
+        std::cerr << "Seed: " << seed << std::endl;
+        std::cerr << "Extrack Experiment: " << std::filesystem::absolute(paths.seed).string() << std::endl;
+#ifdef RL_TOOLS_ENABLE_TENSORBOARD
+        init(device, device.logger, paths.seed);
+#endif
+#ifdef RL_TOOLS_EXTRACK_GIT_DIFF
+        // Save rl-tools library git info
+        std::filesystem::path git_library_path = paths.seed / "git" / "rl_tools";
+        std::filesystem::create_directories(git_library_path);
+        {
+            std::ofstream diff_file(git_library_path / "commit.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::commit;
+            diff_file.close();
+        }
+        {
+            std::ofstream msg_file(git_library_path / "commit_message.txt", std::ios::binary);
+            msg_file.write(
+                reinterpret_cast<const char*>(rl_tools::utils::extrack::git::rl_tools::commit_message_bytes),
+                static_cast<std::streamsize>(rl_tools::utils::extrack::git::rl_tools::commit_message_size)
+            );
+            msg_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "diff.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::diff;
+            diff_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "diff_color.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::diff_color;
+            diff_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "word_diff.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::word_diff;
+            diff_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "word_diff_color.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::word_diff_color;
+            diff_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "diff_staged.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::diff_staged;
+            diff_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "diff_staged_color.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::diff_staged_color;
+            diff_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "word_diff_staged.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::word_diff_staged;
+            diff_file.close();
+        }
+        {
+            std::ofstream diff_file(git_library_path / "word_diff_staged_color.txt");
+            diff_file << rl_tools::utils::extrack::git::rl_tools::word_diff_staged_color;
+            diff_file.close();
+        }
+        {
+            std::ofstream path_file(git_library_path / "path.txt");
+            path_file << rl_tools::utils::extrack::git::rl_tools::path;
+            path_file.close();
+        }
+        add_text(device, device.logger, "git/rl_tools/commit", (std::string("`") + rl_tools::utils::extrack::git::rl_tools::commit + "`").c_str());
+        {
+            std::string commit_message_str(
+                reinterpret_cast<const char*>(rl_tools::utils::extrack::git::rl_tools::commit_message_bytes),
+                rl_tools::utils::extrack::git::rl_tools::commit_message_size
+            );
+            std::string tb_value = std::string("<pre>") + utils::extrack::detail::html_escape(commit_message_str) + "</pre>";
+            add_text(device, device.logger, "git/rl_tools/commit_message", tb_value.c_str());
+        }
+        add_text(device, device.logger, "git/rl_tools/diff", (std::string("```diff\n") + rl_tools::utils::extrack::git::rl_tools::diff + "```").c_str());
+        add_text(device, device.logger, "git/rl_tools/word_diff", (std::string("```diff\n") + rl_tools::utils::extrack::git::rl_tools::word_diff + "```").c_str());
+        add_text(device, device.logger, "git/rl_tools/diff_staged", (std::string("```diff\n") + rl_tools::utils::extrack::git::rl_tools::diff_staged + "```").c_str());
+        add_text(device, device.logger, "git/rl_tools/word_diff_staged", (std::string("```diff\n") + rl_tools::utils::extrack::git::rl_tools::word_diff_staged + "```").c_str());
+        add_text(device, device.logger, "git/rl_tools/path", (std::string("`") + rl_tools::utils::extrack::git::rl_tools::path + "`").c_str());
+        // Save parent project git info if available
+        if(rl_tools::utils::extrack::git::project::available){
+            std::filesystem::path git_project_path = paths.seed / "git" / "project";
+            std::filesystem::create_directories(git_project_path);
+            {
+                std::ofstream diff_file(git_project_path / "commit.txt");
+                diff_file << rl_tools::utils::extrack::git::project::commit;
+                diff_file.close();
+            }
+            {
+                std::ofstream msg_file(git_project_path / "commit_message.txt", std::ios::binary);
+                msg_file.write(
+                    reinterpret_cast<const char*>(rl_tools::utils::extrack::git::project::commit_message_bytes),
+                    static_cast<std::streamsize>(rl_tools::utils::extrack::git::project::commit_message_size)
+                );
+                msg_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "diff.txt");
+                diff_file << rl_tools::utils::extrack::git::project::diff;
+                diff_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "diff_color.txt");
+                diff_file << rl_tools::utils::extrack::git::project::diff_color;
+                diff_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "word_diff.txt");
+                diff_file << rl_tools::utils::extrack::git::project::word_diff;
+                diff_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "word_diff_color.txt");
+                diff_file << rl_tools::utils::extrack::git::project::word_diff_color;
+                diff_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "diff_staged.txt");
+                diff_file << rl_tools::utils::extrack::git::project::diff_staged;
+                diff_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "diff_staged_color.txt");
+                diff_file << rl_tools::utils::extrack::git::project::diff_staged_color;
+                diff_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "word_diff_staged.txt");
+                diff_file << rl_tools::utils::extrack::git::project::word_diff_staged;
+                diff_file.close();
+            }
+            {
+                std::ofstream diff_file(git_project_path / "word_diff_staged_color.txt");
+                diff_file << rl_tools::utils::extrack::git::project::word_diff_staged_color;
+                diff_file.close();
+            }
+            {
+                std::ofstream path_file(git_project_path / "path.txt");
+                path_file << rl_tools::utils::extrack::git::project::path;
+                path_file.close();
+            }
+            add_text(device, device.logger, "git/project/commit", (std::string("`") + rl_tools::utils::extrack::git::project::commit + "`").c_str());
+            {
+                std::string commit_message_str(
+                    reinterpret_cast<const char*>(rl_tools::utils::extrack::git::project::commit_message_bytes),
+                    rl_tools::utils::extrack::git::project::commit_message_size
+                );
+                std::string tb_value = std::string("<pre>") + utils::extrack::detail::html_escape(commit_message_str) + "</pre>";
+                add_text(device, device.logger, "git/project/commit_message", tb_value.c_str());
+            }
+            add_text(device, device.logger, "git/project/diff", (std::string("```diff\n") + rl_tools::utils::extrack::git::project::diff + "```").c_str());
+            add_text(device, device.logger, "git/project/word_diff", (std::string("```diff\n") + rl_tools::utils::extrack::git::project::word_diff + "```").c_str());
+            add_text(device, device.logger, "git/project/diff_staged", (std::string("```diff\n") + rl_tools::utils::extrack::git::project::diff_staged + "```").c_str());
+            add_text(device, device.logger, "git/project/word_diff_staged", (std::string("```diff\n") + rl_tools::utils::extrack::git::project::word_diff_staged + "```").c_str());
+            add_text(device, device.logger, "git/project/path", (std::string("`") + rl_tools::utils::extrack::git::project::path + "`").c_str());
+        }
+#endif
+#ifdef RL_TOOLS_ENABLE_TENSORBOARD
+        {
+#if defined(_WIN32) || defined(_WIN64)
+            auto pid = _getpid();
+#else
+            auto pid = getpid();
+#endif
+            std::ofstream pid_file(paths.seed / "pid.txt");
+            pid_file << pid;
+            pid_file.close();
+            add_text(device, device.logger, "pid", (std::string("`") + std::to_string(pid) + "`").c_str());
+        }
+#endif
+        {
+            std::vector<std::string> arguments = utils::extrack::detail::get_command_line_arguments();
+            std::string arguments_string;
+            for(TI argument_i = 1; argument_i < arguments.size(); argument_i++){
+                if(argument_i > 1){
+                    arguments_string += " ";
+                }
+                arguments_string += utils::extrack::detail::shell_quote(arguments[argument_i]);
+            }
+            std::filesystem::create_directories(paths.seed);
+            std::ofstream cmd_file(paths.seed / "cmd.txt");
+            cmd_file << "cwd=" << std::filesystem::current_path().string() << "\n";
+            if(!arguments.empty()){
+                cmd_file << "executable=" << arguments[0] << "\n";
+            }
+            cmd_file << "args=" << arguments_string << "\n";
+            cmd_file.close();
+            std::string command = arguments.empty() ? "" : utils::extrack::detail::shell_quote(arguments[0]) + (arguments_string.empty() ? "" : " " + arguments_string);
+            add_text(device, device.logger, "cmd", (std::string("`") + command + "`").c_str());
+        }
+    }
+    template <typename DEVICE, typename TI>
+    std::filesystem::path get_step_folder(DEVICE& device, utils::extrack::Config<TI>& config, utils::extrack::Paths& paths, typename DEVICE::index_t step){
+        std::stringstream step_ss;
+        step_ss << std::setw(config.step_width) << std::setfill('0') << step;
+        std::filesystem::path step_folder = paths.seed / "steps" / step_ss.str();
+        std::filesystem::create_directories(step_folder);
+        return step_folder;
+    }
+    template <typename DEVICE>
+    std::filesystem::path get_latest_folder(DEVICE& device, const utils::extrack::Paths& paths){
+        return paths.seed / "latest";
+    }
+    template <typename DEVICE>
+    bool link_latest_artifact(DEVICE& device, const std::filesystem::path& latest_folder, const std::filesystem::path& source_path, const std::filesystem::path& step_folder, typename DEVICE::index_t step){
+        if(latest_folder.empty()){
+            return true;
+        }
+        std::error_code ec;
+        if(!std::filesystem::exists(source_path, ec)){
+            std::cerr << "Latest artifact source does not exist: " << source_path << std::endl;
+            return false;
+        }
+        std::filesystem::create_directories(latest_folder, ec);
+        if(ec){
+            std::cerr << "Failed to create latest artifact folder " << latest_folder << ": " << ec.message() << std::endl;
+            return false;
+        }
+        std::filesystem::path destination_path = latest_folder / source_path.filename();
+        std::filesystem::path temporary_path = latest_folder / (source_path.filename().string() + ".tmp");
+        std::filesystem::remove(temporary_path, ec);
+        ec.clear();
+        std::filesystem::create_hard_link(source_path, temporary_path, ec);
+        if(ec){
+            std::cerr << "Failed to hardlink latest artifact " << source_path << " -> " << destination_path << ": " << ec.message() << std::endl;
+            return false;
+        }
+        std::filesystem::remove(destination_path, ec);
+        if(ec){
+            std::cerr << "Failed to replace latest artifact " << destination_path << ": " << ec.message() << std::endl;
+            std::filesystem::remove(temporary_path, ec);
+            return false;
+        }
+        ec.clear();
+        std::filesystem::rename(temporary_path, destination_path, ec);
+        if(ec){
+            std::cerr << "Failed to move latest artifact into place " << destination_path << ": " << ec.message() << std::endl;
+            std::filesystem::remove(temporary_path, ec);
+            return false;
+        }
+        std::filesystem::path manifest_path = latest_folder / (source_path.filename().string() + ".manifest.txt");
+        std::filesystem::path temporary_manifest_path = latest_folder / (source_path.filename().string() + ".manifest.txt.tmp");
+        {
+            std::ofstream manifest_file(temporary_manifest_path);
+            if(!manifest_file){
+                std::cerr << "Failed to write latest artifact manifest: " << temporary_manifest_path << std::endl;
+                std::filesystem::remove(manifest_path, ec);
+                return false;
+            }
+            manifest_file << "step=" << step << "\n";
+            manifest_file << "path=" << step_folder.string() << "\n";
+            manifest_file << "artifact=" << source_path.string() << "\n";
+        }
+        std::filesystem::remove(manifest_path, ec);
+        if(ec){
+            std::cerr << "Failed to replace latest artifact manifest " << manifest_path << ": " << ec.message() << std::endl;
+            std::filesystem::remove(temporary_manifest_path, ec);
+            return false;
+        }
+        ec.clear();
+        std::filesystem::rename(temporary_manifest_path, manifest_path, ec);
+        if(ec){
+            std::cerr << "Failed to move latest artifact manifest into place " << manifest_path << ": " << ec.message() << std::endl;
+            std::filesystem::remove(temporary_manifest_path, ec);
+            return false;
+        }
+        return true;
+    }
+    template <typename DEVICE>
+    void parse_setup(DEVICE& device, std::string setup, utils::extrack::Path& path){
+        using TI = typename DEVICE::index_t;
+        std::vector<std::string> parts;
+        std::string part;
+        std::stringstream ss(setup);
+        while(std::getline(ss, part, '_')){
+            parts.push_back(part);
+        }
+        utils::assert_exit(device, parts.size() >= 3, "Invalid setup: " + setup);
+        path.commit = parts[0];
+        path.name = parts[1];
+        path.population_variates.clear();
+        for(TI i = 2; i < parts.size(); i++){
+            path.population_variates.push_back(parts[i]);
+        }
+    }
+    template <typename DEVICE>
+    void parse_config(DEVICE& device, std::string config, utils::extrack::Path& path){
+        using TI = typename DEVICE::index_t;
+        std::vector<std::string> parts;
+        std::string part;
+        std::stringstream ss(config);
+        while(std::getline(ss, part, '_')){
+            parts.push_back(part);
+        }
+        utils::assert_exit(device, parts.size() >= 1, "Invalid config: " + config);
+        utils::assert_exit(device, parts.size() == path.population_variates.size(), "Population variates <-> values mismatch" + config);
+        path.population_values.clear();
+        path.attributes.clear();
+        for(TI i = 0; i < parts.size(); i++){
+            path.population_values.push_back(parts[i]);
+            path.attributes[path.population_variates[i]] = parts[i];
+        }
+    }
+    template <typename DEVICE>
+    bool find_latest_run(DEVICE& device, std::filesystem::path experiments_path, utils::extrack::Path& p_query){
+        auto environment_extrack_path = std::getenv("RL_TOOLS_EXTRACK_PATH");
+        if(environment_extrack_path != nullptr && experiments_path.is_relative()){
+            experiments_path = std::filesystem::path(environment_extrack_path) / experiments_path;
+        }
+        std::vector<std::filesystem::directory_entry> experiments{std::filesystem::directory_iterator(experiments_path), std::filesystem::directory_iterator()};
+        std::sort(experiments.begin(), experiments.end());
+        std::vector<std::string> exclude_filenames = {"index.txt", "index_directories.txt", "index_files.txt"};
+        auto in = [](const std::vector<std::string>& exclude_filenames, const std::filesystem::path& path){
+            return std::find(exclude_filenames.begin(), exclude_filenames.end(), path.filename()) != exclude_filenames.end();
+        };
+        bool found = false;
+        utils::extrack::Path query = p_query;
+        utils::extrack::Path current_path;
+        for (auto& experiment : experiments){
+            if(!in(exclude_filenames, experiment.path().filename()) && experiment.is_directory() && (query.experiment.empty() || experiment.path().filename() == query.experiment)){
+                current_path.experiment = experiment.path().filename().string();
+                std::vector<std::filesystem::directory_entry> setups{std::filesystem::directory_iterator(experiment), std::filesystem::directory_iterator()};
+                std::sort(setups.begin(), setups.end());
+                for (auto& setup : setups){
+                    if(setup.is_directory()){
+                        parse_setup(device, setup.path().filename().string(), current_path);
+                        if((query.commit.empty() || query.commit == current_path.commit) && (query.name.empty() || query.name == current_path.name)){
+                            std::vector<std::filesystem::directory_entry> configs{std::filesystem::directory_iterator(setup), std::filesystem::directory_iterator()};
+                            std::sort(configs.begin(), configs.end());
+                            for (auto& config : configs){
+                                parse_config(device, config.path().filename().string(), current_path);
+                                bool matches_query = true;
+                                for (const auto& [key, value] : query.attributes){
+                                    auto it = current_path.attributes.find(key);
+                                    if (it == current_path.attributes.end() || it->second != value) {
+                                        matches_query = false;
+                                        break;
+                                    }
+                                }
+                                if(config.is_directory() && matches_query){
+                                    std::vector<std::filesystem::directory_entry> seeds{std::filesystem::directory_iterator(config), std::filesystem::directory_iterator()};
+                                    std::sort(seeds.begin(), seeds.end());
+                                    for (auto& seed : seeds){
+                                        if(seed.is_directory()){
+                                            if(query.seed.empty() || std::stoull(query.seed) == std::stoull(seed.path().filename())){
+                                                bool has_checkpoint = false;
+                                                auto steps_path = seed.path() / "steps";
+                                                if (std::filesystem::exists(steps_path) && std::filesystem::is_directory(steps_path)){
+                                                    std::vector<std::filesystem::directory_entry> steps{std::filesystem::directory_iterator(steps_path), std::filesystem::directory_iterator()};
+                                                    std::sort(steps.begin(), steps.end());
+                                                    for (auto& step : steps){
+                                                        if(step.is_directory()){
+                                                            if(query.step.empty() || std::stoull(query.step) == std::stoull(step.path().filename())){
+                                                                if(std::filesystem::exists(step.path() / "checkpoint.h5")){
+                                                                    has_checkpoint = true;
+                                                                    current_path.checkpoint_path = step.path() / "checkpoint.h5";
+                                                                }
+                                                            }
+                                                        }
+                                                    }
+                                                }
+                                                current_path.seed = seed.path().filename().string();
+                                                current_path.path = seed.path();
+                                                if(!query.require_checkpoint || has_checkpoint){
+                                                    p_query = current_path;
+                                                    found = true;
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        if(!found){
+            std::cout << "Could not find run for query:" << std::endl;
+            if(!query.experiment.empty()){
+                std::cout << "\tExperiment: " << query.experiment << std::endl;
+            }
+            if(!query.commit.empty()){
+                std::cout << "\tCommit: " << query.commit << std::endl;
+            }
+            if(!query.name.empty()){
+                std::cout << "\tName: " << query.name << std::endl;
+            }
+            if (!query.attributes.empty()){
+                std::cout << "\tAttributes: " << std::endl;
+                for (const auto& [key, value] : query.attributes){
+                    std::cout << "\t\t" << key << ": " << value << std::endl;
+                }
+            }
+            if(!query.seed.empty()){
+                std::cout << "\tSeed: " << query.seed << std::endl;
+            }
+            if(!query.step.empty()){
+                std::cout << "\tStep: " << query.step << std::endl;
+            }
+            if(query.require_checkpoint){
+                std::cout << "\tRequire checkpoint: true" << std::endl;
+            }
+        }
+        return found;
+    }
+    template <typename DEVICE>
+    bool find_latest_checkpoint(DEVICE& device, utils::extrack::Path& p_query){
+        utils::extrack::Path query = p_query;
+        std::vector<std::filesystem::directory_entry> steps{std::filesystem::directory_iterator(query.path / "steps"), std::filesystem::directory_iterator()};
+        std::sort(steps.begin(), steps.end());
+        bool found = false;
+        utils::extrack::Path current_path = query;
+        for (auto& step : steps){
+            if(step.is_directory()){
+                if(std::filesystem::exists(step.path() / "checkpoint.h5")){
+                    current_path.step = step.path().filename();
+                    current_path.checkpoint_path = step.path() / "checkpoint.h5";
+                    p_query = current_path;
+                    found = true;
+                }
+            }
+        }
+        return found;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif