Skip to content

Commit

Permalink
Refactor: Add DeviceFileMeta and DeviceMetaInConfig
Browse files Browse the repository at this point in the history 
Signed-off-by: Li Junlin <[email protected]>
  • Loading branch information
@L-Xiafeng
L-Xiafeng committed Nov 14, 2024
1 parent 98f00ed commit b9a5088
Show file tree
Hide file tree
Showing 5 changed files with 50 additions and 47 deletions.
4 changes: 2 additions & 2 deletions src/Craned/CgroupManager.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1093,8 +1093,8 @@ bool CgroupV1::SetDeviceAccess(const std::unordered_set<SlotId> &devices,
if (set_mknod) op += "m";
std::vector<std::string> deny_limits;
for (const auto &[_, this_device] : Craned::g_this_node_device) {
if (!devices.contains(this_device->dev_id)) {
for (const auto &dev_meta : this_device->device_metas) {
if (!devices.contains(this_device->slot_id)) {
for (const auto &dev_meta : this_device->device_file_metas) {
deny_limits.emplace_back(fmt::format("{} {}:{} {}", dev_meta.op_type,
dev_meta.major, dev_meta.minor,
op));
Expand Down
50 changes: 22 additions & 28 deletions src/Craned/Craned.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,10 +81,8 @@ void ParseConfig(int argc, char** argv) {
}

std::string config_path = parsed_args["config"].as<std::string>();
using DeviceMeta = std::tuple<std::string /*name*/, std::string /*type*/,
std::vector<std::string> /*path*/,
std::optional<std::string> /*EnvInjector*/>;
std::unordered_map<std::string, std::vector<DeviceMeta>> each_node_device;
std::unordered_map<std::string, std::vector<Craned::DeviceMetaInConfig>>
each_node_device;
if (std::filesystem::exists(config_path)) {
try {
YAML::Node config = YAML::LoadFile(config_path);
Expand Down Expand Up @@ -285,7 +283,7 @@ void ParseConfig(int argc, char** argv) {
} else
std::exit(1);

std::vector<DeviceMeta> devices;
std::vector<Craned::DeviceMetaInConfig> devices;
if (node["gres"]) {
for (auto gres_it = node["gres"].begin();
gres_it != node["gres"].end(); ++gres_it) {
Expand All @@ -308,10 +306,10 @@ void ParseConfig(int argc, char** argv) {
std::exit(1);
}
for (const auto& device_path : device_path_list) {
devices.push_back(std::make_tuple(
device_name, device_type, std::vector{device_path},
!env_injector.empty() ? std::optional(env_injector)
: std::nullopt));
devices.push_back(
{device_name, device_type, std::vector{device_path},
!env_injector.empty() ? std::optional(env_injector)
: std::nullopt});
}
}
if (gres_node["DeviceFileList"] &&
Expand All @@ -327,18 +325,18 @@ void ParseConfig(int argc, char** argv) {
device_name, device_type);
std::exit(1);
}
devices.push_back(std::make_tuple(
device_name, device_type,
std::vector(device_path_list.begin(),
device_path_list.end()),
!env_injector.empty() ? std::optional(env_injector)
: std::nullopt));
devices.push_back({device_name, device_type,
std::vector(device_path_list.begin(),
device_path_list.end()),
!env_injector.empty()
? std::optional(env_injector)
: std::nullopt});
}
}
if (!device_file_configured) {
CRANE_ERROR(
"gres {}:{} device DeviceFileRegex or DeviceFileList not "
"configured",
"At least one of DeviceFileRegex or DeviceFileList must be "
"configured for GRES {}:{} device.",
device_name, device_type);
}
}
Expand Down Expand Up @@ -536,30 +534,26 @@ void ParseConfig(int argc, char** argv) {
auto node_res = g_config.CranedRes.at(g_config.Hostname);
auto& devices = each_node_device[g_config.Hostname];
for (auto& dev_arg : devices) {
std::string name, type;
std::optional<std::string> env_injector;
std::vector<std::string> path;
std::tie(name, type, path, env_injector) = dev_arg;

auto& [name, type, path_vec, env_injector] = dev_arg;
auto env_injector_enum =
Craned::GetDeviceEnvInjectorFromStr(env_injector);
Craned::GetDeviceEnvInjectorFromStr(dev_arg.EnvInjectorStr);
if (env_injector_enum == Craned::InvalidInjector) {
CRANE_ERROR("Invalid injector type:{} for device {}.",
env_injector.value_or("EmptyVal"), path);
env_injector.value_or("EmptyVal"), path_vec);
std::exit(1);
}

std::unique_ptr dev = Craned::DeviceManager::ConstructDevice(
name, type, path, env_injector_enum);
name, type, path_vec, env_injector_enum);
if (!dev->Init()) {
CRANE_ERROR("Access Device {} failed.", static_cast<std::string>(*dev));
std::exit(1);
}

dev->dev_id = dev->device_metas.front().path;
dev->slot_id = dev->device_file_metas.front().path;
node_res->dedicated_res.name_type_slots_map[dev->name][dev->type].emplace(
dev->dev_id);
Craned::g_this_node_device[dev->dev_id] = std::move(dev);
dev->slot_id);
Craned::g_this_node_device[dev->slot_id] = std::move(dev);
}
each_node_device.clear();
}
Expand Down
18 changes: 9 additions & 9 deletions src/Craned/DeviceManager.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,30 +46,30 @@ BasicDevice::BasicDevice(const std::string& device_name,
const std::vector<std::string>& device_path,
DeviceEnvInjector env_injector)
: name(device_name), type(device_type), env_injector(env_injector) {
device_metas.reserve(device_path.size());
device_file_metas.reserve(device_path.size());
for (const auto& dev_path : device_path) {
device_metas.emplace_back(DeviceMeta{dev_path, 0, 0, 0});
device_file_metas.emplace_back(DeviceFileMeta{dev_path, 0, 0, 0});
}
}

bool BasicDevice::Init() {
for (auto& device_meta : device_metas) {
for (auto& device_file_meta : device_file_metas) {
const auto& device_major_minor_optype_option =
DeviceManager::GetDeviceFileMajorMinorOpType(device_meta.path);
DeviceManager::GetDeviceFileMajorMinorOpType(device_file_meta.path);
if (!device_major_minor_optype_option.has_value()) return false;
const auto& device_major_minor_optype =
device_major_minor_optype_option.value();

device_meta.major = std::get<0>(device_major_minor_optype);
device_meta.minor = std::get<1>(device_major_minor_optype);
device_meta.op_type = std::get<2>(device_major_minor_optype);
device_file_meta.major = std::get<0>(device_major_minor_optype);
device_file_meta.minor = std::get<1>(device_major_minor_optype);
device_file_meta.op_type = std::get<2>(device_major_minor_optype);
}
return true;
}

BasicDevice::operator std::string() const {
std::vector<std::string> device_files;
for (const auto& device_meta : device_metas) {
for (const auto& device_meta : device_file_metas) {
device_files.push_back(device_meta.path);
}
return fmt::format("{}:{}:{}", name, type,
Expand Down Expand Up @@ -116,7 +116,7 @@ DeviceManager::GetDevEnvListByResInNode(

std::unordered_map<DeviceEnvInjector, int> injector_count_map;
for (const auto& [_, device] : g_this_node_device) {
if (!all_res_slots.contains(device->dev_id)) continue;
if (!all_res_slots.contains(device->slot_id)) continue;
if (device->env_injector == DeviceEnvInjector::CommonDevice) continue;
injector_count_map[device->env_injector]++;
}
Expand Down
24 changes: 16 additions & 8 deletions src/Craned/DeviceManager.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,8 +54,22 @@ constexpr std::array<std::string_view,
DeviceEnvInjector GetDeviceEnvInjectorFromStr(
const std::optional<std::string>& str);

struct DeviceMetaInConfig {
std::string name;
std::string type;
std::vector<std::string> path;
std::optional<std::string> EnvInjectorStr;
};

struct DeviceFileMeta {
std::string path;
unsigned int major;
unsigned int minor;
char op_type;
};

struct BasicDevice {
std::string dev_id;
SlotId slot_id;

// e.g GPU
std::string name;
Expand All @@ -64,13 +78,7 @@ struct BasicDevice {

DeviceEnvInjector env_injector;

struct DeviceMeta {
std::string path;
unsigned int major;
unsigned int minor;
char op_type;
};
std::vector<DeviceMeta> device_metas;
std::vector<DeviceFileMeta> device_file_metas;

BasicDevice(const std::string& device_name, const std::string& device_type,
const std::vector<std::string>& device_path,
Expand Down
1 change: 1 addition & 0 deletions src/Utilities/PublicHeader/include/crane/PublicHeader.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,7 @@ using PartitionId = std::string;
using CranedId = std::string;
using cpu_t = fpm::fixed_24_8;

// TODO: refactor SlotId, it should not be a string of file path.
// Device path. e.g. /dev/nvidia0
using SlotId = std::string;

Expand Down

0 comments on commit b9a5088

Please sign in to comment.