Como um Pod no K8S utiliza GPU: Análise do código fonte do nvidia device plugin

Este artigo analisa o processo de como um Pod no Kubernetes solicita e utiliza recursos de GPU, explorando o mecanismo do device plugin e do nvidia-container-toolkit através do estudo do código fonte.

  1. Visão Geral

Quando um Pod é criado no Kubernetes solicitando recursos de GPU, o sistema precisa resolver duas questões fundamentais: como o cluster percebe a existência das GPUs e como esses recursos específicos são atribuídos a um contêiner.

Este processo envolve componentes especializaods do ecossistema NVIDIA que interagem com o Kubernetes:

  • Device Plugin: Responsável por descobrir e anunciar recursos de GPU ao kubelet
  • Container Toolkit: Gerencia a passagem dos dispositivos e bibliotecas necessárias para dentro do contêiner
  1. Fluxo de Trabalho

2.1 Descoberta de Recursos pelo Kubernetes

O Kubernetes implementa um mecanismo de device plugin que permite a drivers de hardware especiais se registrarem. A implementação da NVIDIA (NVIDIA/k8s-device-plugin) realiza duas tarefas principais:

Primeiro, detecta GPUs disponíveis no nó e as reporta ao kubelet, que atualiza o status do nó no API server. Isso permite ao scheduler tomar decisões informadas ao alocar Pods com requisitos de GPU.

Segundo, durante a alocação de recursos para um Pod, o plugin configura o contêiner com informações necessárias para acesso à GPU. A abordagem padrão utiliza variáveis de ambiente, mas alternativas baseadas em volumes ou anotações CDI também são suportadas.

2.2 Atribuição de GPUs a Contêineres

O nvidia-container-toolkit gerencia a integração real dos dispositivos de GPU no ambiente do contêiner. Esta ferramenta contém três componentes principais que formam a cadeia de execução:

O nvidia-container-runtime atua como um wrapper que modifica a especificação OCI do contêiner, injetando um hook de pré-inicialização que será executado pelo runtime de baixo nível (como o runc).

O nvidia-container-runtime-hook contém a lógica principal: extrai informações de dispositivos de GPU das variáveis de ambiente ou montagens do contêiner e invoca o utilitário de configuração final.

O nvidia-container-cli realiza a operação concreta de montar as bibliotecas e drivers necessários dentro do espaço de nomes do contêiner, garantindo acesso funcional às GPUs designadas.

  1. Análise do Código Fonte do Device Plugin

Examining the allocate function reveals how the plugin prepares container configurations:

// Exemplo reescrito: método de alocação simplificado
func (p *GPUPlugin) HandleAllocationRequest(ctx context.Context, req *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
    response := &pluginapi.AllocateResponse{}
    
    for _, containerReq := range req.ContainerRequests {
        // Validação dos dispositivos solicitados
        if err := p.validateRequestedDevices(containerReq.DevicesIDs); err != nil {
            return nil, err
        }
        
        containerResp, err := p.buildContainerResponse(containerReq.DevicesIDs)
        if err != nil {
            return nil, err
        }
        
        response.ContainerResponses = append(response.ContainerResponses, containerResp)
    }
    
    return response, nil
}

A resposta de alocação pode incluir múltiplas estratégias de configuração:

// Estratégias disponíveis para listar dispositivos
const (
    StrategyEnvironmentVariable = "envvar"
    StrategyVolumeMounts        = "volume-mounts"
    StrategyCDIAnnotations      = "cdi-annotations"
    StrategyCDICRI              = "cdi-cri"
)

func (p *GPUPlugin) buildContainerResponse(deviceIDs []string) (*pluginapi.ContainerAllocateResponse, error) {
    resp := &pluginapi.ContainerAllocateResponse{
        Envs: make(map[string]string),
    }
    
    // Aplica estratégia baseada em variável de ambiente
    if p.shouldUseEnvStrategy() {
        deviceList := strings.Join(deviceIDs, ",")
        resp.Envs["CUDA_VISIBLE_DEVICES"] = deviceList
    }
    
    // Adiciona anotações CDI se necessário
    if p.useCDIStrategy() {
        if err := p.addCDIAnnotations(resp, deviceIDs); err != nil {
            return nil, err
        }
    }
    
    return resp, nil
}

A escolha entre usar UUIDs ou índices numéricos para identificar dispositivos depende da configuração:

const (
    IdentificationByUUID   = "uuid"
    IdentificationByIndex  = "index"
)

// Configuração inicial do plugin
func NewPlugin(config *Config) *GPUPlugin {
    return &GPUPlugin{
        envVariableName:  "CUDA_VISIBLE_DEVICES",
        idStrategy:      config.DeviceIDStrategy,
        socketPath:      config.PluginSocketPath,
    }
}
  1. Análise do Código Fonte do Container Toolkit

4.1 Modificação da Especificação OCI

O container runtime da NVIDIA modifica o spec do contêiner para incluir hooks de pré-inicialização:

// Ponto de entrada do nvidia-container-runtime
func main() {
    rt := runtime.NewRuntime()
    if err := rt.ExecuteWithHookInjection(os.Args); err != nil {
        os.Exit(1)
    }
}

