daggy/libdaggy/tests/unit_executor_slurmexecutor.cpp

#include <unistd.h>

#include <catch2/catch.hpp>
#include <filesystem>
#include <fstream>
#include <iostream>

#include "daggy/Serialization.hpp"
#include "daggy/Utilities.hpp"
#include "daggy/executors/task/SlurmTaskExecutor.hpp"

namespace fs = std::filesystem;

#ifdef DAGGY_ENABLE_SLURM

TEST_CASE("slurm environment", "[slurm][slurm_env]")
{
  daggy::executors::task::SlurmTaskExecutor ex;

  daggy::ConfigValues defaultJobValues{{"minCPUs", "1"},
                                       {"minMemoryMB", "100"},
                                       {"minTmpDiskMB", "0"},
                                       {"priority", "1"},
                                       {"timeLimitSeconds", "200"},
                                       {"userID", std::to_string(getuid())},
                                       {"workDir", fs::current_path().string()},
                                       {"tmpDir", fs::current_path().string()}};
}

TEST_CASE("slurm_execution", "[slurm][slurm_executor]")
{
  daggy::executors::task::SlurmTaskExecutor ex;

  daggy::ConfigValues defaultJobValues{{"minCPUs", "1"},
                                       {"minMemoryMB", "100"},
                                       {"minTmpDiskMB", "0"},
                                       {"priority", "1"},
                                       {"timeLimitSeconds", "200"},
                                       {"userID", std::to_string(getuid())},
                                       {"workDir", fs::current_path().string()},
                                       {"tmpDir", fs::current_path().string()}};

  SECTION("Simple Run")
  {
    daggy::Task task{.job{
        {"command", std::vector<std::string>{"/bin/echo", "abc", "123"}}}};

    task.job.merge(defaultJobValues);

    REQUIRE(ex.validateTaskParameters(task.job));

    auto recFuture = ex.execute(0, "command", task);
    auto rec       = recFuture.get();

    REQUIRE(rec.rc == 0);
    REQUIRE(rec.outputLog.size() >= 6);
    REQUIRE(rec.errorLog.empty());
  }

  SECTION("Simple run with environment")
  {
    // Create the shell script
    auto scriptFile = fs::current_path() / "slurm_simple_env.sh";

    if (fs::exists(scriptFile))
      fs::remove_all(scriptFile);

    std::string valOne   = "funky_times";
    std::string valTwo   = "bleep_bloop";
    std::string valThree = "SOME-DATE";

    std::string testParams{R"({"DATE": ")" + valThree + R"("})"};
    auto params = daggy::configFromJSON(testParams);

    std::ofstream ofh(scriptFile);
    ofh << R"(#!/bin/bash
              echo "${DAGGY_TEST_VAR}"
              echo "${DAGGY_TEST_VAR2}"
              echo "${DATE}"
              )";
    ofh.close();
    fs::permissions(scriptFile, fs::perms::owner_all,
                    fs::perm_options::replace);

    daggy::Task rawTask{.job{
        {"command",
         daggy::executors::task::SlurmTaskExecutor::Command{
             scriptFile.string()}},
        {"environment", std::vector<std::string>{"DAGGY_TEST_VAR=" + valOne,
                                                 "DAGGY_TEST_VAR2=" + valTwo,
                                                 "DATE={{DATE}}"}}}};
    rawTask.job.merge(defaultJobValues);

    auto tasks =
        daggy::expandTaskSet(daggy::TaskSet{{"cmd", rawTask}}, ex, params);
    REQUIRE(tasks.size() == 1);
    auto task = tasks.at("cmd_0");

    REQUIRE(ex.validateTaskParameters(task.job));

    auto recFuture = ex.execute(0, "command", task);
    auto rec       = recFuture.get();

    REQUIRE(rec.rc == 0);
    REQUIRE(rec.outputLog.size() >= 6);
    REQUIRE(rec.outputLog.find(valOne) != std::string::npos);
    REQUIRE(rec.outputLog.find(valTwo) != std::string::npos);
    REQUIRE(rec.outputLog.find(valThree) != std::string::npos);
    REQUIRE(rec.errorLog.empty());

    if (fs::exists(scriptFile))
      fs::remove_all(scriptFile);
  }

  SECTION("Simple Run using commandString")
  {
    daggy::Task task{.job{{"commandString", R"(/bin/echo "abc 123")"}}};
    task.job.merge(defaultJobValues);

    REQUIRE(ex.validateTaskParameters(task.job));

    auto recFuture = ex.execute(0, "command", task);
    auto rec       = recFuture.get();

    REQUIRE(rec.rc == 0);
    REQUIRE(rec.outputLog.size() >= 6);
    REQUIRE(rec.errorLog.empty());
  }

  SECTION("Error Run")
  {
    daggy::Task task{
        .job{{"command", daggy::executors::task::SlurmTaskExecutor::Command{
                             "/bin/expr", "1", "+", "+"}}}};
    task.job.merge(defaultJobValues);

    auto recFuture = ex.execute(0, "command", task);
    auto rec       = recFuture.get();

    REQUIRE(rec.rc != 0);
    REQUIRE(rec.errorLog.size() >= 20);
    REQUIRE(rec.outputLog.empty());
  }

  SECTION("Killing a long task")
  {
    daggy::Task task{
        .job{{"command", daggy::executors::task::SlurmTaskExecutor::Command{
                             "/bin/sleep", "30"}}}};
    task.job.merge(defaultJobValues);

    auto recFuture = ex.execute(0, "command", task);
    ex.stop(0, "command");
    auto rec = recFuture.get();

    REQUIRE(rec.rc == 9);
    REQUIRE(rec.errorLog.empty());
    REQUIRE(rec.outputLog.empty());
    REQUIRE(rec.executorLog == "Job cancelled by user.\n");
  }

  SECTION("Large Output")
  {
    const std::vector<std::string> BIG_FILES{"/share/dict/linux.words",
                                             "/share/dict/cracklib-small",
                                             "/etc/ssh/moduli"};

    for (const auto &bigFile : BIG_FILES) {
      if (!std::filesystem::exists(bigFile))
        continue;

      daggy::Task task{
          .job{{"command", daggy::executors::task::SlurmTaskExecutor::Command{
                               "/bin/cat", bigFile}}}};
      task.job.merge(defaultJobValues);

      auto recFuture = ex.execute(0, "command", task);
      auto rec       = recFuture.get();

      REQUIRE(rec.rc == 0);
      REQUIRE(rec.outputLog.size() == std::filesystem::file_size(bigFile));
      REQUIRE(rec.errorLog.empty());
      break;
    }
  }

  SECTION("Parameter Expansion")
  {
    std::string testParams{R"({"DATE": ["2021-05-06", "2021-05-07" ]})"};
    auto params = daggy::configFromJSON(testParams);

    std::string taskJSON =
        R"({"B": {"job": {"command": ["/bin/echo", "{{DATE}}"]}, "children": ["C"]}})";
    auto tasks = daggy::tasksFromJSON(taskJSON, defaultJobValues);

    auto result = daggy::expandTaskSet(tasks, ex, params);
    REQUIRE(result.size() == 2);
  }

  SECTION("Build with expansion")
  {
    std::string testParams{
        R"({"DATE": ["2021-05-06", "2021-05-07" ], "SOURCE": "name"})"};
    auto params = daggy::configFromJSON(testParams);
    std::string testTasks =
        R"({"A": {"job": {"command": ["/bin/echo", "A"]}, "children": ["B"]}, "B": {"job": {"command": ["/bin/echo", "B", "{{SOURCE}}", "{{DATE}}"]}, "children": ["C"]}, "C": {"job": {"command": ["/bin/echo", "C"]}}})";
    auto tasks = daggy::expandTaskSet(
        daggy::tasksFromJSON(testTasks, defaultJobValues), ex, params);
    REQUIRE(tasks.size() == 4);
  }

  SECTION("Build with expansion using parents instead of children")
  {
    std::string testParams{
        R"({"DATE": ["2021-05-06", "2021-05-07" ], "SOURCE": "name"})"};
    auto params = daggy::configFromJSON(testParams);
    std::string testTasks =
        R"({"A": {"job": {"command": ["/bin/echo", "A"]}}, "B": {"job": {"command": ["/bin/echo", "B", "{{SOURCE}}", "{{DATE}}"]}, "parents": ["A"]}, "C": {"job": {"command": ["/bin/echo", "C"]}, "parents": ["A"]}})";
    auto tasks = daggy::expandTaskSet(
        daggy::tasksFromJSON(testTasks, defaultJobValues), ex, params);

    REQUIRE(tasks.size() == 4);
  }
}
#endif