From 59cb311a65dc0409355a324db962ea87d5a37c71 Mon Sep 17 00:00:00 2001 From: specture724 Date: Sun, 6 Sep 2026 10:59:12 +0800 Subject: [PATCH 1/2] feat: support Docker clients with a host guard bridge (cherry picked from commit ecd29d4b8617eceab1637e45aba0445ec0f0506c) --- .gitignore | 1 + README.md | 2 + deploy/canhazgpu-docker-guard.service | 16 ++ docs/docker.md | 160 ++++++++++++++++ internal/cli/guard.go | 26 +++ internal/cli/root.go | 43 +++++ internal/cli/run.go | 3 + internal/cli/run_supervisor_args_test.go | 4 +- internal/gpu/allocation.go | 55 +++++- internal/gpu/cancel.go | 28 +++ internal/gpu/docker_test.go | 47 +++++ internal/hostbridge/bridge_test.go | 117 ++++++++++++ internal/hostbridge/client.go | 110 +++++++++++ internal/hostbridge/identity.go | 152 +++++++++++++++ internal/hostbridge/peer_linux.go | 24 +++ internal/hostbridge/peer_other.go | 12 ++ internal/hostbridge/server.go | 143 ++++++++++++++ internal/redis_client/client.go | 21 ++- internal/types/types.go | 4 + mkdocs.yml | 3 +- scripts/test-docker.py | 228 +++++++++++++++++++++++ 21 files changed, 1184 insertions(+), 15 deletions(-) create mode 100644 deploy/canhazgpu-docker-guard.service create mode 100644 docs/docker.md create mode 100644 internal/gpu/docker_test.go create mode 100644 internal/hostbridge/bridge_test.go create mode 100644 internal/hostbridge/client.go create mode 100644 internal/hostbridge/identity.go create mode 100644 internal/hostbridge/peer_linux.go create mode 100644 internal/hostbridge/peer_other.go create mode 100644 internal/hostbridge/server.go create mode 100644 scripts/test-docker.py diff --git a/.gitignore b/.gitignore index bc426da..5ae759e 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ canhazgpu .idea/ # Temporary and cache files +__pycache__/ .DS_Store *.tmp *.temp diff --git a/README.md b/README.md index fd8e243..9f17d2a 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,8 @@ In shared development environments with multiple GPUs, researchers and developer You peacefully share a host but want a helper to avoid accidental conflicts. +For development inside Docker with a host guard, see the [Docker setup guide](docs/docker.md). + - You have a single host with GPUs (NVIDIA or AMD) shared by multiple users - You all log in and run commands manually for development and/or testing - You can still talk to each other about playing nice and sharing your GPUs diff --git a/deploy/canhazgpu-docker-guard.service b/deploy/canhazgpu-docker-guard.service new file mode 100644 index 0000000..57a9c79 --- /dev/null +++ b/deploy/canhazgpu-docker-guard.service @@ -0,0 +1,16 @@ +[Unit] +Description=canhazgpu host guard with Docker clients +Wants=network-online.target redis.service +After=network-online.target redis.service docker.service + +[Service] +Type=simple +User=root +RuntimeDirectory=canhazgpu +RuntimeDirectoryMode=0755 +ExecStart=/usr/local/bin/canhazgpu guard --listen /run/canhazgpu/host.sock --enforce +Restart=on-failure +RestartSec=5s + +[Install] +WantedBy=multi-user.target diff --git a/docs/docker.md b/docs/docker.md new file mode 100644 index 0000000..738fbb6 --- /dev/null +++ b/docs/docker.md @@ -0,0 +1,160 @@ +# 宿主机守护与 Docker 开发环境 + +在裸金属宿主机运行一个 root 权限的 `canhazgpu guard`，开发者在自己的 +Docker 中照常使用 `run`、`reserve`、`release`、`queue`、`cancel`、`status`。 +容器可以使用 root 账户及独立的 PID、网络命名空间。 + +宿主机通过 Unix socket 的内核 peer credentials 获取调用进程的**宿主机 PID**， +从宿主机 cgroup 找到 Docker 容器，再读取 `canhazgpu.owner` 标签或管理员的归属配置。 +因此两个 root 容器仍属于两个不同用户。任务及 supervisor 在容器内执行， +Redis 中记录宿主机 PID；查询设备及取消任务由宿主机执行。 +`status` 中的进程 PID 始终使用宿主机编号，内外一致。 + +## 1. 宿主机设置 + +安装同一版本的二进制，并使用现有 Redis 及设备池。尚未初始化时才执行 `admin`： + +```bash +canhazgpu admin --gpus 8 --provider nvidia +sudo canhazgpu guard --listen /run/canhazgpu/host.sock --enforce +``` + +选择机器实际支持的 provider。在含 Ascend 支持的分支上使用 `--provider ascend`。 +不要用 `admin --force` 覆盖正在使用的设备池。已有 guard 时先停止原实例； +同一个 Redis 数据库仍只允许一个 guard。 + +开机启动可使用仓库中的 `deploy/canhazgpu-docker-guard.service`，按实际情况 +设置 Redis 参数及 guard 策略。它使用 `RuntimeDirectory` 管理 socket 目录。 +如已有 guard 服务，只需为其增加 `--listen /run/canhazgpu/host.sock`，并加入 +`RuntimeDirectory=canhazgpu` 和 `RuntimeDirectoryMode=0755`。 + +宿主机与容器客户端都会自动发现 `/run/canhazgpu/host.sock`。 +其他路径用 `--host-socket PATH` 或 `CANHAZGPU_HOST_SOCKET=PATH` 指定。 +`--host-socket off` 可显式使用本机直接连接模式。 +配置了 socket 后，连接失败会直接报错，不会转而使用容器的 root 身份、本地 PID 或本地 Redis。 + +## 2. 创建开发容器 + +在自己的宿主机账户下创建容器，把账户写入标签，挂载 **socket 所在目录**及二进制： + +```bash +docker run -it --name my-dev \ + --label "canhazgpu.owner=$(id -un)" \ + -e CANHAZGPU_HOST_SOCKET=/run/canhazgpu/host.sock \ + --mount type=bind,src=/run/canhazgpu,dst=/run/canhazgpu,readonly \ + --mount type=bind,src=/usr/local/bin/canhazgpu,dst=/usr/local/bin/canhazgpu,readonly \ + --gpus all \ + YOUR_DEVELOPMENT_IMAGE bash +``` + +这里的 `--gpus all` 适用于 NVIDIA Container Toolkit。AMD/Ascend 请保留开发环境 +原有的设备和驱动挂载。**容器必须暴露整个共享设备池，并保持与宿主机相同的设备编号**； +目前不转换仅暴露部分设备或重排设备后的编号。Ascend 通常需要所有 `/dev/davinciN`、 +`davinci_manager`、`devmm_svm`、`hisi_hdc` 和对应的宿主机驱动目录。 +canhazgpu 设置运行时可见设备变量，不代替 Docker 的设备映射。 + +Redis 连接通过同一个 socket 转发到守护进程使用的 Redis，数据库编号和内存阈值 +也从守护进程获取；无需容器联网或暴露 Redis TCP 端口，也无需挂载 Docker socket +或宿主机 `/proc`，无需 `--pid=host`。 +挂载整个 socket 目录可让客户端在守护进程重启后找到新 socket。 + +如果使用 `sudo docker run`，请在 sudo 前取得真实账户： + +```bash +owner=$(id -un) +sudo docker run --label "canhazgpu.owner=$owner" ... +``` + +标签必须是存在的非 root 宿主机账户。标签声明的是容器归属，不是容器内 `$USER`。 +Docker 的创建者账户不能从容器内 root 自动推断。 + +## 3. 已有容器无需重建 + +管理员可创建 JSON 文件，将**完整的容器 ID**映射为宿主机账户： + +```bash +docker inspect --format '{{.Id}}' existing-dev +``` + +例如 `/etc/canhazgpu/docker-owners.json`： + +```json +{ + "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef": "alice" +} +``` + +```bash +sudo canhazgpu guard --listen /run/canhazgpu/host.sock \ + --docker-owners /etc/canhazgpu/docker-owners.json --enforce +``` + +映射优先于标签，启动时加载；修改后重启 guard。它也适用于宿主机默认 Docker +daemon 无法 inspect 的 rootless 容器。已有容器仍须有 socket 目录和二进制挂载； +Docker 不能向运行中的容器补充普通 bind mount，缺少挂载时需按原配置重建。 +无法确认归属的容器不会被默认为宿主机 root，客户端会提示补齐配置； +设备扫描会显示 `unknown`，guard 不会终止身份不明的进程。 + +## 4. 容器内外使用 + +```bash +canhazgpu run --gpus 1 -- python train.py +canhazgpu run --gpu-ids 0 --idle-timeout 0 -- python +canhazgpu status +canhazgpu status --json +canhazgpu queue +canhazgpu cancel TASK_ID +canhazgpu reserve --gpus 1 --duration 2h +canhazgpu release +``` + +终端、标准输入输出、命令退出码和 Ctrl-C 由本地任务保留；supervisor 使用容器内 PID +监控任务，并通过 bridge 维持预约和执行空闲检查。容器内外均按宿主机账户进行排队、 +预约归属、使用归属及取消检查。`--user` 仍可设置显示名，实际归属不会改变。 +容器 root 不具有宿主机管理员身份，不能通过 `cancel --force` 取消其他用户的任务。 +需要管理员取消时，在宿主机使用 `sudo canhazgpu cancel TASK_ID --force`。 + +guard 的违规检测、宽限时间、警告和 SIGINT → SIGTERM → SIGKILL 策略保持一致。 +进程警告可经宿主机 `/proc//fd/2` 到达容器终端。取消容器任务时宿主机同时处理 +主进程和该任务设备上属于该用户的设备进程。若任务直接作为容器 PID 1 运行，建议为镜像 +配置 init（Docker 安装了 docker-init 时可用 `--init`），以便正常处理信号及回收子进程。 + +## 信任范围与运维 + +这仍是共享开发机上的协作调度系统：获准连接 socket 的客户端能访问同一 Redis， +标签/映射用于记账，不是隔离恶意 Docker 管理员的安全边界。 +socket 默认允许本机用户连接；目录应由宿主机管理员控制。只给受信开发容器挂载它。 +需要限制客户端时可通过目录访问权限控制。 + +守护正常退出会移除 socket；若被 SIGKILL 后残留 socket，手工启动会拒绝覆盖， +应确认没有正在使用它的守护后清理。推荐 systemd 的 RuntimeDirectory 自动管理生命周期。 +容器与宿主机应使用同一版 canhazgpu。 + +## 验证 + +```bash +CGO_ENABLED=0 go build -o build/canhazgpu-docker . +python3 scripts/test-docker.py --image YOUR_LOCAL_IMAGE +``` + +集成脚本需要普通宿主机账户、可访问的 Docker、免密 sudo、Redis 及含 `sh`/`sleep` +的本地镜像。它创建三个真实 root 容器（标签归属当前账户与 `nobody`，以及一个通过配置归属的容器）、独立 Redis、 +临时 root 守护和模拟设备查询。它验证归属、宿主机 PID、状态一致、内外取消、 +跨用户拒绝、归属映射、守护重启、排队取消、设备变量、退出码、自动释放、失联报错，以及 guard 对真实 +容器进程的终止。模拟设备不会使用真实加速卡；结束会清理测试资源。 + +实现依据：[Docker labels](https://docs.docker.com/engine/manage-resources/labels/)、 +[Docker process isolation](https://docs.docker.com/engine/containers/run/)、 +[Linux Unix socket peer credentials](https://man7.org/linux/man-pages/man7/unix.7.html)。 + +### 本机验证记录（2026-09-06） + +- Go 完整测试（含独立 Redis 集成测试）及 race 检查通过，`go vet ./...` 通过。 +- Docker 18.09、Linux 5.10、arm64：无网络、独立 PID、非 privileged 的真实 root + 容器通过上述模拟设备集成测试。 +- 在 NPU 分支上使用真实 Ascend 910B1、驱动 25.5.0 和本机 + `quay.io/ascend/vllm-omni:v0.28.0` 镜像验证：先预约物理卡 7，容器内 + `canhazgpu run --gpu-ids 7` 运行 torch_npu 张量任务；实际设备进程占用 130 MB， + 内外状态均归属 `ajhou`，宿主机 PID 一致，容器内取消后进程和预约均释放。 + 此机器的驱动需要 privileged 容器才能成功初始化；该验证仍使用独立 PID 和无网络模式。 + 这属于设备运行环境要求，bridge 本身已在非 privileged 容器中通过验证。 diff --git a/internal/cli/guard.go b/internal/cli/guard.go index 8895f6f..920eb29 100644 --- a/internal/cli/guard.go +++ b/internal/cli/guard.go @@ -3,10 +3,12 @@ package cli import ( "context" "fmt" + "os" "strings" "time" "github.com/russellb/canhazgpu/internal/gpu" + "github.com/russellb/canhazgpu/internal/hostbridge" "github.com/russellb/canhazgpu/internal/redis_client" "github.com/russellb/canhazgpu/internal/types" "github.com/russellb/canhazgpu/internal/utils" @@ -69,6 +71,8 @@ Example usage: func init() { defaults := gpu.DefaultGuardConfig() + guardCmd.Flags().String("listen", "", "Serve container clients on this Unix socket (e.g. /run/canhazgpu/host.sock)") + guardCmd.Flags().String("docker-owners", "", "Host JSON file mapping full Docker container IDs to host accounts") guardCmd.Flags().String("interval", utils.FormatDurationShort(defaults.Interval), "How often to scan the GPUs") guardCmd.Flags().Bool("once", false, "Run a single scan and exit (for cron)") @@ -105,6 +109,10 @@ func runGuard(ctx context.Context) error { } config := getConfig() + config.ContainerOwners, err = hostbridge.LoadOwners(viper.GetString("guard.docker-owners")) + if err != nil { + return fmt.Errorf("Docker owner mappings: %w", err) + } client := redis_client.NewClient(config) defer func() { if err := client.Close(); err != nil { @@ -124,6 +132,24 @@ func runGuard(ctx context.Context) error { engine := gpu.NewAllocationEngine(client, config) notifier := gpu.NewMultiNotifier(viper.GetStringSlice("guard.channels"), viper.GetString("guard.log-file")) guard := gpu.NewGuard(engine, client, config, settings, notifier) + if socket := viper.GetString("guard.listen"); socket != "" { + if os.Geteuid() != 0 { + return fmt.Errorf("the host bridge must run as root to manage container processes") + } + if viper.GetBool("guard.once") { + return fmt.Errorf("--listen requires a continuous guard, not --once") + } + server := &hostbridge.Server{ + Resolver: &hostbridge.Resolver{Owners: config.ContainerOwners}, Config: config, Usage: engine.DetectUsage, + Cancel: func(ctx context.Context, ref, user string, force bool) (any, error) { + return engine.CancelTask(ctx, ref, user, force) + }, + } + if err := server.Start(ctx, socket); err != nil { + return fmt.Errorf("start host bridge: %w", err) + } + defer server.Close() + } if viper.GetBool("guard.once") { pass, err := guard.RunOnce(ctx) diff --git a/internal/cli/root.go b/internal/cli/root.go index badb419..776584f 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -4,8 +4,10 @@ import ( "context" "fmt" "os" + "path/filepath" "strings" + "github.com/russellb/canhazgpu/internal/hostbridge" "github.com/russellb/canhazgpu/internal/types" "github.com/russellb/canhazgpu/internal/utils" "github.com/spf13/cobra" @@ -29,9 +31,11 @@ to requested GPUs while automatically handling cleanup when jobs complete or cra ) func init() { + rootCmd.PersistentPreRunE = prepareHostBridge cobra.OnInitialize(initConfig) // Global flags + rootCmd.PersistentFlags().String("host-socket", "", "Host guard Unix socket (auto-detected at /run/canhazgpu/host.sock; 'off' disables)") rootCmd.PersistentFlags().StringVar(&configFile, "config", "", "config file (default is $HOME/.canhazgpu.yaml)") rootCmd.PersistentFlags().String("redis-host", "localhost", "Redis host") rootCmd.PersistentFlags().Int("redis-port", 6379, "Redis port") @@ -96,6 +100,7 @@ func initConfig() { bindAllFlags() config = &types.Config{ + HostSocket: viper.GetString("host-socket"), RedisHost: viper.GetString("redis.host"), RedisPort: viper.GetInt("redis.port"), RedisDB: viper.GetInt("redis.db"), @@ -156,6 +161,9 @@ func walkCommands(cmd *cobra.Command, fn func(*cobra.Command)) { } func getCurrentUser() string { + if config != nil && config.HostUser != "" { + return config.HostUser + } if user := os.Getenv("USER"); user != "" { return user } @@ -164,3 +172,38 @@ func getCurrentUser() string { } return "unknown" } + +func prepareHostBridge(cmd *cobra.Command, args []string) error { + cfg := getConfig() + if cmd.Name() == "guard" { + if cfg.HostSocket != "" && cfg.HostSocket != "off" { + return fmt.Errorf("guard must run on the host; remove --host-socket / CANHAZGPU_HOST_SOCKET") + } + cfg.HostSocket = "" + return nil + } + if cfg.HostSocket == "off" { + cfg.HostSocket = "" + return nil + } + if cfg.HostSocket == "" { + // The mounted directory survives a guard restart even while its socket + // is absent. Do not switch container clients to local root/PIDs then. + if info, err := os.Stat(filepath.Dir(hostbridge.DefaultSocket)); err == nil && info.IsDir() { + cfg.HostSocket = hostbridge.DefaultSocket + } + } + if cfg.HostSocket == "" { + return nil + } + identity, err := (hostbridge.Client{Socket: cfg.HostSocket}).Identity(cmd.Context()) + if err != nil { + return err + } + cfg.HostPID, cfg.HostUser = identity.PID, identity.User + cfg.RedisDB, cfg.MemoryThreshold = identity.RedisDB, identity.MemoryThreshold + if cmd.Name() == "admin" && !identity.Admin { + return fmt.Errorf("admin through the host bridge requires host root") + } + return nil +} diff --git a/internal/cli/run.go b/internal/cli/run.go index a84a52f..84cd1aa 100644 --- a/internal/cli/run.go +++ b/internal/cli/run.go @@ -345,6 +345,9 @@ func buildSupervisorArgs(executable string, config *types.Config, gpuList string "--redis-port", strconv.Itoa(config.RedisPort), "--redis-db", strconv.Itoa(config.RedisDB), } + if config.HostSocket != "" { + args = append(args, "--host-socket", config.HostSocket) + } if timeout != "" && timeout != "0" { args = append(args, "--timeout", timeout) } diff --git a/internal/cli/run_supervisor_args_test.go b/internal/cli/run_supervisor_args_test.go index 7af437b..4c6caf2 100644 --- a/internal/cli/run_supervisor_args_test.go +++ b/internal/cli/run_supervisor_args_test.go @@ -13,7 +13,7 @@ import ( // line it talks to the default instance and immediately loses the reservation // it was spawned to hold. func TestBuildSupervisorArgs_ForwardsRedisSettings(t *testing.T) { - config := &types.Config{RedisHost: "redis.example", RedisPort: 6380, RedisDB: 15} + config := &types.Config{RedisHost: "redis.example", RedisPort: 6380, RedisDB: 15, HostSocket: "/run/canhazgpu/host.sock", HostPID: 99999} args := buildSupervisorArgs("/usr/local/bin/canhazgpu", config, "1,2", "alice", 4242, "2h", 30*time.Minute) line := strings.Join(args, " ") @@ -26,6 +26,8 @@ func TestBuildSupervisorArgs_ForwardsRedisSettings(t *testing.T) { assert.Contains(t, line, "--gpus 1,2") assert.Contains(t, line, "--user alice") assert.Contains(t, line, "--pid 4242") + assert.Contains(t, line, "--host-socket /run/canhazgpu/host.sock") + assert.NotContains(t, line, "--pid 99999", "the supervisor must monitor the local PID, not the stored host PID") assert.Contains(t, line, "--timeout 2h") assert.Contains(t, line, "--idle-timeout 30m") diff --git a/internal/gpu/allocation.go b/internal/gpu/allocation.go index a4cc4f0..6194ae3 100644 --- a/internal/gpu/allocation.go +++ b/internal/gpu/allocation.go @@ -12,6 +12,7 @@ import ( "time" "github.com/google/uuid" + "github.com/russellb/canhazgpu/internal/hostbridge" "github.com/russellb/canhazgpu/internal/redis_client" "github.com/russellb/canhazgpu/internal/types" "github.com/russellb/canhazgpu/internal/utils" @@ -35,38 +36,41 @@ type AllocationEngine struct { usageMu sync.Mutex usageCache map[int]*types.GPUUsage usageCached time.Time + owners *hostbridge.Resolver } func NewAllocationEngine(client *redis_client.Client, config *types.Config) *AllocationEngine { return &AllocationEngine{ client: client, config: config, + owners: &hostbridge.Resolver{Owners: config.ContainerOwners}, } } func (ae *AllocationEngine) detectGPUUsage(ctx context.Context) (map[int]*types.GPUUsage, error) { ae.usageMu.Lock() + defer ae.usageMu.Unlock() + // A host bridge serves many clients concurrently. Share one driver query + // per cache interval instead of launching simultaneous npu/nvidia-smi calls. if ae.usageCache != nil && time.Since(ae.usageCached) < usageCacheTTL { - cached := ae.usageCache - ae.usageMu.Unlock() - return cached, nil + return ae.usageCache, nil } - ae.usageMu.Unlock() usage, err := ae.queryGPUUsage(ctx) if err != nil { return nil, err } - ae.usageMu.Lock() ae.usageCache = usage ae.usageCached = time.Now() - ae.usageMu.Unlock() return usage, nil } func (ae *AllocationEngine) queryGPUUsage(ctx context.Context) (map[int]*types.GPUUsage, error) { + if ae.config.HostSocket != "" { + return (hostbridge.Client{Socket: ae.config.HostSocket}).Usage(ctx) + } providerName, err := ae.client.GetAvailableProvider(ctx) if err != nil { return nil, fmt.Errorf("failed to get cached provider information: %v", err) @@ -84,7 +88,40 @@ func (ae *AllocationEngine) queryGPUUsage(ctx context.Context) (map[int]*types.G pm = NewProviderManagerFromNames([]string{providerName}) } - return pm.DetectAllGPUUsageWithoutChecks(ctx) + usage, err := pm.DetectAllGPUUsageWithoutChecks(ctx) + if err != nil { + return nil, err + } + for _, device := range usage { + device.Users = make(map[string]bool) + for i := range device.Processes { + process := &device.Processes[i] + // Preserve the provider's native owner when procfs is unavailable. + // A recognized container without a mapping must never inherit root's + // guard exemption or another container's reservation. + identity, err := ae.owners.Resolve(ctx, process.PID) + if identity.ContainerID != "" { + process.User = "unknown" + if err == nil { + process.User = identity.User + } + } + device.Users[process.User] = true + } + } + return usage, nil +} + +// DetectUsage returns the host view used by both guard scans and bridge clients. +func (ae *AllocationEngine) DetectUsage(ctx context.Context) (map[int]*types.GPUUsage, error) { + return ae.detectGPUUsage(ctx) +} + +func (ae *AllocationEngine) taskPID() int { + if ae.config.HostPID > 0 { + return ae.config.HostPID + } + return os.Getpid() } // NewTaskID returns a short handle for a reservation, shown by 'queue' and @@ -107,7 +144,7 @@ func (ae *AllocationEngine) AllocateGPUs(ctx context.Context, request *types.All request.TaskID = NewTaskID() } if request.ReservationType == types.ReservationTypeRun && request.PID == 0 { - request.PID = os.Getpid() + request.PID = ae.taskPID() } // Best effort maintenance so this request sees an up-to-date pool: due @@ -770,7 +807,7 @@ func (ae *AllocationEngine) createQueueEntry(request *QueuedAllocationRequest) * EnqueueTime: types.FlexibleTime{Time: now}, LastHeartbeat: types.FlexibleTime{Time: now}, // The waiting process cleans up its own entry when signalled - PID: os.Getpid(), + PID: ae.taskPID(), } if request.ExpiryTime != nil { diff --git a/internal/gpu/cancel.go b/internal/gpu/cancel.go index a3cb2ae..a207aff 100644 --- a/internal/gpu/cancel.go +++ b/internal/gpu/cancel.go @@ -7,6 +7,7 @@ import ( "syscall" "time" + "github.com/russellb/canhazgpu/internal/hostbridge" "github.com/russellb/canhazgpu/internal/types" ) @@ -37,6 +38,11 @@ const ( // held by a live process are signalled so they clean up after themselves - the // reservation is only released directly when no process is left to do it. func (ae *AllocationEngine) CancelTask(ctx context.Context, ref string, actualUser string, force bool) (*CancelResult, error) { + if ae.config.HostSocket != "" { + var result CancelResult + err := (hostbridge.Client{Socket: ae.config.HostSocket}).Call(ctx, hostbridge.Request{Operation: "cancel", TaskID: ref, Force: force}, &result) + return &result, err + } ref = strings.TrimSpace(ref) if ref == "" { return nil, fmt.Errorf("no task ID given") @@ -74,6 +80,9 @@ func (ae *AllocationEngine) cancelQueuedTask(ctx context.Context, entry *types.Q if !force && owner != actualUser { return nil, fmt.Errorf("task %s belongs to %s (use --force to cancel it anyway)", entry.ShortID(), entry.User) } + if err := ae.checkTaskProcessOwner(ctx, entry.PID, owner); err != nil { + return nil, err + } result := &CancelResult{ TaskID: entry.ShortID(), @@ -100,6 +109,9 @@ func (ae *AllocationEngine) cancelRunningTask(ctx context.Context, task *Running if !force && task.Account() != actualUser { return nil, fmt.Errorf("task %s belongs to %s (use --force to cancel it anyway)", task.TaskID, task.User) } + if err := ae.checkTaskProcessOwner(ctx, task.PID, task.Account()); err != nil { + return nil, err + } result := &CancelResult{ TaskID: task.TaskID, @@ -153,6 +165,22 @@ func (ae *AllocationEngine) cancelRunningTask(ctx context.Context, task *Running return result, nil } +// The host bridge can run cancellation as root. Never trust a stale or forged +// reservation PID to identify a process belonging to a different account. +func (ae *AllocationEngine) checkTaskProcessOwner(ctx context.Context, pid int, owner string) error { + if pid <= 0 || !isProcessAlive(pid) { + return nil + } + identity, err := ae.owners.Resolve(ctx, pid) + if err != nil { + return fmt.Errorf("cannot verify owner of PID %d: %w", pid, err) + } + if identity.User != owner { + return fmt.Errorf("refusing to signal PID %d: current owner %s differs from task account %s", pid, identity.User, owner) + } + return nil +} + // signalProcess sends a signal to a PID and reports whether it was delivered func signalProcess(pid int, signal syscall.Signal) bool { if pid <= 0 { diff --git a/internal/gpu/docker_test.go b/internal/gpu/docker_test.go new file mode 100644 index 0000000..841e232 --- /dev/null +++ b/internal/gpu/docker_test.go @@ -0,0 +1,47 @@ +package gpu + +import ( + "context" + "os/exec" + "testing" + "time" + + "github.com/russellb/canhazgpu/internal/types" + "github.com/stretchr/testify/require" +) + +func TestDockerHostPIDStoredForRunAndQueue(t *testing.T) { + client := setupQueueTestRedis(t) + ctx := context.Background() + require.NoError(t, client.SetGPUCount(ctx, 2)) + require.NoError(t, client.SetAvailableProvider(ctx, "fake")) + engine := NewAllocationEngine(client, &types.Config{HostPID: 7654321, MemoryThreshold: 100}) + request := &types.AllocationRequest{GPUCount: 1, User: "alice", ActualUser: "alice", ReservationType: types.ReservationTypeRun} + ids, err := engine.AllocateGPUs(ctx, request) + require.NoError(t, err) + state, err := client.GetGPUState(ctx, ids[0]) + require.NoError(t, err) + require.Equal(t, 7654321, state.PID) + entry := engine.createQueueEntry(&QueuedAllocationRequest{AllocationRequest: request}) + require.Equal(t, 7654321, entry.PID) +} + +func TestCancelRejectsPIDWhoseOwnerDoesNotMatchReservation(t *testing.T) { + client := setupQueueTestRedis(t) + ctx := context.Background() + require.NoError(t, client.SetGPUCount(ctx, 1)) + engine := NewAllocationEngine(client, &types.Config{MemoryThreshold: 100}) + child := exec.Command("sleep", "60") + require.NoError(t, child.Start()) + defer func() { _ = child.Process.Kill(); _ = child.Wait() }() + reserveForTest(t, engine, 0, &types.GPUState{ + User: "not-the-process-owner", ActualUser: "not-the-process-owner", Type: types.ReservationTypeRun, + TaskID: "bad01234", PID: child.Process.Pid, StartTime: types.FlexibleTime{Time: time.Now()}, + }) + _, err := engine.CancelTask(ctx, "bad01234", "not-the-process-owner", false) + require.ErrorContains(t, err, "refusing to signal") + require.True(t, isProcessAlive(child.Process.Pid)) + state, err := client.GetGPUState(ctx, 0) + require.NoError(t, err) + require.Equal(t, "bad01234", state.TaskID, "do not release a task we failed to stop") +} diff --git a/internal/hostbridge/bridge_test.go b/internal/hostbridge/bridge_test.go new file mode 100644 index 0000000..3c66112 --- /dev/null +++ b/internal/hostbridge/bridge_test.go @@ -0,0 +1,117 @@ +package hostbridge + +import ( + "context" + "fmt" + "io" + "net" + "os" + "os/user" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/russellb/canhazgpu/internal/types" + "github.com/stretchr/testify/require" +) + +func TestContainerID(t *testing.T) { + id := strings.Repeat("a", 64) + for _, path := range []string{"/docker/" + id, "/system.slice/docker-" + id + ".scope", "/user.slice/user-1000.slice/user@1000.service/app.slice/docker-" + id + ".scope", "/docker/" + id + "/child"} { + require.Equal(t, id, ContainerID("0::"+path)) + require.Equal(t, id, ContainerID("5:cpu,memory:"+path+"\n2:pids:/")) + } + for _, value := range []string{"0::/", "0::/docker/abc", "0::/docker/" + id + "suffix"} { + require.Empty(t, ContainerID(value)) + } +} + +func TestOwnerMappingAndDockerLabels(t *testing.T) { + account, err := user.Current() + require.NoError(t, err) + if account.Uid == "0" { + t.Skip("requires a non-root account for owner mapping") + } + id := strings.Repeat("b", 64) + r := &Resolver{Owners: map[string]string{id: account.Username}} + owner, err := r.containerOwner(context.Background(), id) + require.NoError(t, err) + require.Equal(t, account.Username, owner) + _, err = validOwner("root") + require.Error(t, err) + _, err = validOwner("") + require.Error(t, err) + dir := t.TempDir() + script := fmt.Sprintf("#!/bin/sh\nprintf '%%s\\n' '{\"canhazgpu.owner\":\"%s\"}'\n", account.Username) + require.NoError(t, os.WriteFile(filepath.Join(dir, "docker"), []byte(script), 0755)) + t.Setenv("PATH", dir) + r = &Resolver{} + owner, err = r.containerOwner(context.Background(), id) + require.NoError(t, err) + require.Equal(t, account.Username, owner) + require.NoError(t, os.WriteFile(filepath.Join(dir, "docker"), []byte("#!/bin/sh\necho '{}'\n"), 0755)) + _, err = r.containerOwner(context.Background(), strings.Repeat("c", 64)) + require.ErrorContains(t, err, "needs canhazgpu.owner") +} + +func TestBridgeCredentialsUsageAndRedis(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + backend, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer backend.Close() + go func() { + conn, err := backend.Accept() + if err == nil { + defer conn.Close() + _, _ = io.Copy(conn, conn) + } + }() + addr := backend.Addr().(*net.TCPAddr) + server := &Server{ + Resolver: &Resolver{}, Config: &types.Config{RedisHost: "127.0.0.1", RedisPort: addr.Port, RedisDB: 7, MemoryThreshold: 123}, + Usage: func(context.Context) (map[int]*types.GPUUsage, error) { + return map[int]*types.GPUUsage{2: {GPUID: 2, Processes: []types.GPUProcessInfo{{PID: 98765, User: "alice"}}}}, nil + }, + Cancel: func(ctx context.Context, ref, owner string, force bool) (any, error) { + return map[string]string{"owner": owner, "ref": ref}, nil + }, + } + socket := filepath.Join(t.TempDir(), "host.sock") + require.NoError(t, server.Start(ctx, socket)) + defer server.Close() + client := Client{Socket: socket} + identity, err := client.Identity(ctx) + require.NoError(t, err) + account, err := user.Current() + require.NoError(t, err) + require.Equal(t, os.Getpid(), identity.PID) + require.Equal(t, account.Username, identity.User) + require.Equal(t, 7, identity.RedisDB) + usage, err := client.Usage(ctx) + require.NoError(t, err) + require.Equal(t, "alice", usage[2].Processes[0].User) + var result map[string]string + require.NoError(t, client.Call(ctx, Request{Operation: "cancel", TaskID: "abc"}, &result)) + require.Equal(t, identity.User, result["owner"]) + if !identity.Admin { + require.ErrorContains(t, client.Call(ctx, Request{Operation: "cancel", Force: true}, &result), "requires host root") + } + conn, err := client.RedisConn(ctx) + require.NoError(t, err) + defer conn.Close() + require.NoError(t, conn.SetDeadline(time.Now().Add(time.Second))) + _, err = conn.Write([]byte("*1\r\n$4\r\nPING\r\n")) + require.NoError(t, err) + data := make([]byte, 14) + _, err = io.ReadFull(conn, data) + require.NoError(t, err) + require.Equal(t, "*1\r\n$4\r\nPING\r\n", string(data)) + require.ErrorContains(t, client.Call(ctx, Request{Operation: "execute"}, &result), "unknown host operation") +} + +func TestSocketFailureDoesNotFallBack(t *testing.T) { + _, err := (Client{Socket: filepath.Join(t.TempDir(), "missing.sock")}).Identity(context.Background()) + require.ErrorContains(t, err, "connect to host guard") +} diff --git a/internal/hostbridge/client.go b/internal/hostbridge/client.go new file mode 100644 index 0000000..66ebc47 --- /dev/null +++ b/internal/hostbridge/client.go @@ -0,0 +1,110 @@ +// Package hostbridge connects container clients to the host guard. Device IDs +// and stored PIDs always belong to the host; the job and its supervisor remain +// in their original namespaces. +package hostbridge + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "net" + "time" + + "github.com/russellb/canhazgpu/internal/types" +) + +const DefaultSocket = "/run/canhazgpu/host.sock" + +type Identity struct { + User string `json:"user"` + PID int `json:"pid"` + ContainerID string `json:"container_id,omitempty"` + Admin bool `json:"admin"` + RedisDB int `json:"redis_db"` + MemoryThreshold int `json:"memory_threshold"` +} + +type Request struct { + Operation string `json:"operation"` + TaskID string `json:"task_id,omitempty"` + Force bool `json:"force,omitempty"` +} + +type response struct { + Error string `json:"error,omitempty"` + Result json.RawMessage `json:"result,omitempty"` +} + +type Client struct{ Socket string } + +func (c Client) connect(ctx context.Context, req Request) (net.Conn, *bufio.Reader, response, error) { + var reply response + conn, err := (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, "unix", c.Socket) + if err != nil { + return nil, nil, reply, fmt.Errorf("connect to host guard at %s: %w", c.Socket, err) + } + deadline := time.Now().Add(45 * time.Second) + if d, ok := ctx.Deadline(); ok && d.Before(deadline) { + deadline = d + } + _ = conn.SetDeadline(deadline) + stop := context.AfterFunc(ctx, func() { _ = conn.Close() }) + defer stop() + reader := bufio.NewReader(conn) + err = json.NewEncoder(conn).Encode(req) + if err == nil { + var line []byte + line, err = reader.ReadBytes('\n') + if err == nil { + err = json.Unmarshal(line, &reply) + } + } + if err == nil && reply.Error != "" { + err = fmt.Errorf("host guard: %s", reply.Error) + } + if err != nil { + _ = conn.Close() + return nil, nil, reply, err + } + _ = conn.SetDeadline(time.Time{}) + return conn, reader, reply, nil +} + +func (c Client) Call(ctx context.Context, req Request, result any) error { + conn, _, reply, err := c.connect(ctx, req) + if err != nil { + return err + } + defer conn.Close() + return json.Unmarshal(reply.Result, result) +} + +func (c Client) Identity(ctx context.Context) (Identity, error) { + var identity Identity + err := c.Call(ctx, Request{Operation: "identity"}, &identity) + return identity, err +} + +func (c Client) Usage(ctx context.Context) (map[int]*types.GPUUsage, error) { + var usage map[int]*types.GPUUsage + err := c.Call(ctx, Request{Operation: "usage"}, &usage) + return usage, err +} + +// RedisConn tunnels the existing Redis protocol so bridge-network containers +// need neither a TCP port exposed on the host nor a second Redis instance. +func (c Client) RedisConn(ctx context.Context) (net.Conn, error) { + conn, reader, _, err := c.connect(ctx, Request{Operation: "redis"}) + if err != nil { + return nil, err + } + return &bufferedConn{Conn: conn, reader: reader}, nil +} + +type bufferedConn struct { + net.Conn + reader *bufio.Reader +} + +func (c *bufferedConn) Read(p []byte) (int, error) { return c.reader.Read(p) } diff --git a/internal/hostbridge/identity.go b/internal/hostbridge/identity.go new file mode 100644 index 0000000..f29d46e --- /dev/null +++ b/internal/hostbridge/identity.go @@ -0,0 +1,152 @@ +package hostbridge + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "os/user" + "regexp" + "strconv" + "strings" + "sync" + "time" +) + +const OwnerLabel = "canhazgpu.owner" + +var containerPattern = regexp.MustCompile(`(?:^|/)(?:docker[-/])?([a-f0-9]{64})(?:\.scope)?(?:/|$)`) + +// Resolver runs only on the host. Ownership comes from host configuration or +// Docker metadata, never from USER or a caller-supplied container ID. +type Resolver struct { + Owners map[string]string + mu sync.Mutex + cache map[string]ownerCache +} +type ownerCache struct { + user string + expires time.Time +} + +func ContainerID(cgroups string) string { + for _, line := range strings.Split(cgroups, "\n") { + parts := strings.SplitN(line, ":", 3) + if len(parts) != 3 { + continue + } + if m := containerPattern.FindStringSubmatch(parts[2]); len(m) > 1 { + return m[1] + } + } + return "" +} + +func (r *Resolver) Resolve(ctx context.Context, pid int) (Identity, error) { + identity := Identity{PID: pid} + if pid <= 0 { + return identity, fmt.Errorf("invalid PID") + } + cgroups, err := os.ReadFile(fmt.Sprintf("/proc/%d/cgroup", pid)) + if err != nil { + return identity, err + } + identity.ContainerID = ContainerID(string(cgroups)) + if identity.ContainerID != "" { + identity.User, err = r.containerOwner(ctx, identity.ContainerID) + return identity, err + } + status, err := os.ReadFile(fmt.Sprintf("/proc/%d/status", pid)) + if err != nil { + return identity, err + } + for _, line := range strings.Split(string(status), "\n") { + fields := strings.Fields(line) + if len(fields) >= 2 && fields[0] == "Uid:" { + account, err := user.LookupId(fields[1]) + if err != nil { + return identity, err + } + identity.User = account.Username + identity.Admin = fields[1] == "0" + return identity, nil + } + } + return identity, fmt.Errorf("UID missing for PID %d", pid) +} + +func (r *Resolver) containerOwner(ctx context.Context, id string) (string, error) { + if owner := r.Owners[id]; owner != "" { + return validOwner(owner) + } + r.mu.Lock() + cached, ok := r.cache[id] + r.mu.Unlock() + if ok && time.Now().Before(cached.expires) { + return cached.user, nil + } + ctx, cancel := context.WithTimeout(ctx, 3*time.Second) + defer cancel() + out, err := exec.CommandContext(ctx, "docker", "inspect", "--type", "container", "--format", `{{json .Config.Labels}}`, id).Output() + if err != nil { + return "", fmt.Errorf("cannot inspect Docker container %s: %w", id[:12], err) + } + var labels map[string]string + if err := json.Unmarshal(out, &labels); err != nil { + return "", err + } + owner, err := validOwner(labels[OwnerLabel]) + if err != nil { + return "", fmt.Errorf("container %s needs %s= or a host owner mapping: %w", id[:12], OwnerLabel, err) + } + r.mu.Lock() + if r.cache == nil { + r.cache = make(map[string]ownerCache) + } + // Bound memory even on hosts that create many short-lived containers. + if len(r.cache) >= 1024 { + clear(r.cache) + } + r.cache[id] = ownerCache{owner, time.Now().Add(30 * time.Second)} + r.mu.Unlock() + return owner, nil +} + +func validOwner(name string) (string, error) { + if name == "" { + return "", fmt.Errorf("empty owner") + } + account, err := user.Lookup(name) + if err != nil { + return "", fmt.Errorf("unknown host account %q", name) + } + uid, err := strconv.Atoi(account.Uid) + if err != nil || uid == 0 { + return "", fmt.Errorf("container owner must be a non-root host account") + } + return account.Username, nil +} + +func LoadOwners(path string) (map[string]string, error) { + owners := make(map[string]string) + if path == "" { + return owners, nil + } + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + if err := json.Unmarshal(data, &owners); err != nil { + return nil, err + } + for id, name := range owners { + if len(id) != 64 || ContainerID("0::/docker/"+id) != id { + return nil, fmt.Errorf("owner mapping requires full Docker container IDs: %q", id) + } + if _, err := validOwner(name); err != nil { + return nil, err + } + } + return owners, nil +} diff --git a/internal/hostbridge/peer_linux.go b/internal/hostbridge/peer_linux.go new file mode 100644 index 0000000..4f91e82 --- /dev/null +++ b/internal/hostbridge/peer_linux.go @@ -0,0 +1,24 @@ +package hostbridge + +import ( + "net" + + "golang.org/x/sys/unix" +) + +func peerPID(conn *net.UnixConn) (int, error) { + raw, err := conn.SyscallConn() + if err != nil { + return 0, err + } + var cred *unix.Ucred + var credErr error + err = raw.Control(func(fd uintptr) { cred, credErr = unix.GetsockoptUcred(int(fd), unix.SOL_SOCKET, unix.SO_PEERCRED) }) + if err != nil { + return 0, err + } + if credErr != nil { + return 0, credErr + } + return int(cred.Pid), nil +} diff --git a/internal/hostbridge/peer_other.go b/internal/hostbridge/peer_other.go new file mode 100644 index 0000000..279a592 --- /dev/null +++ b/internal/hostbridge/peer_other.go @@ -0,0 +1,12 @@ +//go:build !linux + +package hostbridge + +import ( + "fmt" + "net" +) + +func peerPID(conn *net.UnixConn) (int, error) { + return 0, fmt.Errorf("Docker host bridge requires Linux") +} diff --git a/internal/hostbridge/server.go b/internal/hostbridge/server.go new file mode 100644 index 0000000..9eb5d31 --- /dev/null +++ b/internal/hostbridge/server.go @@ -0,0 +1,143 @@ +package hostbridge + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "net" + "os" + "path/filepath" + "sync" + "time" + + "github.com/russellb/canhazgpu/internal/types" +) + +type Server struct { + Resolver *Resolver + Config *types.Config + Usage func(context.Context) (map[int]*types.GPUUsage, error) + Cancel func(context.Context, string, string, bool) (any, error) + listener *net.UnixListener + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup +} + +// Start refuses to replace an existing socket: another guard may own it. +// The directory is mounted, rather than the inode, so clients reconnect after +// a normal guard restart. Access to this socket implies access to shared Redis. +func (s *Server) Start(ctx context.Context, path string) error { + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + return err + } + listener, err := net.ListenUnix("unix", &net.UnixAddr{Name: path, Net: "unix"}) + if err != nil { + return err + } + if err := os.Chmod(path, 0666); err != nil { + listener.Close() + return err + } + s.listener = listener + s.ctx, s.cancel = context.WithCancel(ctx) + s.wg.Add(1) + go func() { + defer s.wg.Done() + for { + conn, err := listener.AcceptUnix() + if err != nil { + return + } + s.wg.Add(1) + go func() { defer s.wg.Done(); s.serve(conn) }() + } + }() + return nil +} + +func (s *Server) Close() { + s.cancel() + _ = s.listener.Close() + s.wg.Wait() +} + +func (s *Server) serve(conn *net.UnixConn) { + defer conn.Close() + stop := context.AfterFunc(s.ctx, func() { _ = conn.Close() }) + defer stop() + _ = conn.SetDeadline(time.Now().Add(45 * time.Second)) + reader := bufio.NewReader(io.LimitReader(conn, 4096)) + line, err := reader.ReadBytes('\n') + if err != nil { + return + } + var req Request + if err := json.Unmarshal(line, &req); err != nil { + s.reply(conn, nil, err) + return + } + pid, err := peerPID(conn) + if err != nil { + s.reply(conn, nil, err) + return + } + identity, err := s.Resolver.Resolve(s.ctx, pid) + if err != nil { + s.reply(conn, nil, err) + return + } + identity.RedisDB = s.Config.RedisDB + identity.MemoryThreshold = s.Config.MemoryThreshold + requestCtx, cancel := context.WithTimeout(s.ctx, 40*time.Second) + defer cancel() + switch req.Operation { + case "identity": + s.reply(conn, identity, nil) + case "usage": + usage, err := s.Usage(requestCtx) + s.reply(conn, usage, err) + case "cancel": + if req.Force && !identity.Admin { + s.reply(conn, nil, fmt.Errorf("--force requires host root")) + return + } + result, err := s.Cancel(requestCtx, req.TaskID, identity.User, req.Force) + s.reply(conn, result, err) + case "redis": + backend, err := (&net.Dialer{Timeout: 5 * time.Second}).DialContext(s.ctx, "tcp", net.JoinHostPort(s.Config.RedisHost, fmt.Sprint(s.Config.RedisPort))) + if err != nil { + s.reply(conn, nil, err) + return + } + defer backend.Close() + stopBackend := context.AfterFunc(s.ctx, func() { _ = backend.Close() }) + defer stopBackend() + if !s.reply(conn, true, nil) { + return + } + _ = conn.SetDeadline(time.Time{}) + // The client waits for the acknowledgement before sending Redis bytes, + // so the request reader cannot have buffered protocol data. + done := make(chan struct{}) + go func() { _, _ = io.Copy(backend, conn); _ = backend.Close(); close(done) }() + _, _ = io.Copy(conn, backend) + _ = conn.Close() + <-done + default: + s.reply(conn, nil, fmt.Errorf("unknown host operation %q", req.Operation)) + } +} + +func (s *Server) reply(conn net.Conn, result any, err error) bool { + r := response{} + if err == nil { + r.Result, err = json.Marshal(result) + } + if err != nil { + r.Error = err.Error() + } + return json.NewEncoder(conn).Encode(r) == nil +} diff --git a/internal/redis_client/client.go b/internal/redis_client/client.go index c99ee80..ef75c07 100644 --- a/internal/redis_client/client.go +++ b/internal/redis_client/client.go @@ -5,9 +5,11 @@ import ( "encoding/json" "fmt" "math/rand" + "net" "time" "github.com/go-redis/redis/v8" + "github.com/russellb/canhazgpu/internal/hostbridge" "github.com/russellb/canhazgpu/internal/types" ) @@ -18,8 +20,9 @@ type Client struct { func NewClient(config *types.Config) *Client { rdb := redis.NewClient(&redis.Options{ - Addr: fmt.Sprintf("%s:%d", config.RedisHost, config.RedisPort), - DB: config.RedisDB, + Dialer: bridgeDialer(config), + Addr: fmt.Sprintf("%s:%d", config.RedisHost, config.RedisPort), + DB: config.RedisDB, // Connection health settings to detect and recover from stale connections. // This is critical for long-lived processes like the supervisor, where a @@ -69,8 +72,9 @@ func (c *Client) Reconnect() error { _ = c.rdb.Close() c.rdb = redis.NewClient(&redis.Options{ - Addr: fmt.Sprintf("%s:%d", c.config.RedisHost, c.config.RedisPort), - DB: c.config.RedisDB, + Dialer: bridgeDialer(c.config), + Addr: fmt.Sprintf("%s:%d", c.config.RedisHost, c.config.RedisPort), + DB: c.config.RedisDB, DialTimeout: 5 * time.Second, ReadTimeout: 5 * time.Second, @@ -91,6 +95,15 @@ func (c *Client) Reconnect() error { return c.rdb.Ping(ctx).Err() } +func bridgeDialer(config *types.Config) func(context.Context, string, string) (net.Conn, error) { + if config.HostSocket == "" { + return nil + } + return func(ctx context.Context, network, addr string) (net.Conn, error) { + return (hostbridge.Client{Socket: config.HostSocket}).RedisConn(ctx) + } +} + // GPU State Management func (c *Client) SetGPUCount(ctx context.Context, count int) error { diff --git a/internal/types/types.go b/internal/types/types.go index d4778a2..be95946 100644 --- a/internal/types/types.go +++ b/internal/types/types.go @@ -194,6 +194,10 @@ type UsageRecord struct { // Config represents the application configuration type Config struct { + HostSocket string // Optional host guard bridge (shared Redis and host device view) + HostPID int // PID in the host namespace, obtained from Unix peer credentials + HostUser string // Host account owning this process/container + ContainerOwners map[string]string // Host-only overrides keyed by full container ID RedisHost string RedisPort int RedisDB int diff --git a/mkdocs.yml b/mkdocs.yml index 3393e8b..e701402 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -69,6 +69,7 @@ nav: - Scheduled Bookings: usage-schedule.md - Releasing GPUs: usage-release.md - Status Monitoring: usage-status.md + - Docker 开发环境: docker.md - Features: - GPU Validation: features-validation.md - Unreserved Usage Detection: features-unreserved.md @@ -87,4 +88,4 @@ nav: extra: social: - icon: fontawesome/brands/github - link: https://github.com/rbryant/canhazgpu \ No newline at end of file + link: https://github.com/rbryant/canhazgpu diff --git a/scripts/test-docker.py b/scripts/test-docker.py new file mode 100644 index 0000000..9ec7005 --- /dev/null +++ b/scripts/test-docker.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Real Docker/PID integration using simulated devices; requires local image, +Docker access and passwordless sudo for the temporary host guard. No real +accelerator is touched. Every Redis key and container belongs to this test. +""" +import argparse +import json +import os +from pathlib import Path +import pwd +import socket +import subprocess as sp +import tempfile +import time + + +def run(*args, check=True, **kwargs): + return sp.run([str(x) for x in args], text=True, stdout=sp.PIPE, + stderr=sp.PIPE, check=check, **kwargs) + + +def eventually(fn, timeout=25): + end = time.monotonic() + timeout + last = None + while time.monotonic() < end: + try: + result = fn() + if result: + return result + except (AssertionError, ValueError, KeyError, sp.CalledProcessError) as exc: + last = exc + time.sleep(.2) + raise AssertionError(f"condition timed out: {last}") + + +def stop_guard(guard): + children = run('pgrep', '-P', guard.pid, check=False).stdout.split() + for child in children: + run('sudo', '-n', 'kill', '-TERM', child, check=False) + try: + guard.wait(timeout=20) + except sp.TimeoutExpired: + for child in children: + run('sudo', '-n', 'kill', '-KILL', child, check=False) + guard.wait(timeout=10) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--image', required=True, help='Existing Linux image with sh and sleep') + parser.add_argument('--binary', default='build/canhazgpu-docker') + args = parser.parse_args() + binary = Path(args.binary).resolve() + owner = pwd.getpwuid(os.getuid()).pw_name + assert os.getuid() != 0, 'Run as your host account; sudo is used only for the test guard' + other = 'nobody' + assert owner != other + run('sudo', '-n', 'true') + run('docker', 'image', 'inspect', args.image) + with tempfile.TemporaryDirectory(prefix='chg-docker-') as temp: + base = Path(temp) + base.chmod(0o755) + usage_file = base / 'usage.json' + usage_file.write_text('[]') + mock = base / 'nvidia-smi' + mock.write_text('''#!/usr/bin/python3 +import json,sys,os +from pathlib import Path +rows=json.loads(Path(__file__).with_name('usage.json').read_text()) +alive=[] +for gpu,pid in rows: + try: + stat=Path('/proc/%d/stat'%pid).read_text().rsplit(')',1)[1].split() + if stat[0]!='Z': alive.append((gpu,pid)) + except (FileNotFoundError,ProcessLookupError): pass +arg=' '.join(sys.argv[1:]) +if '--query-gpu=' in arg: + for gpu in range(2): + mem=512*sum(g==gpu for g,p in alive) + print('%d, GPU-test%d, Test GPU, %d, 0'%(gpu,gpu,mem)) +elif '--query-compute-apps=' in arg: + for gpu,pid in alive: print('%d, sleep, GPU-test%d, 512 MiB'%(pid,gpu)) +elif '-L' in arg: + print('GPU 0: Test GPU\\nGPU 1: Test GPU') +''') + mock.chmod(0o755) + with socket.socket() as port_socket: + port_socket.bind(('127.0.0.1', 0)) + port = port_socket.getsockname()[1] + redis_log = open(base / 'redis.log', 'w') + guard_log = open(base / 'guard.log', 'w') + redis = sp.Popen(['redis-server', '--bind', '127.0.0.1', '--port', str(port), + '--save', '', '--appendonly', 'no'], stdout=redis_log, stderr=sp.STDOUT) + owners_file = base / 'owners.json' + owners_file.write_text('{}') + containers = [] + guard = None + try: + eventually(lambda: run('redis-cli', '-p', port, 'ping', check=False).stdout.strip() == 'PONG') + run(binary, '--host-socket', 'off', '--redis-port', port, 'admin', '--gpus', '2', '--provider', 'fake') + run('redis-cli', '-p', port, 'set', 'canhazgpu:provider', 'nvidia') + guard_argv = ['sudo', '-n', 'env', 'PATH='+temp+':'+os.environ['PATH'], str(binary), + 'guard', '--redis-port', str(port), '--listen', str(base / 'host.sock'), + '--interval', '1s', '--grace', '2s', '--confirmations', '1', '--enforce', + '--max-warnings', '1', '--warn-interval', '1s', '--kill-grace', '1s', + '--channels', 'log', '--docker-owners', str(owners_file)] + guard = sp.Popen(guard_argv, stdout=guard_log, stderr=sp.STDOUT) + eventually(lambda: (base / 'host.sock').exists()) + for account in (owner, other, ''): + cid = run('docker', 'run', '--rm', '-d', '--network', 'none', + '--label', 'canhazgpu.owner='+account, + '-v', temp+':/run/canhazgpu:ro', + '-v', str(binary)+':/usr/local/bin/canhazgpu:ro', + '--entrypoint', '/bin/sh', args.image, '-c', 'exec sleep 600').stdout.strip() + containers.append(cid) + a, b, unmapped = containers + + def cli(cid, *argv, **kwargs): + return run('docker', 'exec', cid, 'canhazgpu', *argv, **kwargs) + + def state(gpu): + return json.loads(run('redis-cli', '-p', port, 'get', f'canhazgpu:gpu:{gpu}').stdout.strip() or '{}') or {} + + def host(*argv, **kwargs): + return run(binary, '--host-socket', base / 'host.sock', *argv, **kwargs) + + def start(cid, gpu): + run('docker', 'exec', '-d', cid, 'canhazgpu', 'run', '--gpu-ids', gpu, + '--idle-timeout', '0', '--', 'sleep', '600') + return eventually(lambda: state(gpu).get('pid') and state(gpu)) + + def usage(rows): + tmp = base / 'usage.next' + tmp.write_text(json.dumps(rows)) + tmp.replace(usage_file) + + sa = start(a, 0) + sb = start(b, 1) + assert sa['actual_user'] == owner and sb['actual_user'] == other + assert sa['pid'] != sb['pid'] + assert a in Path('/proc/%d/cgroup' % sa['pid']).read_text() + assert b in Path('/proc/%d/cgroup' % sb['pid']).read_text() + usage([[0, sa['pid']], [1, sb['pid']]]) + def attributed(): + states = json.loads(cli(a, 'status', '--json', '--no-schedule').stdout) + return all(s.get('processes') and s['processes'][0]['user'] == u + for s, u in zip(states, (owner, other))) + eventually(attributed) + for output in (host('status', '--json', '--no-schedule'), cli(b, 'status', '--json', '--no-schedule')): + states = json.loads(output.stdout) + assert [s['user'] for s in states] == [owner, other] + assert not any(s.get('foreign_users') for s in states) + print('PASS: two root containers retain distinct host owners and host PIDs; status agrees', flush=True) + missing_owner = cli(unmapped, 'status', check=False) + assert missing_owner.returncode != 0 and 'needs canhazgpu.owner' in missing_owner.stderr + owners_file.write_text(json.dumps({unmapped: owner})) + stop_guard(guard) + stopped = cli(a, 'status', check=False) + assert stopped.returncode != 0 and 'connect to host guard' in stopped.stderr + guard = sp.Popen(guard_argv, stdout=guard_log, stderr=sp.STDOUT) + eventually(lambda: (base / 'host.sock').exists()) + states = json.loads(cli(unmapped, 'status', '--json', '--no-schedule').stdout) + assert [s['user'] for s in states] == [owner, other] + assert Path('/proc/%d' % sa['pid']).exists() + print('PASS: missing owners rejected; existing container mapping and guard restart work', flush=True) + assert cli(b, 'cancel', sa['task_id'], check=False).returncode != 0 + assert cli(b, 'cancel', sa['task_id'], '--force', check=False).returncode != 0 + assert Path('/proc/%d' % sa['pid']).exists() + host('cancel', sa['task_id']) + cli(b, 'cancel', sb['task_id']) + eventually(lambda: not state(0).get('user') and not state(1).get('user')) + print('PASS: host and container cancellation kill container processes; foreign cancellation rejected', flush=True) + usage([]) + eventually(lambda: all(s['status'] == 'AVAILABLE' for s in json.loads(host('status', '--json', '--no-schedule').stdout))) + # Finishing a command preserves its exit status, env and supervisor cleanup. + done = cli(a, 'run', '--gpu-ids', '0', '--idle-timeout', '0', '--', + 'sh', '-c', 'test "$CUDA_VISIBLE_DEVICES" = 0; exit 7', check=False) + assert done.returncode == 7, done + eventually(lambda: not state(0).get('user')) + # Queue PID is also in host namespace and a different owner can cancel its own queue entry. + sa = start(a, 0) + run('docker', 'exec', '-d', b, 'canhazgpu', 'run', '--gpu-ids', '0', '--', 'sleep', '600') + def queued(): + keys = run('redis-cli', '-p', port, '--scan', '--pattern', 'canhazgpu:queue:entry:*').stdout.splitlines() + if not keys: return None + return json.loads(run('redis-cli', '-p', port, 'get', keys[0]).stdout) + entry = eventually(queued) + assert entry['actual_user'] == other + assert b in Path('/proc/%d/cgroup' % entry['pid']).read_text() + cli(b, 'cancel', entry['id'][:8]) + eventually(lambda: not queued()) + host('cancel', sa['task_id']) + print('PASS: exit code, visible devices, automatic release, and queued cancellation', flush=True) + # A bypassing process is actually killed by the host guard. + run('docker', 'exec', '-d', a, 'sleep', '601') + def rogue_pid(): + for line in run('docker', 'top', a, '-eo', 'pid,args').stdout.splitlines(): + fields = line.split(None, 1) + if len(fields) == 2 and fields[1] == 'sleep 601': return int(fields[0]) + rogue = eventually(rogue_pid) + usage([[0, rogue]]) + eventually(lambda: not Path('/proc/%d' % rogue).exists()) + assert owner in (base / 'guard.log').read_text() + assert 'SIGINT' in (base / 'guard.log').read_text() + print('PASS: host guard attributes and terminates an unreserved process inside Docker', flush=True) + # A missing bridge is an error, never a fallback to container-local Redis or PIDs. + missing = run('docker', 'exec', '-e', 'CANHAZGPU_HOST_SOCKET=/missing.sock', a, + 'canhazgpu', 'status', check=False) + assert missing.returncode != 0 and 'connect to host guard' in missing.stderr + print('PASS: disconnected clients fail explicitly', flush=True) + except Exception as exc: + if isinstance(exc, sp.CalledProcessError): + print(exc.stderr, flush=True) + guard_log.flush() + print((base / 'guard.log').read_text(), flush=True) + raise + finally: + for cid in containers: + run('docker', 'rm', '-f', cid, check=False) + if guard: + stop_guard(guard) + redis.terminate() + redis.wait(timeout=10) + guard_log.close() + redis_log.close() + +if __name__ == '__main__': + main() From 1d87b5956fe8a626f7604d8793662f5820f2cad5 Mon Sep 17 00:00:00 2001 From: specture724 Date: Mon, 7 Sep 2026 10:39:10 +0800 Subject: [PATCH 2/2] fix: load Docker owner mappings without a running guard (cherry picked from commit 29eea794fdb503b6a850fc8512367744cca113c4) --- docs/docker.md | 5 ++++ internal/cli/docker_owners_test.go | 38 ++++++++++++++++++++++++++++++ internal/cli/guard.go | 5 ---- internal/cli/root.go | 23 ++++++++++++++++++ 4 files changed, 66 insertions(+), 5 deletions(-) create mode 100644 internal/cli/docker_owners_test.go diff --git a/docs/docker.md b/docs/docker.md index 738fbb6..4c69a5c 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -95,6 +95,11 @@ Docker 不能向运行中的容器补充普通 bind mount，缺少挂载时需 无法确认归属的容器不会被默认为宿主机 root，客户端会提示补齐配置； 设备扫描会显示 `unknown`，guard 不会终止身份不明的进程。 +`/etc/canhazgpu/docker-owners.json` 会被宿主机所有命令自动读取。 +因此即使 guard 停止，直接执行 `canhazgpu status` 也能显示映射中的账户； +不需要为了查看归属而启动强制执行。其他路径可使用全局 `--docker-owners PATH` +或 `CANHAZGPU_DOCKER_OWNERS=PATH`。独立 status 每次调用读取文件，guard 修改映射后需重启。 + ## 4. 容器内外使用 ```bash diff --git a/internal/cli/docker_owners_test.go b/internal/cli/docker_owners_test.go new file mode 100644 index 0000000..e28a053 --- /dev/null +++ b/internal/cli/docker_owners_test.go @@ -0,0 +1,38 @@ +package cli + +import ( + "encoding/json" + "os" + "os/user" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestLoadContainerOwnersExplicitFile(t *testing.T) { + account, err := user.Current() + require.NoError(t, err) + if account.Uid == "0" { + t.Skip("owner mapping requires a non-root account") + } + id := strings.Repeat("a", 64) + path := filepath.Join(t.TempDir(), "owners.json") + data, err := json.Marshal(map[string]string{id: account.Username}) + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, data, 0600)) + owners, err := loadContainerOwners(path) + require.NoError(t, err) + require.Equal(t, account.Username, owners[id]) + _, err = loadContainerOwners(path + ".missing") + require.Error(t, err, "explicit missing files must not silently discard mappings") + require.NoError(t, os.WriteFile(path, []byte("not JSON"), 0600)) + _, err = loadContainerOwners(path) + require.Error(t, err) +} + +func TestDockerOwnerFlagIsGlobal(t *testing.T) { + require.NotNil(t, rootCmd.PersistentFlags().Lookup("docker-owners")) + require.Nil(t, guardCmd.LocalNonPersistentFlags().Lookup("docker-owners")) +} diff --git a/internal/cli/guard.go b/internal/cli/guard.go index 920eb29..1163cb3 100644 --- a/internal/cli/guard.go +++ b/internal/cli/guard.go @@ -72,7 +72,6 @@ Example usage: func init() { defaults := gpu.DefaultGuardConfig() guardCmd.Flags().String("listen", "", "Serve container clients on this Unix socket (e.g. /run/canhazgpu/host.sock)") - guardCmd.Flags().String("docker-owners", "", "Host JSON file mapping full Docker container IDs to host accounts") guardCmd.Flags().String("interval", utils.FormatDurationShort(defaults.Interval), "How often to scan the GPUs") guardCmd.Flags().Bool("once", false, "Run a single scan and exit (for cron)") @@ -109,10 +108,6 @@ func runGuard(ctx context.Context) error { } config := getConfig() - config.ContainerOwners, err = hostbridge.LoadOwners(viper.GetString("guard.docker-owners")) - if err != nil { - return fmt.Errorf("Docker owner mappings: %w", err) - } client := redis_client.NewClient(config) defer func() { if err := client.Close(); err != nil { diff --git a/internal/cli/root.go b/internal/cli/root.go index 776584f..7897f90 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -36,6 +36,7 @@ func init() { // Global flags rootCmd.PersistentFlags().String("host-socket", "", "Host guard Unix socket (auto-detected at /run/canhazgpu/host.sock; 'off' disables)") + rootCmd.PersistentFlags().String("docker-owners", "", "Host JSON file mapping Docker container IDs to accounts (default: /etc/canhazgpu/docker-owners.json)") rootCmd.PersistentFlags().StringVar(&configFile, "config", "", "config file (default is $HOME/.canhazgpu.yaml)") rootCmd.PersistentFlags().String("redis-host", "localhost", "Redis host") rootCmd.PersistentFlags().Int("redis-port", 6379, "Redis port") @@ -175,6 +176,17 @@ func getCurrentUser() string { func prepareHostBridge(cmd *cobra.Command, args []string) error { cfg := getConfig() + ownersPath := viper.GetString("docker-owners") + if ownersPath == "" && cmd.Name() == "guard" { + // Preserve existing guard configuration files after promoting the flag + // to a global option that standalone status can also use. + ownersPath = viper.GetString("guard.docker-owners") + } + owners, err := loadContainerOwners(ownersPath) + if err != nil { + return fmt.Errorf("Docker owner mappings: %w", err) + } + cfg.ContainerOwners = owners if cmd.Name() == "guard" { if cfg.HostSocket != "" && cfg.HostSocket != "off" { return fmt.Errorf("guard must run on the host; remove --host-socket / CANHAZGPU_HOST_SOCKET") @@ -207,3 +219,14 @@ func prepareHostBridge(cmd *cobra.Command, args []string) error { } return nil } + +func loadContainerOwners(path string) (map[string]string, error) { + if path != "" { + return hostbridge.LoadOwners(path) + } + owners, err := hostbridge.LoadOwners("/etc/canhazgpu/docker-owners.json") + if os.IsNotExist(err) { + return map[string]string{}, nil + } + return owners, err +}