// Método que injeta o hook no spec
func (r *Runtime) injectNVIDIAHook(containerSpec *specs.Spec) error {
    if r.hasExistingHook(containerSpec) {
        return nil
    }
    
    hookPath := r.resolveHookBinaryPath()
    containerSpec.Hooks.Prestart = append(containerSpec.Hooks.Prestart, specs.Hook{
        Path: hookPath,
        Args: []string{filepath.Base(hookPath), "prestart"},
    })
    
    return nil
}

4.2 Processamento do Hook de Pré-inicialização

O hook executa durante a criação do contêiner e segue esta sequência:

func handlePrestartHook() {
    hookConfig, err := loadHookConfiguration()
    if err != nil {
        log.Fatal("Falha ao carregar configuração:", err)
    }
    
    containerState := parseContainerState()
    ociSpec := loadOCISpec(containerState.BundlePath)
    
    gpuDevices := resolveGPUDevices(ociSpec.Process.Env, ociSpec.Mounts)
    if len(gpuDevices) == 0 {
        return // Não é um contêiner que requer GPU
    }
    
    configureCommand := buildCLICommand(gpuDevices, hookConfig)
    executeConfiguration(configureCommand)
}

A resolução de dispositivos segue uma hierarquia de preferências:

func resolveGPUDevices(envVars []string, mounts []Mount) []string {
    // 1. Tenta obter de montagens de volumes
    if devices := getDevicesFromVolumeMounts(mounts); len(devices) > 0 {
        return devices
    }
    
    // 2. Tenta obter de variáveis de ambiente
    if devices := getDevicesFromEnvironment(envVars); len(devices) > 0 {
        return devices
    }
    
    // 3. Fallback para imagens legado
    if isLegacyCUDAImage(envVars) {
        return []string{"all"}
    }
    
    return nil
}

func getDevicesFromEnvironment(envVars []string) []string {
    for _, env := range envVars {
        if strings.HasPrefix(env, "CUDA_VISIBLE_DEVICES=") {
            deviceString := strings.TrimPrefix(env, "CUDA_VISIBLE_DEVICES=")
            return strings.Split(deviceString, ",")
        }
    }
    return nil
}

4.3 Execução do CLI de Configuração

O componente final realiza a montagem dos recursos necessários:

// Simulação da estrutura de comandos do nvidia-container-cli
type CLICommand struct {
    Executable string
    Action     string
    Devices    []string
    Capabilities []string
    PID        int
    RootfsPath string
}

func buildCLICommand(devices []string, config *HookConfig) *CLICommand {
    cmd := &CLICommand{
        Executable: config.CLIPath,
        Action:     "configure",
        Devices:    devices,
        PID:        config.ContainerPID,
        RootfsPath: config.RootfsPath,
    }
    
    if config.DriverCapabilities != "" {
        cmd.Capabilities = strings.Split(config.DriverCapabilities, ",")
    }
    
    return cmd
}

func (cmd *CLICommand) Execute() error {
    args := []string{cmd.Action}
    
    // Adiciona argumentos para dispositivos
    for _, dev := range cmd.Devices {
        args = append(args, "--device="+dev)
    }
    
    // Adiciona capacidades do driver
    for _, cap := range cmd.Capabilities {
        args = append(args, "--capability="+cap)
    }
    
    args = append(args, "--pid="+strconv.Itoa(cmd.PID))
    args = append(args, cmd.RootfsPath)
    
    return syscall.Exec(cmd.Executable, args, os.Environ())
}

O processo de montagem final utiliza a técnica de bind mount para inserir bibliotecas específicas no contêiner:

// Exemplo conceitual de como as bibliotecas são montadas
func mountDriverLibraries(containerPID int, capabilities []string) error {
    requiredLibraries := resolveRequiredLibraries(capabilities)
    
    for _, lib := range requiredLibraries {
        hostPath := "/usr/lib/x86_64-linux-gnu/" + lib
        containerPath := "/usr/lib/x86_64-linux-gnu/" + lib
        
        if err := bindMount(hostPath, containerPath, containerPID); err != nil {
            return err
        }
    }
    
    return nil
}
  1. Síntese do Processo

O fluxo completo de alocação de GPU em um Pod pode ser resumido nestas etapas:

  1. Descoberta: O device plugin registra GPUs disponíveis nos nós do cluster
  2. Agendamento: O scheduler do Kubernetes coloca o Pod em um nó com recursos suficientes
  3. Configuração: O plugin aloca GPUs específicas e configura variáveis de ambiente no contêiner
  4. Inicialização: O runtime do contêiner injeta hooks de configuração durante a inicialização
  5. Montagem: As bibliotecas e drivers necessários são montados no filesystem do contêiner

Uma questão comum resolvida por esta implementação: Pods sem requisição explícita de GPU podem ainda assim ter acesso a todas as GPUs disponíveis se a imagem do contêiner for reconhecida como "legacy" pelo toolkit.

Tags: kubernetes GPU NVIDIA device-plugin container-runtime

Publicado em 7-6 22:32