From 922584a5fed503bcf088fa5f4472715ac7112388 Mon Sep 17 00:00:00 2001 From: specture724 Date: Tue, 1 Sep 2026 22:55:44 +0800 Subject: [PATCH 1/4] feat: add Huawei Ascend NPU support --- README.md | 19 +- docs/admin-troubleshooting.md | 41 +++- docs/commands.md | 23 ++- docs/configuration.md | 4 +- docs/dev-architecture.md | 5 +- docs/dev-contributing.md | 10 +- docs/dev-testing.md | 16 +- docs/features-guard.md | 2 +- docs/features-idle-timeout.md | 2 +- docs/features-validation.md | 47 +++-- docs/index.md | 5 +- docs/installation.md | 27 ++- docs/quick-quick-start.md | 5 +- docs/quickstart-zh.md | 10 +- docs/quickstart.md | 10 +- docs/usage-reserve.md | 26 ++- docs/usage-run.md | 27 ++- docs/usage-status.md | 6 +- internal/cli/admin.go | 26 ++- internal/cli/reserve.go | 19 +- internal/cli/run.go | 33 ++-- internal/cli/run_env_test.go | 35 ++++ internal/cli/run_test.go | 13 +- internal/gpu/allocation.go | 4 +- internal/gpu/allocation_test.go | 4 +- internal/gpu/ascend_provider.go | 277 +++++++++++++++++++++++++++ internal/gpu/ascend_provider_test.go | 70 +++++++ internal/gpu/guard.go | 2 +- internal/gpu/provider.go | 36 +++- internal/gpu/test_helpers_test.go | 13 +- internal/gpu/validation_test.go | 4 +- internal/types/types.go | 6 +- 32 files changed, 702 insertions(+), 125 deletions(-) create mode 100644 internal/cli/run_env_test.go create mode 100644 internal/gpu/ascend_provider.go create mode 100644 internal/gpu/ascend_provider_test.go diff --git a/README.md b/README.md index fd8e243..469ecf0 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ In shared development environments with multiple GPUs, researchers and developer You peacefully share a host but want a helper to avoid accidental conflicts. -- You have a single host with GPUs (NVIDIA or AMD) shared by multiple users +- You have a single host with NVIDIA GPUs, AMD GPUs, or Huawei Ascend NPUs 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 @@ -35,7 +35,8 @@ canhazgpu admin --gpus 8 canhazgpu status # Run vLLM with an automatic 2 GPU reservation. -# - CUDA_VISIBLE_DEVICES is set in the environment before running the command. +# - The provider visibility variable is set before running the command +# (CUDA_VISIBLE_DEVICES for NVIDIA/AMD; ASCEND_RT_VISIBLE_DEVICES for Ascend). # - If GPUs are unavailable, waits in queue until they become available. canhazgpu run --gpus 2 -- vllm serve my/model --tensor-parallel-size 2 @@ -61,9 +62,12 @@ canhazgpu reserve --gpus 1 --duration 4h # Reserve specific GPU IDs manually canhazgpu reserve --gpu-ids 0,2 --duration 2h -# Reserve GPUs and set CUDA_VISIBLE_DEVICES in one step (for scripting) +# Reserve NVIDIA/AMD GPUs and set CUDA_VISIBLE_DEVICES in one step (for scripting) export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --short) +# For Ascend, use the CANN visibility variable instead. +export ASCEND_RT_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --short) + # Release manual reservations when done canhazgpu release @@ -83,8 +87,8 @@ canhazgpu web --port 8080 - **MRU-per-user allocation**: Smart GPU affinity using most recently used per-user strategy with LRU fallback - **Specific GPU reservation**: Reserve exact GPU IDs when needed (e.g., --gpu-ids 1,3) - **Unreserved usage detection**: Identifies GPUs in use without proper reservations -- **Real-time validation**: Uses nvidia-smi or amd-smi to verify actual GPU usage -- **Multi-provider support**: Supports both NVIDIA and AMD GPUs with automatic detection +- **Real-time validation**: Uses nvidia-smi, amd-smi, or npu-smi to verify actual device usage +- **Multi-provider support**: Supports NVIDIA, AMD, and Huawei Ascend devices with automatic detection - **Flexible reservations**: Support for both command execution and manual reservations - **Reservation reporting**: Track and analyze GPU reservation patterns over time by user - **Web dashboard**: Real-time monitoring interface with status and reservation reports @@ -121,6 +125,7 @@ For detailed usage, configuration, and administration: - **GPUs** with appropriate management tools: - **NVIDIA GPUs**: nvidia-smi available - **AMD GPUs**: amd-smi available (ROCm 5.7+) + - **Huawei Ascend NPUs** (including 910B1): npu-smi available and the CANN runtime configured for the user - **System access** to `/proc` filesystem or `ps` command ## Installation @@ -146,11 +151,13 @@ Then initialize the GPU pool: canhazgpu admin --gpus $(nvidia-smi -L | wc -l) # For NVIDIA # OR canhazgpu admin --gpus $(amd-smi list --json | jq 'length') # For AMD +# OR (example: eight Ascend NPUs) +canhazgpu admin --gpus 8 --provider ascend # For Huawei Ascend ``` ## How It Works -1. **Validation**: Uses nvidia-smi or amd-smi to detect actual GPU usage and identify conflicts +1. **Validation**: Uses nvidia-smi, amd-smi, or npu-smi to detect actual device usage and identify conflicts 2. **Coordination**: Uses Redis for distributed state management and race condition prevention 3. **Queueing**: FCFS queue; the first entry whose full request can be satisfied is allocated, so no GPUs are held by a job that cannot start yet 4. **Allocation**: MRU-per-user (Most Recently Used per user) strategy provides GPU affinity with LRU fallback for fair distribution diff --git a/docs/admin-troubleshooting.md b/docs/admin-troubleshooting.md index 13b437a..090b536 100644 --- a/docs/admin-troubleshooting.md +++ b/docs/admin-troubleshooting.md @@ -93,6 +93,8 @@ redis-cli get "canhazgpu:provider" canhazgpu admin --gpus 8 --provider nvidia --force # OR canhazgpu admin --gpus 8 --provider amd --force +# OR +canhazgpu admin --gpus 8 --provider ascend --force # Let system auto-detect canhazgpu admin --gpus 8 --force @@ -110,6 +112,43 @@ canhazgpu admin --gpus 4 --provider nvidia # Use AMD provider for AMD GPUs canhazgpu admin --gpus 2 --provider amd + +# Use Huawei Ascend provider for Ascend NPUs +canhazgpu admin --gpus 8 --provider ascend +``` + +## Huawei Ascend NPU Issues + +### npu-smi Permission Denied + +**Symptoms:** +```bash +❯ npu-smi info +DrvMngGetConsoleLogLevel failed. (ret=4) +dcmi module initialize failed. ret is -8005 +``` + +**Cause:** The account cannot read the Ascend device nodes. On typical CANN +installations the nodes are owned by the configured runtime group, commonly +`HwHiAiUser`. + +**Solution:** An administrator must add the account to that group, then the +user must start a completely new login session: + +```bash +sudo usermod -aG HwHiAiUser + +# After logging out and back in +id -nG +npu-smi info +``` + +Read `/etc/ascend_install.info` to confirm the site's `UserGroup`; do not +assume `HwHiAiUser` if the installation uses a different group. Once +`npu-smi info` succeeds for the user, initialize the pool with: + +```bash +canhazgpu admin --gpus 8 --provider ascend ``` ## NVIDIA GPU Issues @@ -623,4 +662,4 @@ ps aux | grep -E "(redis|nvidia|python)" - `dmesg` output for hardware issues - Any custom monitoring logs -This troubleshooting guide covers the most common issues encountered in production deployments of canhazgpu. Most problems can be resolved by following these systematic approaches. \ No newline at end of file +This troubleshooting guide covers the most common issues encountered in production deployments of canhazgpu. Most problems can be resolved by following these systematic approaches. diff --git a/docs/commands.md b/docs/commands.md index 867759c..a6001b1 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -13,7 +13,7 @@ Commands: release Release manually reserved GPUs held by the current user report Generate GPU usage reports reserve Reserve GPUs manually for a specified duration - run Reserve GPUs and run a command with CUDA_VISIBLE_DEVICES set + run Reserve GPUs and run a command with device visibility set schedule Show the GPU booking schedule for a day status Show current GPU allocation status violations Show GPU usage that bypassed the reservation system @@ -103,7 +103,7 @@ canhazgpu admin --gpus [--force] [--provider ] **Options:** - `--gpus`: Number of GPUs available on this machine (required) - `--force`: Force reinitialization even if already initialized -- `--provider`: GPU provider type (`nvidia`, `amd`, or `fake`). Auto-detected if not specified. +- `--provider`: Device provider type (`nvidia`, `amd`, `ascend`, or `fake`). Auto-detected if not specified. **Examples:** ```bash @@ -116,6 +116,9 @@ canhazgpu admin --gpus 8 --provider nvidia # Use AMD GPUs canhazgpu admin --gpus 4 --provider amd +# Use Huawei Ascend NPUs +canhazgpu admin --gpus 8 --provider ascend + # Use fake provider for development/testing (no real GPUs required) canhazgpu admin --gpus 4 --provider fake @@ -125,7 +128,7 @@ canhazgpu admin --gpus 4 --force !!! tip "Fake Provider for Development" Use `--provider fake` to develop and test canhazgpu on systems without actual GPUs. - The fake provider simulates GPU behavior without requiring nvidia-smi or amd-smi. + The fake provider simulates GPU behavior without requiring nvidia-smi, amd-smi, or npu-smi. !!! warning "Destructive Operation" Using `--force` will clear all existing reservations. Use with caution in production. @@ -305,11 +308,11 @@ canhazgpu run --wait 30m --gpus 4 -- python train.py ``` **Behavior:** -1. Validates actual GPU availability using nvidia-smi +1. Validates actual device availability using the configured provider 2. Excludes GPUs that are in use without reservation 3. If GPUs unavailable, waits in queue (unless `--nonblock` is set) 4. Reserves the requested number of GPUs using MRU-per-user allocation (with LRU fallback) -5. Sets `CUDA_VISIBLE_DEVICES` to the allocated GPU IDs +5. Sets `CUDA_VISIBLE_DEVICES` for NVIDIA/AMD or `ASCEND_RT_VISIBLE_DEVICES` for Ascend to the allocated device IDs 6. Runs your command 7. Automatically releases GPUs when the command finishes 8. Maintains a heartbeat while running to keep the reservation active @@ -398,9 +401,13 @@ canhazgpu reserve --start 14:00 --end 16:00 --gpus 2 ``` **Important Note:** -Unlike the `run` command, `reserve` does NOT automatically set `CUDA_VISIBLE_DEVICES`. You can use `--short` for easy shell integration: +Unlike the `run` command, `reserve` does NOT automatically set the provider visibility variable. You can use `--short` for easy shell integration: ```bash +# NVIDIA or AMD export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --short) + +# Huawei Ascend +export ASCEND_RT_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --short) ``` **Use Cases:** @@ -822,7 +829,7 @@ When no remote hosts are configured, the dashboard shows the traditional single- All allocation commands (`run` and `reserve`) automatically: -1. **Scan for unreserved usage** using nvidia-smi +1. **Scan for unreserved usage** using the configured provider 2. **Exclude unreserved GPUs** from the available pool 3. **Hold back GPUs needed by scheduled bookings** during the window the reservation would cover 4. **Provide detailed error messages** if insufficient GPUs remain @@ -844,4 +851,4 @@ When multiple GPUs are available, the system uses **Most Recently Used per User* ### Status Integration -The `status` command shows comprehensive information about all reservation types and validates actual usage against reservations, making it easy to identify and resolve conflicts. \ No newline at end of file +The `status` command shows comprehensive information about all reservation types and validates actual usage against reservations, making it easy to identify and resolve conflicts. diff --git a/docs/configuration.md b/docs/configuration.md index dd2eb73..61e727a 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -51,7 +51,7 @@ guard: channels: ["process", "tty", "log"] log-file: "" exclude-users: ["root"] - exclude-commands: ["Xorg", "nvidia-smi", "amd-smi", "dcgm-exporter", "nvidia-persistenced"] + exclude-commands: ["Xorg", "nvidia-smi", "amd-smi", "npu-smi", "dcgm-exporter", "nvidia-persistenced"] notify-holder: true # Default settings for 'run' command @@ -241,4 +241,4 @@ web: host: "0.0.0.0" ``` -This configuration provides sensible defaults while allowing easy customization for different environments and use cases. \ No newline at end of file +This configuration provides sensible defaults while allowing easy customization for different environments and use cases. diff --git a/docs/dev-architecture.md b/docs/dev-architecture.md index c8299d4..80262ba 100644 --- a/docs/dev-architecture.md +++ b/docs/dev-architecture.md @@ -407,11 +407,12 @@ Current MRU-per-user allocation could be enhanced with: ### 2. GPU Provider System -The system supports multiple GPU providers through a unified interface: +The system supports multiple accelerator providers through a unified interface: **Available Providers:** - **NVIDIA**: Uses nvidia-smi for NVIDIA GPU management - **AMD**: Uses amd-smi (ROCm 5.7+) for AMD GPU management +- **Huawei Ascend**: Uses npu-smi and CANN logical device IDs - **Fake**: Simulated provider for development and testing without real GPUs **Provider Architecture:** @@ -575,4 +576,4 @@ type UsageRecord struct { - Large GPU count scenarios - Memory leak detection -This architecture provides a robust, scalable foundation for GPU resource management while maintaining simplicity and ease of deployment. \ No newline at end of file +This architecture provides a robust, scalable foundation for GPU resource management while maintaining simplicity and ease of deployment. diff --git a/docs/dev-contributing.md b/docs/dev-contributing.md index 2504262..e23175c 100644 --- a/docs/dev-contributing.md +++ b/docs/dev-contributing.md @@ -17,7 +17,7 @@ cd canhazgpu # System requirements # - Go 1.25+ # - Redis server -# - NVIDIA drivers with nvidia-smi +# - A supported accelerator tool (nvidia-smi, amd-smi, or npu-smi) # Go dependencies (automatic) go mod download @@ -451,7 +451,7 @@ func (e *AllocationEngine) AllocateGPUs(ctx context.Context, req *AllocationRequ ### 3. Adding New GPU Providers -The system already supports NVIDIA, AMD, and Fake providers. To add support for new GPU hardware: +The system already supports NVIDIA, AMD, Huawei Ascend, and Fake providers. To add support for new accelerator hardware: 1. **Implement the GPUProvider interface:** ```go @@ -602,7 +602,7 @@ Relates to #456 - Operating system and version - Go version (for development issues) - Redis version -- GPU driver version (NVIDIA or AMD) +- Accelerator driver version (NVIDIA, AMD, or Ascend) - Complete error messages - Steps to reproduce - Expected vs actual behavior @@ -616,7 +616,7 @@ Clear description of the bug - OS: Ubuntu 22.04 - Go: 1.23.0 (if building from source) - Redis: 7.0.0 -- GPU Provider: nvidia / amd +- GPU Provider: nvidia / amd / ascend - GPU Driver: 535.129.03 - canhazgpu version: 1.0.0 @@ -695,4 +695,4 @@ Contributors are recognized in: - Release notes - Documentation acknowledgments -Thank you for contributing to canhazgpu! Your efforts help make GPU resource management better for everyone. \ No newline at end of file +Thank you for contributing to canhazgpu! Your efforts help make GPU resource management better for everyone. diff --git a/docs/dev-testing.md b/docs/dev-testing.md index 17150e6..27c4b81 100644 --- a/docs/dev-testing.md +++ b/docs/dev-testing.md @@ -23,7 +23,7 @@ This guide explains how to run tests for canhazgpu and understand the testing in #### Integration Tests (Slower) - Run with: `make test` or `make test-integration` - Duration: 5-30 seconds per test -- Dependencies: Redis server, nvidia-smi (optional) +- Dependencies: Redis server and a supported provider tool (`nvidia-smi`, `amd-smi`, or `npu-smi`; optional) - Tests real system interactions ## Running Tests @@ -72,8 +72,8 @@ When running full tests (`make test`), these tests may take time: 2. **GPU Validation Tests** (5-10 seconds) - `TestDetectGPUUsage_Integration` - - Calls nvidia-smi command - - Logs: Indicates nvidia-smi availability + - Calls the available provider tool + - Logs: Indicates provider availability 3. **Heartbeat Manager Tests** (1-3 seconds) - `TestHeartbeatManager_Wait` @@ -83,7 +83,7 @@ When running full tests (`make test`), these tests may take time: 4. **GPU Allocation Tests** (2-10 seconds) - `TestAllocationEngine_AllocateGPUs_Structure` - - Combines Redis + nvidia-smi validation + - Combines Redis + provider validation - Logs: Indicates each phase ### Test Logging @@ -107,8 +107,8 @@ Integration tests include verbose logging to explain timing: - Tests automatically skip if unavailable - Uses database 15 (test database) -2. **nvidia-smi** (optional) - - Used for GPU detection tests +2. **Provider tool** (optional) + - `nvidia-smi`, `amd-smi`, or `npu-smi` is used for device detection tests - Tests gracefully handle missing command - Expected to fail on non-GPU systems @@ -126,7 +126,7 @@ SKIP: Redis not available for testing: dial tcp :6379: connect: connection refus ``` **Solution**: Start Redis server or run `make test-short` -### nvidia-smi Not Found +### Provider Tool Not Found ``` nvidia-smi not available or failed: exec: "nvidia-smi": executable file not found ``` @@ -208,4 +208,4 @@ make test make test-coverage ``` -This testing infrastructure ensures reliable GPU allocation while providing fast feedback during development. \ No newline at end of file +This testing infrastructure ensures reliable GPU allocation while providing fast feedback during development. diff --git a/docs/features-guard.md b/docs/features-guard.md index e532a3d..22db803 100644 --- a/docs/features-guard.md +++ b/docs/features-guard.md @@ -88,7 +88,7 @@ Safety rails: - **Escalation ladder**: SIGINT → SIGTERM → SIGKILL, `--kill-grace` apart, so a job gets the chance to shut down cleanly - **Circuit breaker**: at most `--max-kills-per-hour` terminations (default 3); beyond that the guard only warns, since a storm of kills is more likely a bug than a room full of offenders - **Unknown owners are never terminated**: if the process owner cannot be determined it might be a system process -- **Allow lists**: `--exclude-users` (default `root`) and `--exclude-commands` (default `Xorg,nvidia-smi,amd-smi,dcgm-exporter,nvidia-persistenced`) +- **Allow lists**: `--exclude-users` (default `root`) and `--exclude-commands` (default `Xorg,nvidia-smi,amd-smi,npu-smi,dcgm-exporter,nvidia-persistenced`) - **Dry run**: `--dry-run` records what would have happened, including in `canhazgpu violations` ## Avoiding false positives diff --git a/docs/features-idle-timeout.md b/docs/features-idle-timeout.md index 7253b5f..3390235 100644 --- a/docs/features-idle-timeout.md +++ b/docs/features-idle-timeout.md @@ -76,7 +76,7 @@ The same rules apply as for manual reservations: **Reservations without an idle timeout are exempt.** That includes reservations created before this feature existed, `canhazgpu run --idle-timeout 0`, and any reservation whose stored timeout is zero. They stay tied to the lifetime of their process. -**Reservations are only released when usage can actually be checked.** If `nvidia-smi`/`amd-smi` cannot be queried, idle detection is skipped for that pass rather than guessed at. The same applies to usage that cannot be attributed to an owner. +**Reservations are only released when usage can actually be checked.** If `nvidia-smi`, `amd-smi`, or `npu-smi` cannot be queried, idle detection is skipped for that pass rather than guessed at. The same applies to usage that cannot be attributed to an owner. **Memory counts as usage.** A GPU holding a loaded model with no active kernel is "in use" as far as canhazgpu is concerned — this matches how the rest of the tool defines usage. Lower `--memory-threshold` if you want stricter accounting. diff --git a/docs/features-validation.md b/docs/features-validation.md index 1b2a0e8..00d88a6 100644 --- a/docs/features-validation.md +++ b/docs/features-validation.md @@ -1,11 +1,14 @@ -# GPU Validation +# Device Validation -canhazgpu integrates with nvidia-smi to provide real-time validation of GPU usage, ensuring that reservations match actual resource utilization and detecting unreserved usage. +canhazgpu uses the configured provider (`nvidia-smi`, `amd-smi`, or `npu-smi`) +to validate actual device usage in real time, ensuring that reservations match +resource utilization and detecting unreserved usage. ## How Validation Works -### nvidia-smi Integration -The system uses nvidia-smi to query actual GPU processes and memory usage: +### Provider Integration +The NVIDIA provider queries actual GPU processes and memory usage with +`nvidia-smi`: ```bash # canhazgpu internally runs commands like: @@ -13,6 +16,11 @@ nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits nvidia-smi --query-compute-apps=pid,process_name,gpu_uuid,used_memory --format=csv,noheader ``` +AMD uses `amd-smi` JSON output. Huawei Ascend uses `npu-smi info` and its +logical NPU IDs. Ascend devices can report a persistent HBM baseline even +when idle, so canhazgpu derives Ascend memory usage from the process table +rather than treating that baseline as a workload. + ### Process Owner Detection For each GPU process, canhazgpu identifies the owner: @@ -125,9 +133,9 @@ Identifies stale reservations that could be released early to improve resource a ## Validation in Allocation ### Pre-Allocation Scanning -Before any GPU allocation, canhazgpu: +Before any device allocation, canhazgpu: -1. **Scans all GPUs** using nvidia-smi +1. **Scans all devices** using the configured provider 2. **Identifies unreserved usage** and excludes those GPUs 3. **Updates available GPU pool** with only truly available GPUs 4. **Proceeds with allocation** using the validated pool @@ -187,7 +195,7 @@ The default threshold is **100 MB** and can be adjusted with `--memory-threshold ### Caching and Efficiency - Validation runs only during allocation and status commands -- nvidia-smi queries are batched for efficiency +- Provider queries are batched for efficiency - Process information is gathered in parallel where possible ### Impact on System Performance @@ -217,20 +225,26 @@ Validation complements the heartbeat system: ## Troubleshooting Validation -### nvidia-smi Not Available +### Provider Tool Not Available ```bash Error: nvidia-smi command not found ``` -Ensure NVIDIA drivers are properly installed: +Ensure the management tool for the configured provider is available: ```bash -# Test nvidia-smi availability +# NVIDIA nvidia-smi -# If not available, install NVIDIA drivers -sudo apt install nvidia-driver-* # Ubuntu +# AMD +amd-smi list + +# Huawei Ascend +npu-smi info ``` +For Ascend permission failures, add the account to the CANN runtime group +(commonly `HwHiAiUser`) and start a new login session before retrying. + ### Permission Issues ```bash Warning: Could not determine owner for PID 12345 @@ -244,11 +258,12 @@ This may occur when: The system will still function but with less detailed process information. ### Memory Reporting Discrepancies -Different tools may report slightly different GPU memory usage: -- nvidia-smi vs. CUDA runtime memory reports +Different tools may report slightly different device memory usage: +- Provider tools vs. framework runtime memory reports - Shared memory vs. process-specific memory - Memory allocated vs. memory actually used -canhazgpu uses nvidia-smi reporting for consistency across all processes. +canhazgpu uses the configured provider's reporting consistently for all +processes. -GPU validation ensures that canhazgpu maintains accurate, real-time awareness of GPU resource utilization, preventing conflicts and enabling efficient resource sharing. \ No newline at end of file +GPU validation ensures that canhazgpu maintains accurate, real-time awareness of GPU resource utilization, preventing conflicts and enabling efficient resource sharing. diff --git a/docs/index.md b/docs/index.md index 2844a8d..f6b39c8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -41,7 +41,8 @@ canhazgpu web --port 8080 - ✅ **Heartbeat monitoring**: Detects crashed processes and reclaims GPUs - ✅ **Unreserved usage detection**: Identifies GPUs in use without proper reservations - ✅ **User accountability**: Shows which users are running unreserved processes -- ✅ **Real-time validation**: Uses nvidia-smi to verify actual GPU usage +- ✅ **Real-time validation**: Uses nvidia-smi, amd-smi, or npu-smi to verify actual device usage +- ✅ **Multi-provider support**: Supports NVIDIA, AMD, and Huawei Ascend devices - ✅ **Smart allocation**: Automatically excludes unreserved GPUs from allocation - ✅ **Usage reporting**: Track and analyze GPU usage patterns over time - ✅ **Web dashboard**: Real-time monitoring interface with status and reports @@ -101,4 +102,4 @@ canhazgpu web --port 8080 - **[Architecture](dev-architecture.md)** - System design overview - **[Contributing](dev-contributing.md)** - Contribution guidelines - **[Testing](dev-testing.md)** - Testing procedures -- **[Release Process](dev-release.md)** - Release management \ No newline at end of file +- **[Release Process](dev-release.md)** - Release management diff --git a/docs/installation.md b/docs/installation.md index 7fcaefa..f8367dd 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -7,6 +7,7 @@ - **GPUs** with appropriate management tools: - **NVIDIA GPUs**: nvidia-smi available - **AMD GPUs**: amd-smi available (ROCm 5.7+) + - **Huawei Ascend NPUs** (including 910B1): npu-smi available and the CANN runtime configured for the user - **System access** to `/proc` filesystem or `ps` command for user detection ## Dependencies @@ -71,6 +72,28 @@ amd-smi list If not installed, install ROCm drivers for your system: +### Huawei Ascend NPUs + +Ensure the CANN runtime can query the devices for the same account that will +run canhazgpu: + +```bash +npu-smi info +# Should display the NPU table and process table +``` + +On installations that restrict access to the CANN runtime group, an +administrator must add each user to the configured group. For the common +default group this is: + +```bash +sudo usermod -aG HwHiAiUser +``` + +The user must fully log out and start a new login session before the group is +effective. Verify with `id -nG` and then rerun `npu-smi info`. Check +`/etc/ascend_install.info` for a site-specific `UserGroup` value. + ## Install canhazgpu ### Option 1: Homebrew (Recommended) @@ -184,6 +207,8 @@ canhazgpu run --gpus 1 -- nvidia-smi -- # For NVIDIA # Shows nvidia-smi options canhazgpu run --gpus 1 -- amd-smi -- # For AMD # Shows amd-smi options +canhazgpu run --gpus 1 -- npu-smi -- # For Huawei Ascend +# Shows npu-smi options ``` ### Manual Installation @@ -216,4 +241,4 @@ Should show available commands. ## Next Steps - **[Quick Start Guide](quickstart.md)** - Initialize and start using canhazgpu -- **[Configuration](configuration.md)** - Set up defaults and customize behavior \ No newline at end of file +- **[Configuration](configuration.md)** - Set up defaults and customize behavior diff --git a/docs/quick-quick-start.md b/docs/quick-quick-start.md index e663c34..86a1e82 100644 --- a/docs/quick-quick-start.md +++ b/docs/quick-quick-start.md @@ -18,7 +18,7 @@ canhazgpu status -vv # 显示所有进程 canhazgpu run --gpus 1 -- python train.py ``` -`run` 会自动预约 GPU、设置 `CUDA_VISIBLE_DEVICES`、跑完自动释放。注意 `--` 必须写，它分隔 canhazgpu 的参数和你的命令。 +`run` 会自动预约设备、为 NVIDIA/AMD 设置 `CUDA_VISIBLE_DEVICES` 或为 Ascend 设置 `ASCEND_RT_VISIBLE_DEVICES`、跑完自动释放。注意 `--` 必须写，它分隔 canhazgpu 的参数和你的命令。 常用选项： @@ -57,6 +57,9 @@ canhazgpu cancel 45b590f7 # 排队中的直接出队；跑着的会被 SIGTERM # 预约 1 张 GPU 4 小时，并把卡号写进环境变量 export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 1 --duration 4h --short) +# 华为 Ascend +export ASCEND_RT_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 1 --duration 4h --short) + jupyter notebook ``` diff --git a/docs/quickstart-zh.md b/docs/quickstart-zh.md index 077f8a6..26b6b0f 100644 --- a/docs/quickstart-zh.md +++ b/docs/quickstart-zh.md @@ -22,6 +22,9 @@ canhazgpu admin --gpus $(nvidia-smi -L | wc -l) # AMD 主机 canhazgpu admin --gpus $(amd-smi list --json | jq 'length') + +# 华为 Ascend 主机（示例：8 张逻辑 NPU） +canhazgpu admin --gpus 8 --provider ascend ``` 不要在繁忙的主机上不带 `--force` 运行它；它会清空所有 reservation。 @@ -63,7 +66,7 @@ canhazgpu status --json | jq -r '.[] | select(.status == "AVAILABLE") | .gpu_id' ## 2. 跑任务：run -`run` 是启动任何 GPU 负载的首选方式：它负责预约 GPU、设置 `CUDA_VISIBLE_DEVICES`、运行你的命令，并在命令退出后自动释放。一定要用 `--` 分隔 canhazgpu 的参数和你的命令： +`run` 是启动任何加速器负载的首选方式：它负责预约设备、为 NVIDIA/AMD 设置 `CUDA_VISIBLE_DEVICES` 或为 Ascend 设置 `ASCEND_RT_VISIBLE_DEVICES`、运行你的命令，并在命令退出后自动释放。一定要用 `--` 分隔 canhazgpu 的参数和你的命令： ```bash canhazgpu run --gpus 1 -- python train.py @@ -99,7 +102,7 @@ canhazgpu run --gpus 2 --timeout 12h --note "bert-finetune" -- \ ## 3. 交互式工作：reserve -当你不只跑一条命令、而是需要一段时间的 GPU（notebook、调试、多步骤实验）时，用 `reserve`。它是 **manual** reservation：有时长、不会自动设置 `CUDA_VISIBLE_DEVICES`，直到过期、闲置被回收或你手动释放。 +当你不只跑一条命令、而是需要一段时间的设备（notebook、调试、多步骤实验）时，用 `reserve`。它是 **manual** reservation：有时长、不会自动设置 provider 对应的可见设备变量，直到过期、闲置被回收或你手动释放。 ```bash # 1 张 GPU，4 小时 @@ -111,6 +114,9 @@ canhazgpu reserve --gpu-ids 0,2 --duration 2h # 预约并一步设置环境变量 export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --duration 3h --short) +# 华为 Ascend +export ASCEND_RT_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --duration 3h --short) + jupyter notebook ``` diff --git a/docs/quickstart.md b/docs/quickstart.md index 4721a07..fa74824 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -22,6 +22,9 @@ canhazgpu admin --gpus $(nvidia-smi -L | wc -l) # AMD host canhazgpu admin --gpus $(amd-smi list --json | jq 'length') + +# Huawei Ascend host (example: eight logical NPUs) +canhazgpu admin --gpus 8 --provider ascend ``` Do not run this on a busy host without `--force`; it resets all reservations. @@ -63,7 +66,7 @@ canhazgpu status --json | jq -r '.[] | select(.status == "AVAILABLE") | .gpu_id' ## 2. Run a job: `run` -`run` is the recommended way to start anything GPU-heavy. It reserves GPUs, sets `CUDA_VISIBLE_DEVICES`, runs your command, and releases the GPUs when the command exits. Use the `--` separator so canhazgpu does not try to parse your command's flags: +`run` is the recommended way to start anything accelerator-heavy. It reserves devices, sets `CUDA_VISIBLE_DEVICES` for NVIDIA/AMD or `ASCEND_RT_VISIBLE_DEVICES` for Ascend, runs your command, and releases the devices when the command exits. Use the `--` separator so canhazgpu does not try to parse your command's flags: ```bash canhazgpu run --gpus 1 -- python train.py @@ -99,7 +102,7 @@ canhazgpu run --gpus 2 --timeout 12h --note "bert-finetune" -- \ ## 3. Reserve for interactive work: `reserve` -Use `reserve` when you need GPUs for a while without running a single command: notebooks, debugging, multi-step experiments. The reservation is **manual**: it has a duration, does not set `CUDA_VISIBLE_DEVICES` for you, and stays until it expires, goes idle, or you release it. +Use `reserve` when you need devices for a while without running a single command: notebooks, debugging, multi-step experiments. The reservation is **manual**: it has a duration, does not set the provider visibility variable for you, and stays until it expires, goes idle, or you release it. ```bash # 1 GPU for 4 hours @@ -111,6 +114,9 @@ canhazgpu reserve --gpu-ids 0,2 --duration 2h # Reserve and set the environment in one step export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --duration 3h --short) +# Huawei Ascend +export ASCEND_RT_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --duration 3h --short) + jupyter notebook ``` diff --git a/docs/usage-reserve.md b/docs/usage-reserve.md index 55b2fe2..3374c2e 100644 --- a/docs/usage-reserve.md +++ b/docs/usage-reserve.md @@ -129,9 +129,12 @@ Perfect for Jupyter notebooks, IPython sessions, or iterative model development: canhazgpu reserve --duration 4h # Note the GPU IDs from the output, e.g., "Reserved 1 GPU(s): [2]" -# Manually set CUDA_VISIBLE_DEVICES +# Manually set the NVIDIA/AMD visibility variable export CUDA_VISIBLE_DEVICES=2 +# On Huawei Ascend, use the CANN visibility variable instead +export ASCEND_RT_VISIBLE_DEVICES=2 + # Start Jupyter with the reserved GPU jupyter notebook @@ -146,9 +149,12 @@ Reserve GPUs while you prepare and test your batch jobs: canhazgpu reserve --gpus 2 --duration 2h # Note the GPU IDs from the output, e.g., "Reserved 2 GPU(s): [1, 3]" -# Manually set CUDA_VISIBLE_DEVICES +# Manually set the NVIDIA/AMD visibility variable export CUDA_VISIBLE_DEVICES=1,3 +# On Huawei Ascend, use the CANN visibility variable instead +export ASCEND_RT_VISIBLE_DEVICES=1,3 + # Test your scripts with the reserved GPUs python test_distributed.py @@ -191,7 +197,7 @@ canhazgpu release ## How Manual Reservations Work ### Allocation Process -1. **Validation**: Checks actual GPU usage with nvidia-smi +1. **Validation**: Checks actual device usage with the configured provider 2. **Conflict Detection**: Excludes GPUs in unreserved use 3. **LRU Selection**: Chooses least recently used GPUs 4. **Time-based Expiry**: Sets expiration time based on duration @@ -212,8 +218,11 @@ Reserved 2 GPU(s): [1, 3] for 4h 0m 0s 1 │ ● IN_USE │ alice │ 30s │ MANUAL │ expires in 3h 59m │ no usage detected │ - │ 0% 3 │ ● IN_USE │ alice │ 30s │ MANUAL │ expires in 3h 59m │ no usage detected │ - │ 0% -# Manually set CUDA_VISIBLE_DEVICES +# Manually set the NVIDIA/AMD visibility variable export CUDA_VISIBLE_DEVICES=1,3 + +# On Huawei Ascend, use the CANN visibility variable instead +export ASCEND_RT_VISIBLE_DEVICES=1,3 python your_script.py ``` @@ -314,13 +323,14 @@ canhazgpu release # Clean up immediately ### Shell Scripts -The `--short` flag outputs only the GPU IDs, making it easy to set `CUDA_VISIBLE_DEVICES` in scripts: +The `--short` flag outputs only device IDs. Set the variable for the configured +provider in scripts: ```bash #!/bin/bash set -e -# Reserve GPUs and set environment variable in one step +# Reserve NVIDIA/AMD GPUs and set the environment variable in one step export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --duration 3h --short) echo "Using GPUs: $CUDA_VISIBLE_DEVICES" @@ -344,7 +354,7 @@ import subprocess import os def reserve_gpus(count=1, duration="2h"): - """Reserve GPUs and set CUDA_VISIBLE_DEVICES""" + """Reserve NVIDIA/AMD GPUs and set CUDA_VISIBLE_DEVICES.""" result = subprocess.run([ "canhazgpu", "reserve", "--gpus", str(count), @@ -373,4 +383,4 @@ finally: release_gpus() ``` -Manual reservations provide fine-grained control over GPU allocation, making them perfect for interactive development and planned work sessions. \ No newline at end of file +Manual reservations provide fine-grained control over GPU allocation, making them perfect for interactive development and planned work sessions. diff --git a/docs/usage-run.md b/docs/usage-run.md index b5ad856..f2266ae 100644 --- a/docs/usage-run.md +++ b/docs/usage-run.md @@ -80,10 +80,10 @@ canhazgpu run --gpus 1 -- jupyter notebook --ip=0.0.0.0 --port=8888 When you run `canhazgpu run --gpus 2 -- python train.py`, here's what happens: -1. **GPU Validation**: Uses nvidia-smi to check actual GPU usage +1. **Device Validation**: Uses the configured provider to check actual device usage 2. **Conflict Detection**: Identifies GPUs in use without proper reservations 3. **Allocation**: Reserves 2 GPUs using the MRU-per-user strategy (with LRU fallback) -4. **Environment Setup**: Sets `CUDA_VISIBLE_DEVICES` to the allocated GPU IDs (e.g., "0,3") +4. **Environment Setup**: Sets the provider visibility variable to the allocated device IDs (e.g., "0,3") 5. **Command Execution**: Runs `python train.py` with the GPU environment 6. **Heartbeat**: Maintains reservation with periodic heartbeats while running 7. **Idle Watch**: Releases the reservation if no holder GPU usage is detected within the idle timeout @@ -91,12 +91,18 @@ When you run `canhazgpu run --gpus 2 -- python train.py`, here's what happens: ## Environment Variables -The `run` command automatically sets: +The `run` command automatically sets the visibility variable for the configured +provider: -- `CUDA_VISIBLE_DEVICES`: Comma-separated list of allocated GPU IDs -- Your command sees only the reserved GPUs as GPU 0, 1, 2, etc. +| Provider | Variable | +| --- | --- | +| NVIDIA or AMD | `CUDA_VISIBLE_DEVICES` | +| Huawei Ascend | `ASCEND_RT_VISIBLE_DEVICES` | -Example: If GPUs 1 and 3 are allocated, `CUDA_VISIBLE_DEVICES=1,3` is set, and your PyTorch code will see them as `cuda:0` and `cuda:1`. +Each value is a comma-separated list of allocated logical device IDs. For +example, NVIDIA/AMD jobs allocated devices 1 and 3 receive +`CUDA_VISIBLE_DEVICES=1,3`; Ascend jobs receive +`ASCEND_RT_VISIBLE_DEVICES=1,3`. ## Advanced Usage @@ -202,7 +208,7 @@ This indicates high contention. Try again in a few seconds. ### Resource Planning - **Estimate GPU needs**: Start with fewer GPUs and scale up if needed -- **Monitor memory usage**: Use `nvidia-smi` during training to optimize allocation +- **Monitor memory usage**: Use the provider tool during training (`nvidia-smi`, `amd-smi`, or `npu-smi info`) - **Test with small datasets**: Verify your code works before requesting many GPUs ### Command Structure @@ -253,9 +259,12 @@ WantedBy=multi-user.target ### Resource Usage ```bash -# Monitor GPU usage while job runs +# Monitor NVIDIA GPU usage while job runs watch -n 5 nvidia-smi +# On Huawei Ascend, inspect NPU usage instead +watch -n 5 npu-smi info + # Check heartbeat status canhazgpu status # Look for "last heartbeat" info ``` @@ -270,4 +279,4 @@ canhazgpu run --gpus 1 -- python train.py 2>&1 | tee training.log & tail -f training.log ``` -The `run` command provides a robust, automatic way to manage GPU reservations for your workloads while ensuring fair resource sharing across your team. \ No newline at end of file +The `run` command provides a robust, automatic way to manage GPU reservations for your workloads while ensuring fair resource sharing across your team. diff --git a/docs/usage-status.md b/docs/usage-status.md index 753cfd0..dc3e66a 100644 --- a/docs/usage-status.md +++ b/docs/usage-status.md @@ -107,7 +107,7 @@ For programmatic integration, use the `--json` or `-j` flag to get structured JS | `unreserved_users` | array | List of users with unreserved processes | | `process_info` | string | Process details for unreserved usage, including how long each process has been running | | `processes` | array | Processes currently using the GPU. Default: `pid` + `elapsed_seconds`; `process_name` is included with `-v` (max 2) or `-vv` (all) | -| `utilization_percent` | integer | GPU utilization reported by the provider (0-100, e.g. nvidia-smi `utilization.gpu`) | +| `utilization_percent` | integer | Device utilization reported by the provider (0-100, e.g. NVIDIA `utilization.gpu` or Ascend `AICore(%)`) | | `idle_timeout` | string | Idle timeout of a manual reservation (e.g. `15m`) | | `idle_for` | string | How long the holder has not used the GPU | | `booking_id` | string | Scheduled booking that created this reservation | @@ -116,7 +116,7 @@ For programmatic integration, use the `--json` or `-j` flag to get structured JS | `foreign_memory_mb` | integer | Memory those processes hold | | `error` | string | Error message (for ERROR status) | -The last column, **UTIL**, is the GPU utilization percentage reported by the provider (nvidia-smi `utilization.gpu` for NVIDIA; best-effort for AMD), from 0 to 100%. +The last column, **UTIL**, is the device utilization percentage reported by the provider (NVIDIA `utilization.gpu`, AMD best effort, or Ascend `AICore(%)`), from 0 to 100%. ## Status Information Explained @@ -518,4 +518,4 @@ else: print("Timeout waiting for GPUs") ``` -The `status` command is your primary tool for understanding GPU resource utilization, identifying conflicts, and coordinating with your team. Regular monitoring helps maintain efficient resource usage and prevents conflicts. \ No newline at end of file +The `status` command is your primary tool for understanding GPU resource utilization, identifying conflicts, and coordinating with your team. Regular monitoring helps maintain efficient resource usage and prevents conflicts. diff --git a/internal/cli/admin.go b/internal/cli/admin.go index 98cae7b..138a2ff 100644 --- a/internal/cli/admin.go +++ b/internal/cli/admin.go @@ -33,7 +33,7 @@ Use --force to reinitialize an existing pool (this will clear all reservations). func init() { adminCmd.Flags().IntP("gpus", "g", 0, "Number of GPUs available on this machine (required)") adminCmd.Flags().Bool("force", false, "Force reinitialization even if already initialized") - adminCmd.Flags().StringP("provider", "p", "", "GPU provider to use (nvidia, amd, or fake). If not specified, auto-detect available provider. Use 'fake' for development/testing without real GPUs") + adminCmd.Flags().StringP("provider", "p", "", "GPU provider to use (nvidia, amd, ascend, or fake). If not specified, auto-detect available provider. Use 'fake' for development/testing without real GPUs") if err := adminCmd.MarkFlagRequired("gpus"); err != nil { // This should not happen in practice, but handle it panic(fmt.Sprintf("Failed to mark gpus flag as required: %v", err)) @@ -59,18 +59,15 @@ func runAdmin(ctx context.Context, gpuCount int, force bool, explicitProvider st // Determine which provider to use var providerName string if explicitProvider != "" { - // Use explicitly specified provider - fmt.Printf("Using explicitly specified GPU provider: %s\n", explicitProvider) - - // Validate provider name - if explicitProvider != "nvidia" && explicitProvider != "amd" && explicitProvider != "fake" { - return fmt.Errorf("invalid provider '%s'. Valid providers are: nvidia, amd, fake", explicitProvider) + providerName = gpu.CanonicalProviderName(explicitProvider) + if providerName == "" { + return fmt.Errorf("invalid provider '%s'. Valid providers are: nvidia, amd, ascend, fake", explicitProvider) } + fmt.Printf("Using explicitly specified GPU provider: %s\n", providerName) // For fake provider, skip availability check - if explicitProvider == "fake" { + if providerName == "fake" { fmt.Println("Using fake GPU provider for development/testing") - providerName = explicitProvider } else { // Validate that the specified provider is available pm := gpu.NewProviderManager() @@ -78,17 +75,18 @@ func runAdmin(ctx context.Context, gpuCount int, force bool, explicitProvider st available := false for _, provider := range availableProviders { - if provider.Name() == explicitProvider { + if provider.Name() == providerName { available = true break } } if !available { - return fmt.Errorf("provider '%s' is not available on this system", explicitProvider) + if providerName == gpu.AscendProviderName { + return fmt.Errorf("provider '%s' is not available on this system; ensure the current user can run 'npu-smi info'", providerName) + } + return fmt.Errorf("provider '%s' is not available on this system", providerName) } - - providerName = explicitProvider } } else { // Auto-detect available provider @@ -97,7 +95,7 @@ func runAdmin(ctx context.Context, gpuCount int, force bool, explicitProvider st availableProviders := pm.GetAvailableProviders() if len(availableProviders) == 0 { - return fmt.Errorf("no GPU providers available (nvidia-smi, amd-smi not found)") + return fmt.Errorf("no GPU providers available (nvidia-smi, amd-smi, or npu-smi unavailable)") } if len(availableProviders) > 1 { diff --git a/internal/cli/reserve.go b/internal/cli/reserve.go index 21dcdfa..cf5b18c 100644 --- a/internal/cli/reserve.go +++ b/internal/cli/reserve.go @@ -69,9 +69,10 @@ Time formats supported by --start and --end: - +2h (relative to now) IMPORTANT: Unlike 'canhazgpu run', this command does NOT automatically set -CUDA_VISIBLE_DEVICES. After reserving, you must manually set the environment -variable based on the GPU IDs shown in the output: - export CUDA_VISIBLE_DEVICES=1,3 +the device visibility variable. After reserving, manually set the variable +for your provider based on the device IDs shown in the output: + export CUDA_VISIBLE_DEVICES=1,3 # NVIDIA or AMD + export ASCEND_RT_VISIBLE_DEVICES=1,3 # Ascend Example usage: canhazgpu reserve --gpus 2 --duration 4h @@ -82,7 +83,7 @@ Example usage: canhazgpu reserve --wait 30m --gpus 4 --duration 2h # Wait up to 30 minutes canhazgpu reserve --start 14:00 --end 16:00 --gpus 2 # Book a time slot canhazgpu reserve --start 'tomorrow 09:00' --duration 4h --gpus 8 - export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --short) # For scripting + export CUDA_VISIBLE_DEVICES=$(canhazgpu reserve --gpus 2 --short) # NVIDIA or AMD scripting The reserved GPUs must be manually released with 'canhazgpu release' or will automatically expire after the specified duration.`, @@ -218,6 +219,11 @@ func runReserve(ctx context.Context, opts reserveOptions) error { if err := client.Ping(ctx); err != nil { return fmt.Errorf("failed to connect to Redis: %v", err) } + providerName, err := client.GetAvailableProvider(ctx) + if err != nil { + return fmt.Errorf("failed to get cached provider information: %v", err) + } + visibleDevicesEnv := gpu.VisibleDevicesEnvVar(providerName) // Create allocation engine engine := gpu.NewAllocationEngine(client, config) @@ -258,7 +264,7 @@ func runReserve(ctx context.Context, opts reserveOptions) error { // Sort GPU IDs for consistent ordering in output and environment variable sort.Ints(allocatedGPUs) - // Build list for CUDA_VISIBLE_DEVICES + // Build the allocated device ID list for the provider visibility variable. ids := make([]string, len(allocatedGPUs)) for i, id := range allocatedGPUs { ids[i] = strconv.Itoa(id) @@ -279,7 +285,8 @@ func runReserve(ctx context.Context, opts reserveOptions) error { } fmt.Printf( - "\nRun the following command to run only on these GPUs:\nexport CUDA_VISIBLE_DEVICES=%s\n", + "\nRun the following command to run only on these GPUs:\nexport %s=%s\n", + visibleDevicesEnv, strings.Join(ids, ","), ) diff --git a/internal/cli/run.go b/internal/cli/run.go index a84a52f..72b03a6 100644 --- a/internal/cli/run.go +++ b/internal/cli/run.go @@ -31,12 +31,12 @@ func (e *ExitCodeError) Error() string { var runCmd = &cobra.Command{ Use: "run", - Short: "Reserve GPUs and run a command with CUDA_VISIBLE_DEVICES set", - Long: `Reserve GPUs and run a command with CUDA_VISIBLE_DEVICES automatically set. + Short: "Reserve GPUs and run a command with device visibility set", + Long: `Reserve GPUs and run a command with the provider device visibility variable automatically set. The command will: 1. Reserve the requested number of GPUs (or specific GPU IDs) -2. Set CUDA_VISIBLE_DEVICES to the allocated GPU IDs +2. Set CUDA_VISIBLE_DEVICES (NVIDIA/AMD) or ASCEND_RT_VISIBLE_DEVICES (Ascend) to the allocated device IDs 3. Run your command with full interactive terminal support 4. Automatically release GPUs when the command finishes 5. Maintain a heartbeat while running to keep the reservation active @@ -198,6 +198,11 @@ func runRun(ctx context.Context, gpuCount int, gpuIDs []int, timeoutStr string, _ = client.Close() return fmt.Errorf("failed to connect to Redis: %v", err) } + providerName, err := client.GetAvailableProvider(ctx) + if err != nil { + _ = client.Close() + return fmt.Errorf("failed to get cached provider information: %v", err) + } // Create allocation engine engine := gpu.NewAllocationEngine(client, config) @@ -308,14 +313,8 @@ func runRun(ctx context.Context, gpuCount int, gpuIDs []int, timeoutStr string, return fmt.Errorf("command not found: %s", command[0]) } - // Set up environment with CUDA_VISIBLE_DEVICES, replacing any existing value - var env []string - for _, e := range os.Environ() { - if !strings.HasPrefix(e, "CUDA_VISIBLE_DEVICES=") { - env = append(env, e) - } - } - env = append(env, fmt.Sprintf("CUDA_VISIBLE_DEVICES=%s", gpuListStr)) + // Set the provider-specific visibility variable, replacing only its existing value. + env := withVisibleDevicesEnv(os.Environ(), providerName, gpuListStr) // Exec the user's command - this replaces the current process // The supervisor will continue running and monitor our PID @@ -330,6 +329,18 @@ func runRun(ctx context.Context, gpuCount int, gpuIDs []int, timeoutStr string, return fmt.Errorf("failed to exec command: %v", err) } +func withVisibleDevicesEnv(environment []string, providerName string, deviceIDs string) []string { + variable := gpu.VisibleDevicesEnvVar(providerName) + prefix := variable + "=" + env := make([]string, 0, len(environment)+1) + for _, entry := range environment { + if !strings.HasPrefix(entry, prefix) { + env = append(env, entry) + } + } + return append(env, fmt.Sprintf("%s=%s", variable, deviceIDs)) +} + // buildSupervisorArgs builds the command line for the supervisor we spawn. The // Redis settings are passed explicitly: the supervisor is a fresh process, so // without them it would fall back to the default Redis instead of the one this diff --git a/internal/cli/run_env_test.go b/internal/cli/run_env_test.go new file mode 100644 index 0000000..26fc2ee --- /dev/null +++ b/internal/cli/run_env_test.go @@ -0,0 +1,35 @@ +package cli + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestWithVisibleDevicesEnv(t *testing.T) { + tests := []struct { + name string + base []string + provider string + want []string + }{ + { + name: "NVIDIA replaces CUDA visibility", + base: []string{"PATH=/bin", "CUDA_VISIBLE_DEVICES=7", "OTHER=value"}, + provider: "nvidia", + want: []string{"PATH=/bin", "OTHER=value", "CUDA_VISIBLE_DEVICES=1,3"}, + }, + { + name: "Ascend replaces Ascend visibility only", + base: []string{"PATH=/bin", "CUDA_VISIBLE_DEVICES=7", "ASCEND_RT_VISIBLE_DEVICES=4"}, + provider: "ascend", + want: []string{"PATH=/bin", "CUDA_VISIBLE_DEVICES=7", "ASCEND_RT_VISIBLE_DEVICES=1,3"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, withVisibleDevicesEnv(tt.base, tt.provider, "1,3")) + }) + } +} diff --git a/internal/cli/run_test.go b/internal/cli/run_test.go index c415903..f0cbdd6 100644 --- a/internal/cli/run_test.go +++ b/internal/cli/run_test.go @@ -27,9 +27,13 @@ func isAmdSmiAvailable() bool { return err == nil } +func isAscendSmiAvailable() bool { + return exec.Command("npu-smi", "info").Run() == nil +} + // isAnyGPUProviderAvailable checks if any GPU provider is available func isAnyGPUProviderAvailable() bool { - return isNvidiaSmiAvailable() || isAmdSmiAvailable() + return isNvidiaSmiAvailable() || isAmdSmiAvailable() || isAscendSmiAvailable() } // TestIsNvidiaSmiAvailable tests the helper function itself @@ -50,6 +54,11 @@ func TestIsAmdSmiAvailable(t *testing.T) { // since it depends on the test environment } +func TestIsAscendSmiAvailable(t *testing.T) { + available := isAscendSmiAvailable() + t.Logf("npu-smi availability: %v", available) +} + func TestRunCommand_FailureCleanup(t *testing.T) { if testing.Short() { t.Skip("Skipping integration test in short mode") @@ -138,7 +147,7 @@ func TestRunCommand_Structure(t *testing.T) { func TestRunRun_Validation(t *testing.T) { if !isAnyGPUProviderAvailable() { - t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi not found)") + t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi, npu-smi unavailable)") } tests := []struct { diff --git a/internal/gpu/allocation.go b/internal/gpu/allocation.go index a4cc4f0..692ee01 100644 --- a/internal/gpu/allocation.go +++ b/internal/gpu/allocation.go @@ -385,8 +385,8 @@ type GPUStatusInfo struct { ProcessInfo string Error string ModelInfo *ModelInfo `json:"model_info,omitempty"` // Detected AI model information - Provider string `json:"provider,omitempty"` // GPU provider (e.g., "NVIDIA", "AMD") - GPUModel string `json:"gpu_model,omitempty"` // GPU model (e.g., "H100", "RTX 4090") + Provider string `json:"provider,omitempty"` // Accelerator provider (e.g., "NVIDIA", "AMD", "Ascend") + GPUModel string `json:"gpu_model,omitempty"` // Device model (e.g., "H100", "MI300X", "910B1") Note string `json:"note,omitempty"` // Optional note describing the reservation purpose IdleTimeout time.Duration `json:"idle_timeout,omitempty"` IdleFor time.Duration `json:"idle_for,omitempty"` // How long the reservation has been without GPU usage diff --git a/internal/gpu/allocation_test.go b/internal/gpu/allocation_test.go index 54dd686..584e255 100644 --- a/internal/gpu/allocation_test.go +++ b/internal/gpu/allocation_test.go @@ -73,7 +73,7 @@ func TestAllocationEngine_GetGPUStatus_Structure(t *testing.T) { } if !isAnyGPUProviderAvailable() { - t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi not found)") + t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi, npu-smi unavailable)") } t.Log("Starting integration test - this may take time if Redis is not available") @@ -112,7 +112,7 @@ func TestAllocationEngine_AllocateGPUs_Structure(t *testing.T) { } if !isAnyGPUProviderAvailable() { - t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi not found)") + t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi, npu-smi unavailable)") } t.Log("Starting GPU allocation integration test - may take 10+ seconds") diff --git a/internal/gpu/ascend_provider.go b/internal/gpu/ascend_provider.go new file mode 100644 index 0000000..82692e3 --- /dev/null +++ b/internal/gpu/ascend_provider.go @@ -0,0 +1,277 @@ +package gpu + +import ( + "bufio" + "context" + "fmt" + "os/exec" + "strconv" + "strings" + + "github.com/russellb/canhazgpu/internal/types" +) + +const ( + // AscendProviderName is the canonical provider name stored in Redis. + AscendProviderName = "ascend" + + // AscendVisibleDevicesEnv is the CANN runtime variable that restricts a + // process to a set of logical Ascend device IDs. + AscendVisibleDevicesEnv = "ASCEND_RT_VISIBLE_DEVICES" +) + +// AscendProvider implements GPUProvider for Huawei Ascend NPUs using npu-smi. +// Despite the GPUProvider name, the reservation model applies equally to NPUs. +type AscendProvider struct{} + +// NewAscendProvider creates an Ascend NPU provider. +func NewAscendProvider() *AscendProvider { + return &AscendProvider{} +} + +// Name returns the provider name used in configuration and Redis. +func (a *AscendProvider) Name() string { + return AscendProviderName +} + +// IsAvailable verifies that npu-smi can query the devices for the current +// user. Checking only that the binary exists would incorrectly report success +// when the user is not in the Ascend runtime group. +func (a *AscendProvider) IsAvailable() bool { + return exec.Command("npu-smi", "info").Run() == nil +} + +// DetectGPUUsage reads device and process information from one npu-smi call. +func (a *AscendProvider) DetectGPUUsage(ctx context.Context) (map[int]*types.GPUUsage, error) { + usage, err := a.queryGPUUsage(ctx) + if err != nil { + return nil, err + } + + for _, npuUsage := range usage { + for i := range npuUsage.Processes { + process := &npuUsage.Processes[i] + user, err := getProcessOwner(process.PID) + if err != nil { + user = "unknown" + } + process.User = user + process.ElapsedSeconds = getProcessElapsedSeconds(process.PID) + npuUsage.Users[user] = true + } + } + + return usage, nil +} + +// GetGPUCount returns the number of logical Ascend devices reported by +// npu-smi. Logical IDs are the IDs accepted by ASCEND_RT_VISIBLE_DEVICES. +func (a *AscendProvider) GetGPUCount(ctx context.Context) (int, error) { + usage, err := a.queryGPUUsage(ctx) + if err != nil { + return 0, err + } + return len(usage), nil +} + +func (a *AscendProvider) queryGPUUsage(ctx context.Context) (map[int]*types.GPUUsage, error) { + output, err := exec.CommandContext(ctx, "npu-smi", "info").CombinedOutput() + if err != nil { + if details := strings.TrimSpace(string(output)); details != "" { + return nil, fmt.Errorf("npu-smi info failed: %w: %s", err, details) + } + return nil, fmt.Errorf("npu-smi info failed: %w", err) + } + + usage, err := parseAscendSMIOutput(string(output)) + if err != nil { + return nil, fmt.Errorf("failed to parse npu-smi info: %w", err) + } + return usage, nil +} + +type ascendDeviceInfo struct { + model string + utilizationPercent int +} + +// parseAscendSMIOutput parses the human-readable table emitted by +// "npu-smi info". npu-smi reports a persistent HBM allocation even when a +// device has no user process, so MemoryMB is derived from its process table +// instead of the HBM-Usage column. +func parseAscendSMIOutput(output string) (map[int]*types.GPUUsage, error) { + devices := make(map[int]*ascendDeviceInfo) + processes := make(map[int][]types.GPUProcessInfo) + inProcessTable := false + currentNPU := -1 + + scanner := bufio.NewScanner(strings.NewReader(output)) + scanner.Buffer(make([]byte, 1024), 1024*1024) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + lowerLine := strings.ToLower(line) + if strings.Contains(lowerLine, "process id") && strings.Contains(lowerLine, "process memory") { + inProcessTable = true + continue + } + + fields := splitSMITableRow(line) + if len(fields) == 0 { + continue + } + + if inProcessTable { + if npuID, process, ok := parseAscendProcessRow(fields); ok { + processes[npuID] = append(processes[npuID], process) + } + continue + } + + npuID, model, ok := parseAscendDeviceRow(fields) + if !ok { + continue + } + + if model != "" { + device, exists := devices[npuID] + if !exists { + device = &ascendDeviceInfo{} + devices[npuID] = device + } + device.model = model + currentNPU = npuID + continue + } + + // The detail row starts with the Chip ID, not the logical NPU ID. + // It belongs to the preceding model row. + if device, exists := devices[currentNPU]; exists { + if utilization, ok := parseAscendUtilization(fields); ok { + device.utilizationPercent = utilization + } + } + } + + if err := scanner.Err(); err != nil { + return nil, err + } + if len(devices) == 0 { + return nil, fmt.Errorf("no Ascend NPUs found in npu-smi output") + } + + usage := make(map[int]*types.GPUUsage, len(devices)) + for npuID, device := range devices { + npuProcesses := processes[npuID] + memoryMB := 0 + for _, process := range npuProcesses { + memoryMB += process.MemoryMB + } + + usage[npuID] = &types.GPUUsage{ + GPUID: npuID, + MemoryMB: memoryMB, + UtilizationPercent: device.utilizationPercent, + Processes: npuProcesses, + Users: make(map[string]bool), + Provider: "Ascend", + Model: device.model, + } + } + + return usage, nil +} + +func splitSMITableRow(line string) []string { + if !strings.HasPrefix(line, "|") || !strings.HasSuffix(line, "|") { + return nil + } + + parts := strings.Split(line, "|") + if len(parts) < 3 { + return nil + } + + fields := make([]string, 0, len(parts)-2) + for _, part := range parts[1 : len(parts)-1] { + fields = append(fields, strings.TrimSpace(part)) + } + return fields +} + +func parseAscendDeviceRow(fields []string) (int, string, bool) { + if len(fields) == 0 { + return 0, "", false + } + + parts := strings.Fields(fields[0]) + if len(parts) == 0 { + return 0, "", false + } + + npuID, err := strconv.Atoi(parts[0]) + if err != nil || npuID < 0 { + return 0, "", false + } + return npuID, strings.Join(parts[1:], " "), true +} + +func parseAscendUtilization(fields []string) (int, bool) { + if len(fields) == 0 { + return 0, false + } + + // In npu-smi 25.5 the last table cell combines AICore, memory, and HBM + // counters. AICore is its first number. The first cell is a Chip ID and + // must not be treated as utilization. + value, ok := firstNonNegativeInt(fields[len(fields)-1]) + if ok && value <= 100 { + return value, true + } + return 0, false +} + +func parseAscendProcessRow(fields []string) (int, types.GPUProcessInfo, bool) { + if len(fields) < 4 { + return 0, types.GPUProcessInfo{}, false + } + + // Current 25.5 output combines NPU and Chip in its first cell, followed + // by Process id, Process name, and Process memory. Older releases can use + // separate NPU and Chip cells, so accept both layouts. + npuID, ok := firstNonNegativeInt(fields[0]) + if !ok { + return 0, types.GPUProcessInfo{}, false + } + + pidIndex := 1 + if len(fields) >= 5 { + pidIndex = 2 + } + pid, ok := firstNonNegativeInt(fields[pidIndex]) + if !ok || pid == 0 { + return 0, types.GPUProcessInfo{}, false + } + + nameIndex := pidIndex + 1 + memoryIndex := nameIndex + 1 + if memoryIndex >= len(fields) { + return 0, types.GPUProcessInfo{}, false + } + memoryMB, _ := firstNonNegativeInt(fields[memoryIndex]) + + return npuID, types.GPUProcessInfo{ + PID: pid, + ProcessName: strings.TrimSpace(fields[nameIndex]), + MemoryMB: memoryMB, + }, true +} + +func firstNonNegativeInt(value string) (int, bool) { + for _, field := range strings.Fields(value) { + number, err := strconv.Atoi(field) + if err == nil && number >= 0 { + return number, true + } + } + return 0, false +} diff --git a/internal/gpu/ascend_provider_test.go b/internal/gpu/ascend_provider_test.go new file mode 100644 index 0000000..32904c2 --- /dev/null +++ b/internal/gpu/ascend_provider_test.go @@ -0,0 +1,70 @@ +package gpu + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ascendSMIOutput = ` ++------------------------------------------------------------------------------------------------+ +| npu-smi 25.5.0 Version: 25.5.0 | ++---------------------------+---------------+----------------------------------------------------+ +| NPU Name | Health | Power(W) Temp(C) Hugepages-Usage(page)| +| Chip | Bus-Id | AICore(%) Memory-Usage(MB) HBM-Usage(MB) | ++===========================+===============+====================================================+ +| 0 910B1 | OK | 92.4 48 0 / 0 | +| 0 | 0000:C1:00.0 | 73 0 / 0 3452 / 65536 | ++===========================+===============+====================================================+ +| 1 910B1 | OK | 97.8 50 0 / 0 | +| 0 | 0000:01:00.0 | 12 0 / 0 3443 / 65536 | ++===========================+===============+====================================================+ ++---------------------------+---------------+----------------------------------------------------+ +| NPU Chip | Process id | Process name | Process memory(MB) | ++===========================+===============+====================================================+ +| 0 0 | 2241871 | vLLM-OmniDiff | 8118 | ++===========================+===============+====================================================+ +| No running processes found in NPU 1 | ++===========================+===============+====================================================+ +` + +func TestParseAscendSMIOutput(t *testing.T) { + usage, err := parseAscendSMIOutput(ascendSMIOutput) + require.NoError(t, err) + require.Len(t, usage, 2) + + assert.Equal(t, 0, usage[0].GPUID) + assert.Equal(t, "910B1", usage[0].Model) + assert.Equal(t, "Ascend", usage[0].Provider) + assert.Equal(t, 73, usage[0].UtilizationPercent) + assert.Equal(t, 8118, usage[0].MemoryMB) + require.Len(t, usage[0].Processes, 1) + assert.Equal(t, 2241871, usage[0].Processes[0].PID) + assert.Equal(t, "vLLM-OmniDiff", usage[0].Processes[0].ProcessName) + assert.Equal(t, 8118, usage[0].Processes[0].MemoryMB) + + assert.Equal(t, 1, usage[1].GPUID) + assert.Equal(t, "910B1", usage[1].Model) + assert.Equal(t, 12, usage[1].UtilizationPercent) + assert.Zero(t, usage[1].MemoryMB, "persistent HBM baseline is not user process memory") + assert.Empty(t, usage[1].Processes) +} + +func TestAscendProviderNamesAndVisibility(t *testing.T) { + assert.Equal(t, AscendProviderName, NewAscendProvider().Name()) + + for _, providerName := range []string{"ascend", "npu", "ascend-npu", " ASCEND "} { + assert.Equal(t, AscendProviderName, CanonicalProviderName(providerName)) + assert.Equal(t, AscendVisibleDevicesEnv, VisibleDevicesEnvVar(providerName)) + } + assert.Equal(t, "nvidia", CanonicalProviderName(" NVIDIA ")) + assert.Equal(t, "CUDA_VISIBLE_DEVICES", VisibleDevicesEnvVar("amd")) + assert.Empty(t, CanonicalProviderName("unknown")) +} + +func TestNewProviderManagerFromNamesAscend(t *testing.T) { + manager := NewProviderManagerFromNames([]string{"npu"}) + require.Len(t, manager.providers, 1) + assert.IsType(t, &AscendProvider{}, manager.providers[0]) +} diff --git a/internal/gpu/guard.go b/internal/gpu/guard.go index 3039e60..ebd25a2 100644 --- a/internal/gpu/guard.go +++ b/internal/gpu/guard.go @@ -62,7 +62,7 @@ func DefaultGuardConfig() GuardConfig { KillGrace: types.DefaultGuardKillGrace, MaxKillsPerHour: types.DefaultGuardMaxKillsPerHour, ExcludeUsers: []string{"root"}, - ExcludeCommands: []string{"Xorg", "nvidia-smi", "amd-smi", "dcgm-exporter", "nvidia-persistenced"}, + ExcludeCommands: []string{"Xorg", "nvidia-smi", "amd-smi", "npu-smi", "dcgm-exporter", "nvidia-persistenced"}, Maintenance: true, NotifyHolder: true, } diff --git a/internal/gpu/provider.go b/internal/gpu/provider.go index 32229c6..05001ad 100644 --- a/internal/gpu/provider.go +++ b/internal/gpu/provider.go @@ -3,13 +3,15 @@ package gpu import ( "context" "fmt" + "strings" "github.com/russellb/canhazgpu/internal/types" ) -// GPUProvider defines the interface for GPU providers (NVIDIA, AMD, etc.) +// GPUProvider defines the interface for accelerator providers (NVIDIA, AMD, +// Ascend, etc.). The historical name is retained for API compatibility. type GPUProvider interface { - // Name returns the name of the provider (e.g., "nvidia", "amd") + // Name returns the name of the provider (e.g., "nvidia", "amd", "ascend") Name() string // IsAvailable checks if the provider's tools are available on the system @@ -33,6 +35,7 @@ func NewProviderManager() *ProviderManager { providers: []GPUProvider{ NewNVIDIAProvider(), NewAMDProvider(), + NewAscendProvider(), }, } } @@ -42,11 +45,13 @@ func NewProviderManagerFromNames(providerNames []string) *ProviderManager { var providers []GPUProvider for _, name := range providerNames { - switch name { + switch CanonicalProviderName(name) { case "nvidia": providers = append(providers, NewNVIDIAProvider()) case "amd": providers = append(providers, NewAMDProvider()) + case AscendProviderName: + providers = append(providers, NewAscendProvider()) case "fake": // Create with 0 GPUs; count will be set from Redis when used providers = append(providers, NewFakeProvider(0)) @@ -58,6 +63,31 @@ func NewProviderManagerFromNames(providerNames []string) *ProviderManager { } } +// CanonicalProviderName returns the name stored in Redis for a supported +// provider. Ascend aliases are accepted for operator convenience. +func CanonicalProviderName(providerName string) string { + normalized := strings.ToLower(strings.TrimSpace(providerName)) + switch normalized { + case "nvidia", "amd", "fake": + return normalized + case AscendProviderName, "npu", "ascend-npu": + return AscendProviderName + default: + return "" + } +} + +// VisibleDevicesEnvVar returns the runtime environment variable used to limit +// a launched process to its allocated devices. +func VisibleDevicesEnvVar(providerName string) string { + switch CanonicalProviderName(providerName) { + case AscendProviderName: + return AscendVisibleDevicesEnv + default: + return "CUDA_VISIBLE_DEVICES" + } +} + // NewProviderManagerWithFake creates a provider manager with a fake provider // having the specified GPU count func NewProviderManagerWithFake(gpuCount int) *ProviderManager { diff --git a/internal/gpu/test_helpers_test.go b/internal/gpu/test_helpers_test.go index 933e4cf..50ef93b 100644 --- a/internal/gpu/test_helpers_test.go +++ b/internal/gpu/test_helpers_test.go @@ -19,9 +19,15 @@ func isAmdSmiAvailable() bool { return err == nil } +// isAscendSmiAvailable checks that npu-smi can query the devices for the +// current user, rather than merely checking that its binary is on PATH. +func isAscendSmiAvailable() bool { + return NewAscendProvider().IsAvailable() +} + // isAnyGPUProviderAvailable checks if any GPU provider is available func isAnyGPUProviderAvailable() bool { - return isNvidiaSmiAvailable() || isAmdSmiAvailable() + return isNvidiaSmiAvailable() || isAmdSmiAvailable() || isAscendSmiAvailable() } // TestIsNvidiaSmiAvailable tests the helper function itself @@ -41,3 +47,8 @@ func TestIsAmdSmiAvailable(t *testing.T) { // This test just documents the current state, doesn't assert a specific value // since it depends on the test environment } + +func TestIsAscendSmiAvailable(t *testing.T) { + available := isAscendSmiAvailable() + t.Logf("npu-smi availability: %v", available) +} diff --git a/internal/gpu/validation_test.go b/internal/gpu/validation_test.go index 8339ca0..020b64f 100644 --- a/internal/gpu/validation_test.go +++ b/internal/gpu/validation_test.go @@ -83,14 +83,14 @@ func TestDetectGPUUsage_Integration(t *testing.T) { } t.Log("Starting GPU provider integration test - may take 5-10 seconds or timeout") - t.Log("This test uses the new GPU Provider system (NVIDIA/AMD)") + t.Log("This test uses the GPU provider system (NVIDIA/AMD/Ascend)") // Use the new GPU Provider system pm := NewProviderManager() availableProviders := pm.GetAvailableProviders() if len(availableProviders) == 0 { - t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi not found)") + t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi, npu-smi unavailable)") } t.Logf("Found %d available provider(s):", len(availableProviders)) diff --git a/internal/types/types.go b/internal/types/types.go index d4778a2..2a581bd 100644 --- a/internal/types/types.go +++ b/internal/types/types.go @@ -89,15 +89,15 @@ func (ft FlexibleTime) ToTime() time.Time { return ft.Time } -// GPUUsage represents actual GPU usage detected via nvidia-smi +// GPUUsage represents actual accelerator usage detected by a provider. type GPUUsage struct { GPUID int `json:"gpu_id"` MemoryMB int `json:"memory_mb"` UtilizationPercent int `json:"utilization_percent,omitempty"` // GPU utilization reported by the provider (0-100) Processes []GPUProcessInfo `json:"processes"` Users map[string]bool `json:"users"` - Provider string `json:"provider"` // "nvidia" or "amd" - Model string `json:"model"` // GPU model name (e.g., "H100", "RTX 4090") or "AMD" + Provider string `json:"provider"` // e.g., "NVIDIA", "AMD", or "Ascend" + Model string `json:"model"` // Device model name (e.g., "H100", "MI300X", or "910B1") } // GPUProcessInfo represents a process using a GPU From c302fba9c3895956e20709886aa082e6cf53932b Mon Sep 17 00:00:00 2001 From: specture724 Date: Tue, 1 Sep 2026 23:18:41 +0800 Subject: [PATCH 2/4] chore: add systemd guard service --- deploy/canhazgpu-guard.service | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 deploy/canhazgpu-guard.service diff --git a/deploy/canhazgpu-guard.service b/deploy/canhazgpu-guard.service new file mode 100644 index 0000000..65694fc --- /dev/null +++ b/deploy/canhazgpu-guard.service @@ -0,0 +1,14 @@ +[Unit] +Description=canhazgpu accelerator reservation guard +Wants=network-online.target redis.service +After=network-online.target redis.service + +[Service] +Type=simple +User=root +ExecStart=/usr/local/bin/canhazgpu guard --interval 15s +Restart=on-failure +RestartSec=5s + +[Install] +WantedBy=multi-user.target From a6cdb63c242ad8c5f02bbcfaef9e625afa5eb7d9 Mon Sep 17 00:00:00 2001 From: specture724 Date: Tue, 1 Sep 2026 23:37:14 +0800 Subject: [PATCH 3/4] test: isolate Redis integration databases --- internal/cli/run.go | 15 +++++--- internal/cli/run_test.go | 57 +++++++++++++++++-------------- internal/gpu/allocation_test.go | 10 +++--- internal/gpu/booking_test.go | 9 +++-- internal/gpu/heartbeat_test.go | 14 ++++---- internal/gpu/queue_test.go | 6 ++-- internal/gpu/test_helpers_test.go | 4 +++ 7 files changed, 65 insertions(+), 50 deletions(-) diff --git a/internal/cli/run.go b/internal/cli/run.go index 72b03a6..892d84a 100644 --- a/internal/cli/run.go +++ b/internal/cli/run.go @@ -160,10 +160,7 @@ func runRun(ctx context.Context, gpuCount int, gpuIDs []int, timeoutStr string, timeoutStr = "" } - // If neither is specified, default to 1 GPU - if gpuCount == 0 && len(gpuIDs) == 0 { - gpuCount = 1 - } + gpuCount = normalizeRunGPUCount(gpuCount, gpuIDs) config := getConfig() @@ -329,6 +326,16 @@ func runRun(ctx context.Context, gpuCount int, gpuIDs []int, timeoutStr string, return fmt.Errorf("failed to exec command: %v", err) } +// normalizeRunGPUCount applies the CLI default without changing a request for +// specific device IDs. Keeping this separate makes the behavior testable +// without invoking the exec-based run path. +func normalizeRunGPUCount(gpuCount int, gpuIDs []int) int { + if gpuCount == 0 && len(gpuIDs) == 0 { + return 1 + } + return gpuCount +} + func withVisibleDevicesEnv(environment []string, providerName string, deviceIDs string) []string { variable := gpu.VisibleDevicesEnvVar(providerName) prefix := variable + "=" diff --git a/internal/cli/run_test.go b/internal/cli/run_test.go index f0cbdd6..8685a87 100644 --- a/internal/cli/run_test.go +++ b/internal/cli/run_test.go @@ -15,6 +15,9 @@ import ( "github.com/stretchr/testify/require" ) +// cliTestRedisDB is isolated from Redis-backed tests in other Go packages. +const cliTestRedisDB = 13 + // isNvidiaSmiAvailable checks if nvidia-smi command is available func isNvidiaSmiAvailable() bool { _, err := exec.LookPath("nvidia-smi") @@ -70,7 +73,7 @@ func TestRunCommand_FailureCleanup(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, // Test database + RedisDB: cliTestRedisDB, } client := redis_client.NewClient(config) @@ -145,47 +148,49 @@ func TestRunCommand_Structure(t *testing.T) { assert.Equal(t, "", timeoutFlag.DefValue) } -func TestRunRun_Validation(t *testing.T) { - if !isAnyGPUProviderAvailable() { - t.Skip("Skipping test: no GPU providers available (nvidia-smi, amd-smi, npu-smi unavailable)") - } - +func TestNormalizeRunGPUCount(t *testing.T) { tests := []struct { name string gpuCount int - command []string - wantErr bool + gpuIDs []int + want int }{ { - name: "Zero GPU count (defaults to 1)", + name: "zero count defaults to one GPU", gpuCount: 0, - command: []string{"echo", "test"}, - wantErr: false, + want: 1, }, { - name: "Negative GPU count", - gpuCount: -1, - command: []string{"echo", "test"}, - wantErr: true, + name: "specific IDs leave count unchanged", + gpuCount: 0, + gpuIDs: []int{1, 3}, + want: 0, + }, + { + name: "explicit count is preserved", + gpuCount: 2, + want: 2, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - - err := runRun(ctx, tt.gpuCount, nil, "", "0", "", "", true, "", tt.command) - - if tt.wantErr { - assert.Error(t, err) - } else { - assert.NoError(t, err) - } + assert.Equal(t, tt.want, normalizeRunGPUCount(tt.gpuCount, tt.gpuIDs)) }) } } +func TestRunRequestRejectsNegativeGPUCount(t *testing.T) { + request := &types.AllocationRequest{ + GPUCount: normalizeRunGPUCount(-1, nil), + User: "testuser", + ActualUser: "testuser", + ReservationType: types.ReservationTypeRun, + } + + assert.Error(t, request.Validate()) +} + func TestExitCodeHandling(t *testing.T) { // Test that we can properly detect exit codes from failed commands // This tests the logic that was fixed to ensure cleanup happens @@ -223,7 +228,7 @@ func TestRunCommand_HeartbeatCleanup_Integration(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: cliTestRedisDB, } client := redis_client.NewClient(config) diff --git a/internal/gpu/allocation_test.go b/internal/gpu/allocation_test.go index 584e255..0503d62 100644 --- a/internal/gpu/allocation_test.go +++ b/internal/gpu/allocation_test.go @@ -15,7 +15,7 @@ func TestAllocationEngine_Structure(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, MemoryThreshold: types.MemoryThresholdMB, } redisClient := redis_client.NewClient(config) @@ -82,7 +82,7 @@ func TestAllocationEngine_GetGPUStatus_Structure(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, MemoryThreshold: types.MemoryThresholdMB, } redisClient := redis_client.NewClient(config) @@ -120,7 +120,7 @@ func TestAllocationEngine_AllocateGPUs_Structure(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, MemoryThreshold: types.MemoryThresholdMB, } redisClient := redis_client.NewClient(config) @@ -164,7 +164,7 @@ func TestAllocationEngine_ReleaseGPUs_Structure(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, MemoryThreshold: types.MemoryThresholdMB, } redisClient := redis_client.NewClient(config) @@ -416,7 +416,7 @@ func TestReleaseSpecificGPUs(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, MemoryThreshold: types.MemoryThresholdMB, } redisClient := redis_client.NewClient(config) diff --git a/internal/gpu/booking_test.go b/internal/gpu/booking_test.go index e7c5f8f..9f04169 100644 --- a/internal/gpu/booking_test.go +++ b/internal/gpu/booking_test.go @@ -182,15 +182,14 @@ func TestBookingShortIDAndStatus(t *testing.T) { // setupBookingTestEngine creates an allocation engine backed by a Redis test // database with a fake GPU pool of the requested size. // -// DB 14 is used rather than the usual DB 15 because these tests clear the state -// they touch, while other tests in this package rely on state left behind in -// DB 15. Only canhazgpu keys are removed - never the whole database - so -// pointing these tests at a shared Redis cannot destroy unrelated data. +// The gpu package uses a dedicated database because tests in other packages +// run concurrently. Only canhazgpu keys are removed here, never unrelated +// keys that may share this database. func setupBookingTestEngine(t *testing.T, gpuCount int) (*AllocationEngine, *redis_client.Client, context.Context) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 14, + RedisDB: gpuTestRedisDB, MemoryThreshold: types.MemoryThresholdMB, } diff --git a/internal/gpu/heartbeat_test.go b/internal/gpu/heartbeat_test.go index 8a69cf6..2c749e1 100644 --- a/internal/gpu/heartbeat_test.go +++ b/internal/gpu/heartbeat_test.go @@ -14,7 +14,7 @@ func TestHeartbeatManager_Structure(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } redisClient := redis_client.NewClient(config) @@ -38,7 +38,7 @@ func TestHeartbeatManager_StartStop(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } redisClient := redis_client.NewClient(config) @@ -111,7 +111,7 @@ func TestHeartbeatManager_Wait(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } redisClient := redis_client.NewClient(config) @@ -175,7 +175,7 @@ func TestHeartbeatManager_SendHeartbeat(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } redisClient := redis_client.NewClient(config) @@ -201,7 +201,7 @@ func TestHeartbeatManager_DoubleStop(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } redisClient := redis_client.NewClient(config) @@ -279,7 +279,7 @@ func TestHeartbeatManager_ReleaseGPUs(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } redisClient := redis_client.NewClient(config) @@ -328,7 +328,7 @@ func TestHeartbeatManager_ReservationLoss(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } client := redis_client.NewClient(config) diff --git a/internal/gpu/queue_test.go b/internal/gpu/queue_test.go index 3d8cbbf..87816a0 100644 --- a/internal/gpu/queue_test.go +++ b/internal/gpu/queue_test.go @@ -16,7 +16,7 @@ func setupQueueTestRedis(t *testing.T) *redis_client.Client { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, // Use test database + RedisDB: gpuTestRedisDB, } client := redis_client.NewClient(config) @@ -376,7 +376,7 @@ func TestQueueAllocatesFirstFullySatisfiableEntry(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, MemoryThreshold: 100, } engine := NewAllocationEngine(client, config) @@ -473,7 +473,7 @@ func TestQueueStatus(t *testing.T) { config := &types.Config{ RedisHost: "localhost", RedisPort: 6379, - RedisDB: 15, + RedisDB: gpuTestRedisDB, } engine := NewAllocationEngine(client, config) diff --git a/internal/gpu/test_helpers_test.go b/internal/gpu/test_helpers_test.go index 50ef93b..bc7a397 100644 --- a/internal/gpu/test_helpers_test.go +++ b/internal/gpu/test_helpers_test.go @@ -5,6 +5,10 @@ import ( "testing" ) +// gpuTestRedisDB is dedicated to Redis-backed tests in this package. Go runs +// packages concurrently, so it must not overlap with another package's test DB. +const gpuTestRedisDB = 14 + // isNvidiaSmiAvailable checks if nvidia-smi command is available // This is used by tests to skip tests that require nvidia-smi when it's not present func isNvidiaSmiAvailable() bool { From 83a0562c6f9f61831a8ae5154b9f8447517007b8 Mon Sep 17 00:00:00 2001 From: specture724 Date: Wed, 2 Sep 2026 00:15:18 +0800 Subject: [PATCH 4/4] chore: enforce guard policy --- deploy/canhazgpu-guard.service | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deploy/canhazgpu-guard.service b/deploy/canhazgpu-guard.service index 65694fc..3c24515 100644 --- a/deploy/canhazgpu-guard.service +++ b/deploy/canhazgpu-guard.service @@ -6,7 +6,7 @@ After=network-online.target redis.service [Service] Type=simple User=root -ExecStart=/usr/local/bin/canhazgpu guard --interval 15s +ExecStart=/usr/local/bin/canhazgpu guard --enforce --exclude-users "" --interval 1s --grace 5s --confirmations 1 --max-warnings 1 --warn-interval 5s --kill-grace 5s --max-kills-per-hour 0 --channels log --log-file /home/ajhou/.cache/gpu.log Restart=on-failure RestartSec=5s