Add CDI for WSLC GPU

src/linux/init/WSLCInit.cpp

@@ -32,6 +32,10 @@ Module Name:
 #include <pty.h>
 #include "mountutilcpp.h"
 #include <filesystem>
+#include <iostream>
+#include "JsonUtils.h"
+#include "cdi_schema.h"
+#include "lxfsshares.h"
 
 extern int InitializeLogging(bool SetStderr, wil::LogFunction* ExceptionCallback) noexcept;
 
@@ -63,19 +67,105 @@ struct WSLCState
 
 static WSLCState g_state;
 
-int CreateCaptureCrashSymlink()
+void WriteWslcCdiSpec()
 try
 {
-    THROW_LAST_ERROR_IF(symlink("/init", "/" LX_INIT_WSL_CAPTURE_CRASH) < 0);
+    wsl::shared::cdi::DeviceNode dxg{};
+    dxg.path = "/dev/dxg";
+    dxg.permissions = "rwm";
+
+    wsl::shared::cdi::Mount libs{};
+    libs.hostPath = LXSS_LIB_PATH;
+    libs.containerPath = LXSS_LIB_PATH;
+    libs.options = {"ro", "rbind"};
+
+    wsl::shared::cdi::Mount drivers{};
+    drivers.hostPath = LXSS_GPU_DRIVERS_PATH;
+    drivers.containerPath = LXSS_GPU_DRIVERS_PATH;
+    drivers.options = {"ro", "rbind"};
+
+    wsl::shared::cdi::Hook hook{};
+    hook.hookName = "createContainer";
+    hook.path = "/" LX_INIT_WSLC_GPU_HOOK;
+    hook.args = {LX_INIT_WSLC_GPU_HOOK};
+
+    wsl::shared::cdi::Device gpu{};
+    gpu.name = "gpu";
+    gpu.containerEdits.deviceNodes.push_back(std::move(dxg));
+    gpu.containerEdits.mounts.push_back(std::move(libs));
+    gpu.containerEdits.mounts.push_back(std::move(drivers));
+    gpu.containerEdits.hooks.push_back(std::move(hook));
+
+    wsl::shared::cdi::Spec spec{};
+    spec.cdiVersion = "0.6.0";
+    spec.kind = LX_WSLC_CDI_KIND;
+    spec.devices.push_back(std::move(gpu));
+
+    THROW_LAST_ERROR_IF(UtilMkdirPath("/etc/cdi", 0755) < 0);
+    THROW_LAST_ERROR_IF(
+        WriteToFile("/etc/cdi/microsoft.com-wslc.json", nlohmann::json(spec).dump().c_str(), O_WRONLY | O_CLOEXEC | O_CREAT | O_TRUNC) < 0);
+}
+CATCH_LOG()
+
+void WriteDockerDaemonConfig()
+try
+{
+    constexpr auto c_daemonConfigPath = "/etc/docker/daemon.json";
+
+    THROW_ERRNO_IF(EEXIST, std::filesystem::exists(c_daemonConfigPath));
+
+    nlohmann::json config = nlohmann::json::object();
+    config["features"]["cdi"] = true;
+
+    THROW_LAST_ERROR_IF(UtilMkdirPath("/etc/docker", 0755) < 0);
+    THROW_LAST_ERROR_IF(WriteToFile(c_daemonConfigPath, config.dump().c_str(), O_WRONLY | O_CLOEXEC | O_CREAT | O_TRUNC) < 0);
+}
+CATCH_LOG()
+
+int WslcGpuHookEntry()
+try
+{
+    // OCI runtime hooks receive the container state as JSON on stdin.
+    const std::string stateJson{std::istreambuf_iterator<char>(std::cin), {}};
+    THROW_ERRNO_IF(EINVAL, stateJson.empty());
+
+    const auto state = nlohmann::json::parse(stateJson);
+    const std::filesystem::path bundle = state.at("bundle").get<std::string>();
+    THROW_ERRNO_IF(EINVAL, !bundle.is_absolute());
+
+    // Read the OCI spec's root.path from <bundle>/config.json. This is either an absolute path to
+    // the overlay-merged rootfs or a path relative to the bundle directory.
+    const auto spec = nlohmann::json::parse(UtilReadFileContent((bundle / "config.json").native()));
+    std::filesystem::path rootfsPath = spec.at("root").at("path").get<std::string>();
+    if (rootfsPath.is_relative())
+    {
+        rootfsPath = bundle / rootfsPath;
+    }
+
+    rootfsPath = std::filesystem::canonical(rootfsPath);
+    THROW_ERRNO_IF(EINVAL, rootfsPath == "/");
+    THROW_ERRNO_IF(ENOTDIR, !std::filesystem::is_directory(rootfsPath));
+
+    const auto confDir = rootfsPath / "etc/ld.so.conf.d";
+    const auto confPath = confDir / "ld.wsl.conf";
+
+    THROW_LAST_ERROR_IF(UtilMkdirPath(confDir.c_str(), 0755) < 0);
+    THROW_LAST_ERROR_IF(WriteToFile(confPath.c_str(), LXSS_LIB_PATH "\n", O_WRONLY | O_CLOEXEC | O_CREAT | O_TRUNC | O_NOFOLLOW) < 0);
+
+    // chroots into the rootfs and uses the container's own /etc/ld.so.conf chain,
+    // writing /etc/ld.so.cache inside the container.
+    const char* const ldArgv[] = {LDCONFIG_COMMAND, "-r", rootfsPath.c_str(), nullptr};
+    THROW_LAST_ERROR_IF(UtilCreateProcessAndWait(ldArgv[0], ldArgv) < 0);
 
     return 0;
 }
 CATCH_RETURN_ERRNO()
 
 void WSLCEnableCrashDumpCollection()
 {
-    if (CreateCaptureCrashSymlink() < 0)
+    if (symlink("/init", "/" LX_INIT_WSL_CAPTURE_CRASH) < 0 && errno != EEXIST)
     {
+        LOG_ERROR("symlink(/init, /" LX_INIT_WSL_CAPTURE_CRASH ") failed {}", errno);
         return;
     }
 
@@ -619,8 +709,12 @@ void HandleMessageImpl(
         {
             THROW_LAST_ERROR_IF(Chroot(target) < 0);
 
-            // Recreate the crash dump symlink inside the new root.
-            CreateCaptureCrashSymlink();
+            // Recreate the /init symlinks inside the new root.
+            THROW_LAST_ERROR_IF(symlink("/init", "/" LX_INIT_WSL_CAPTURE_CRASH) < 0 && errno != EEXIST);
+            THROW_LAST_ERROR_IF(symlink("/init", "/" LX_INIT_WSLC_GPU_HOOK) < 0 && errno != EEXIST);
+
+            WriteWslcCdiSpec();
+            WriteDockerDaemonConfig();
         }
 
         response.Result = 0;

src/linux/init/init.cpp

@@ -162,6 +162,8 @@ wil::unique_fd UnmarshalConsoleFromServer(int MessageFd, LXBUS_IPC_CONSOLE_ID Co
 
 int WslInitWatcher(int Argc, char** Argv);
 
+int WslcGpuHookEntry();
+
 int WslEntryPoint(int Argc, char* Argv[])
 {
     //
@@ -228,6 +230,10 @@ int WslEntryPoint(int Argc, char* Argv[])
         {
             ExitCode = WslInitWatcher(Argc, Argv);
         }
+        else if (strcmp(BaseName, LX_INIT_WSLC_GPU_HOOK) == 0)
+        {
+            ExitCode = WslcGpuHookEntry();
+        }
         else
         {
             // Handle the special case for import result messages, everything else is sent to the binfmt interpreter.

src/linux/init/util.cpp

@@ -3338,7 +3338,7 @@ uint16_t UtilWinAfToLinuxAf(uint16_t WinAddressFamily)
     return LinuxAddressFamily;
 }
 
-int WriteToFile(const char* Path, const char* Content, int permissions)
+int WriteToFile(const char* Path, const char* Content, int OpenFlags, int Permissions)
 
 /*++
 
@@ -3352,14 +3352,18 @@ Routine Description:
 
     Content - Supplies the content to be written to the file.
 
+    Permissions - Supplies the file mode used when O_CREAT causes the file to be created.
+
+    OpenFlags - Supplies the flags passed to open(). Defaults to O_WRONLY | O_CLOEXEC | O_CREAT.
+
 Return Value:
 
     0 on success, -1 on failure.
 
 --*/
 
 {
-    wil::unique_fd Fd{open(Path, (O_WRONLY | O_CLOEXEC | O_CREAT), permissions)};
+    wil::unique_fd Fd{open(Path, OpenFlags, Permissions)};
     if (!Fd)
     {
         int errnoPrev = errno;

src/linux/init/util.h

@@ -313,6 +313,6 @@ std::string UtilReadFileContent(std::string_view path);
 
 uint16_t UtilWinAfToLinuxAf(uint16_t AddressFamily);
 
-int WriteToFile(const char* Path, const char* Content, int permissions = 0644);
+int WriteToFile(const char* Path, const char* Content, int OpenFlags = O_WRONLY | O_CLOEXEC | O_CREAT, int Permissions = 0644);
 
 int ProcessCreateProcessMessage(wsl::shared::Transaction& Transaction, gsl::span<gsl::byte> Buffer);

src/shared/inc/cdi_schema.h

@@ -0,0 +1,74 @@
+/*++
+
+Copyright (c) Microsoft. All rights reserved.
+
+Module Name:
+
+    cdi_schema.h
+
+Abstract:
+
+    Schema for Container Device Interface (CDI) specs.
+    See https://github.com/cncf-tags/container-device-interface/blob/main/SPEC.md
+
+--*/
+
+#pragma once
+
+#include "JsonUtils.h"
+
+namespace wsl::shared::cdi {
+
+struct DeviceNode
+{
+    std::string path;
+    std::string permissions;
+
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(DeviceNode, path, permissions);
+};
+
+struct Mount
+{
+    std::string hostPath;
+    std::string containerPath;
+    std::vector<std::string> options;
+
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(Mount, hostPath, containerPath, options);
+};
+
+struct Hook
+{
+    std::string hookName;
+    std::string path;
+    std::vector<std::string> args;
+
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(Hook, hookName, path, args);
+};
+
+struct ContainerEdits
+{
+    std::vector<DeviceNode> deviceNodes;
+    std::vector<Mount> mounts;
+    std::vector<Hook> hooks;
+
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(ContainerEdits, deviceNodes, mounts, hooks);
+};
+
+struct Device
+{
+    std::string name;
+    ContainerEdits containerEdits;
+
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(Device, name, containerEdits);
+};
+
+struct Spec
+{
+    std::string cdiVersion;
+    std::string kind;
+    std::vector<Device> devices;
+
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(Spec, cdiVersion, kind, devices);
+};
+
+} // namespace wsl::shared::cdi

src/shared/inc/lxfsshares.h

@@ -23,6 +23,7 @@ typedef struct _LXSS_SHARED_DIRECTORY
 
 #define LXSS_LIB_PREFIX "/usr/lib/wsl"
 #define LXSS_LIB_PATH LXSS_LIB_PREFIX "/lib"
+#define LXSS_GPU_DRIVERS_PATH LXSS_LIB_PREFIX "/drivers"
 #define LXSS_GPU_DRIVERS_SHARE "drivers"
 #define LXSS_GPU_LIB_SHARE "lib"
 #define LXSS_GPU_INBOX_LIB_SHARE LXSS_GPU_LIB_SHARE "_inbox"
@@ -32,4 +33,4 @@ typedef struct _LXSS_SHARED_DIRECTORY
 // Shared directories for GPU compute support.
 //
 
-constexpr LXSS_SHARED_DIRECTORY g_gpuShares[] = {{LXSS_GPU_DRIVERS_SHARE, LXSS_LIB_PREFIX "/drivers"}, {LXSS_GPU_LIB_SHARE, LXSS_LIB_PATH}};
+constexpr LXSS_SHARED_DIRECTORY g_gpuShares[] = {{LXSS_GPU_DRIVERS_SHARE, LXSS_GPU_DRIVERS_PATH}, {LXSS_GPU_LIB_SHARE, LXSS_LIB_PATH}};

src/shared/inc/lxinitshared.h

@@ -247,6 +247,11 @@ Module Name:
 
 #define LX_INIT_WSL_INIT_WATCHER "init-watcher"
 
+#define LX_INIT_WSLC_GPU_HOOK "wsl-gpu-hook"
+
+#define LX_WSLC_CDI_KIND "microsoft.com/wslc"
+#define LX_WSLC_GPU_CDI_DEVICE LX_WSLC_CDI_KIND "=gpu"
+
 //
 // WSL2-specific environment variables.
 //

src/windows/inc/docker_schema.h

@@ -215,6 +215,14 @@ struct Ulimit
     NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(Ulimit, Name, Soft, Hard);
 };
 
+struct DeviceRequest
+{
+    std::string Driver;
+    std::vector<std::string> DeviceIDs;
+
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(DeviceRequest, Driver, DeviceIDs);
+};
+
 struct HostConfig
 {
     std::vector<Mount> Mounts;
@@ -230,14 +238,15 @@ struct HostConfig
     // the field — so we don't bother with std::optional here.
     std::int64_t ShmSize{};
     std::optional<std::vector<DeviceMapping>> Devices;
+    std::optional<std::vector<DeviceRequest>> DeviceRequests;
 
     // Per-container resource limits. 0 means "no limit" (Docker default).
     std::int64_t Memory{};
     std::int64_t NanoCpus{};
     std::optional<std::vector<Ulimit>> Ulimits;
 
     NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(
-        HostConfig, Mounts, PortBindings, NetworkMode, Init, Dns, DnsSearch, DnsOptions, Binds, Tmpfs, Devices, ShmSize, Memory, NanoCpus, Ulimits);
+        HostConfig, Mounts, PortBindings, NetworkMode, Init, Dns, DnsSearch, DnsOptions, Binds, Tmpfs, Devices, DeviceRequests, ShmSize, Memory, NanoCpus, Ulimits);
 };
 
 struct EndpointSettings

src/windows/wslcsession/WSLCContainer.cpp

@@ -494,29 +494,6 @@ void ProcessAdditionalNetworks(
     }
 }
 
-void ConfigureLdPathForGpu(std::vector<std::string>& Env)
-{
-    static constexpr std::string_view ldLibraryPathPrefix = "LD_LIBRARY_PATH=";
-    auto it = std::ranges::find_if(Env, [](const std::string& e) { return e.starts_with(ldLibraryPathPrefix); });
-
-    if (it != Env.end())
-    {
-        // If the user already has an LD_LIBRARY_PATH, append the GPU library paths to it.
-        auto ldPath = it->substr(ldLibraryPathPrefix.size());
-        if (!ldPath.empty() && !ldPath.ends_with(":"))
-        {
-            it->append(":");
-        }
-
-        it->append(WSLCVirtualMachine::c_gpuLibrariesPath);
-    }
-    else
-    {
-        // Otherwise create a new entry.
-        Env.emplace_back(std::format("LD_LIBRARY_PATH={}", WSLCVirtualMachine::c_gpuLibrariesPath));
-    }
-}
-
 } // namespace
 
 ContainerPortMapping::ContainerPortMapping(VMPortMapping&& VmMapping, uint16_t ContainerPort) :
@@ -1203,11 +1180,6 @@ void WSLCContainerImpl::Exec(const WSLCProcessOptions* Options, LPCSTR DetachKey
         request.DetachKeys = DetachKeys;
     }
 
-    if (WI_IsFlagSet(m_containerFlags, WSLCContainerFlagsGpu))
-    {
-        ConfigureLdPathForGpu(request.Env);
-    }
-
     try
     {
         auto result = m_dockerClient.CreateExec(m_id, request);
@@ -1595,17 +1567,8 @@ std::unique_ptr<WSLCContainerImpl> WSLCContainerImpl::Create(
             !virtualMachine.FeatureEnabled(WslcFeatureFlagsGPU),
             "WSLCContainerFlagsGpu requires GPU support enabled on the session");
 
-        if (!request.HostConfig.Binds.has_value())
-        {
-            request.HostConfig.Binds = std::vector<std::string>{};
-        }
-
-        request.HostConfig.Binds->push_back(std::format("{0}:{0}:ro", WSLCVirtualMachine::c_gpuLibrariesPath));
-        request.HostConfig.Binds->push_back(std::format("{0}:{0}:ro", WSLCVirtualMachine::c_gpuDriversPath));
-
-        request.HostConfig.Devices = {{"/dev/dxg", "/dev/dxg", "rwm"}};
-
-        ConfigureLdPathForGpu(request.Env);
+        // Request the WSL GPU device via CDI.
+        request.HostConfig.DeviceRequests = std::vector<common::docker_schema::DeviceRequest>{{"cdi", {LX_WSLC_GPU_CDI_DEVICE}}};
     }
 
     // Prepare port mappings from container options.