Novas capacidades do Ray no SageMaker HyperPod

Ray e HyperPod: uma integração que simplifica o trabalho com IA distribuída

A AWS anunciou novas capacidades do Ray no Amazon SageMaker HyperPod, integrando o framework open-source diretamente com a infraestrutura especializada para treinamento e serving de modelos de fundação. O Ray é um framework open-source amplamente utilizado por cientistas de dados para escalar workloads Python distribuídos em clusters de GPUs — do treinamento distribuído com Ray Train ao serving de modelos com Ray Serve.

No Kubernetes, clusters Ray são gerenciados pelo KubeRay, um operador open-source que cuida do ciclo de vida dos clusters por meio de recursos customizados (RayCluster, RayJob, RayService). O SageMaker HyperPod, por sua vez, oferece infraestrutura otimizada para aprendizado de máquina (ML) em larga escala no Amazon Elastic Kubernetes Service (EKS), com monitoramento de saúde dos nós e recuperação automática integrados.

Antes dessa integração, rodar Ray no Kubernetes exigia que cientistas de dados escrevessem manifestos YAML, gerenciassem rebuilds de imagens Docker a cada mudança de dependência, configurassem kubectl port-forward para acessar o Ray Dashboard e instalassem Prometheus e Grafana manualmente para observabilidade. Com o novo lançamento, tudo isso passa a ser gerenciado diretamente pelo SageMaker Studio.

O que muda na prática

A partir de agora, cientistas de dados conseguem, dentro do SageMaker Studio: criar e gerenciar clusters Ray, abrir o Ray Dashboard e dashboards de observabilidade no Amazon Managed Grafana, conectar workspaces JupyterLab ou Code Editor ao cluster, submeter jobs distribuídos e configurar detecção de jobs travados (hung job detection).

No nível de aplicação, os jobs de treinamento Ray ganham tolerância automática a falhas por meio do monitoramento de saúde dos nós do HyperPod, além de checkpointing em camadas para retomada mais rápida. A integração com o SageMaker JumpStart permite carregar pesos de modelos diretamente em endpoints Ray Serve, com offloading de KV cache para armazenamento em camadas no serving de requisições com contexto longo. Tudo isso funciona com o KubeRay open-source e as APIs padrão do Ray — scripts e workflows existentes rodam sem modificação.

Pré-requisitos

Para usar essas capacidades, é necessário ter um cluster Amazon SageMaker HyperPod com orquestração Amazon EKS com os seguintes componentes instalados:

  • SageMaker Spaces EKS add-on: habilita workspaces JupyterLab e Code Editor no HyperPod que se conectam a clusters Ray para desenvolvimento interativo.
  • HyperPod Observability EKS add-on: coleta métricas de workloads Ray e provisiona dashboards Grafana no Amazon Managed Grafana.
  • KubeRay operator: gerencia RayCluster, RayJob e RayService como recursos nativos do Kubernetes.
  • HyperPod Ray Endpoint Operator (Helm chart): gera endpoints públicos autenticados para acesso ao dashboard e submissão remota de jobs.

Também é necessário um domínio SageMaker Studio, que fornece a interface de console para criar clusters Ray, visualizar workloads, abrir dashboards e gerenciar HyperPod Spaces. Para instruções completas de configuração, consulte o guia de introdução ao Ray no HyperPod.

Experiência no SageMaker Studio

O SageMaker Studio agora oferece um ambiente completo de desenvolvimento Ray. Cientistas de dados podem criar, gerenciar e monitorar clusters Ray diretamente pelo console, sem escrever manifestos Kubernetes ou executar comandos kubectl. É possível conferir a experiência completa neste demo interativo.

Para criar um cluster, basta acessar o HyperPod no Studio, selecionar o cluster, ir até a aba Tasks e escolher o tipo RayCluster. O formulário de criação solicita nome do cluster, tipos de instância para o nó principal (head) e workers, quantidade de workers e imagem de container. Por padrão, os clusters utilizam a imagem SageMaker Distribution, que já vem com Ray instalado e é mantida pela AWS com atualizações regulares de segurança. É possível também especificar uma imagem customizada para workloads com dependências adicionais.

Para quem prefere kubectl ou precisa de customizações avançadas, um editor YAML inline no Studio expõe o manifesto Kubernetes completo. O operador KubeRay também se integra ao HyperPod task governance, permitindo que administradores definam cotas de computação e prioridades de agendamento para workloads Ray junto com outros jobs de treinamento.

Durante a criação, é recomendável habilitar os endpoints remotos para acessar o Ray Dashboard, submeter jobs e recuperar logs de qualquer lugar com acesso à internet, de forma segura e sem necessidade de kubectl port-forward local. Saiba mais em acesso ao Ray Dashboard.

Submissão remota de jobs

Para workloads de produção, é possível submeter jobs remotamente para clusters Ray a partir do Studio, do laptop ou de pipelines de Integração Contínua e Entrega Contínua (CI/CD), usando o pacote Python toolkit-for-ray-on-sagemaker-ai. O pacote cuida da resolução de endpoints e geração de credenciais para a API do EKS via autenticação IAM, permitindo o uso das APIs padrão de submissão de jobs do Ray com um resolvedor de endereços compatível com o SageMaker:

$ aws eks update-kubeconfig --name <eks-cluster-name> --region <region>
$ pip install toolkit-for-ray-on-sagemaker-ai
$ ray job submit --address sagemaker_ray://<ray-cluster-name>/<namespace> \
    --working-dir <your-code-directory> \
    --python your-code.py
# To list ray jobs
$ ray job list --address sagemaker_ray://<ray-cluster-name>/<namespace>

Desenvolvimento interativo com SageMaker Spaces

Cientistas de dados podem conectar um cluster Ray a um espaço HyperPod JupyterLab ou Code Editor diretamente pelo Studio. O espaço entra no cluster como um nó worker de zero-computação, dando ao notebook acesso nativo completo como Ray driver. O cluster é selecionado em uma lista durante a criação do espaço.

Após criar o espaço e conectá-lo ao cluster Ray, basta abrir o JupyterLab ou Code Editor no navegador e chamar ray.init(address="auto") para estar conectado ao cluster e executar workloads distribuídos como se estivesse no nó principal.

O parâmetro runtime_env do Ray permite injetar dependências Python em tempo de execução, sem precisar reconstruir imagens de container. Workers podem ser escalados para cima ou para baixo sem recriar o cluster. Por exemplo, um cientista de dados pode começar prototipando em um notebook com um único worker e depois escalar para quatro workers GPU alterando apenas uma linha no ScalingConfig. O treinamento roda de forma distribuída no cluster Ray enquanto o notebook permanece interativo para monitoramento de progresso, ajuste de hiperparâmetros ou inspeção de resultados intermediários. Saiba mais em IDEs e Notebooks com Ray.

Imagem original — fonte: Aws

Observabilidade pronta para uso

Configurar observabilidade do Ray no Kubernetes com Amazon Managed Service for Prometheus e Amazon Managed Grafana era um processo de múltiplas etapas: instalar Helm charts, criar PodMonitors e ServiceMonitors, configurar roles IAM para assinatura SigV4 e importar arquivos JSON de dashboards manualmente.

O HyperPod Observability EKS add-on agora cuida de tudo isso automaticamente. Ele descobre os pods do nó principal e workers do Ray, coleta suas métricas e provisiona quatro dashboards Grafana pré-construídos no Amazon Managed Grafana: Ray Core, Ray Data, Ray Train e Ray Serve — todos organizados em uma pasta Ray e com suporte a filtragem por cluster específico. Os dashboards aparecem junto com os dashboards de infraestrutura do HyperPod (GPU, EFA, task governance), centralizando métricas de workloads Ray e saúde do cluster em um único lugar.

Treinamento resiliente

O SageMaker HyperPod oferece três camadas de resiliência para workloads de treinamento Ray: recuperação automática de nós em caso de falha de hardware, detecção de jobs travados quando o treinamento para de progredir, e checkpointing em camadas para recuperação rápida em ambos os cenários.

Recuperação automática de nós

O SageMaker HyperPod monitora continuamente a saúde dos nós e pode substituir nós com falha automaticamente, permitindo que jobs de treinamento Ray rodem pelo tempo necessário sem reinicializações manuais. Quando um nó é substituído, o Ray reagenda os pods worker no novo nó saudável. Se o código de treinamento salva checkpoints periodicamente e inclui lógica para retomar do último checkpoint, o job continua de onde parou — sem nenhuma alteração no código existente. Basta configurar um número suficiente de tentativas no FailureConfig do RayJob para que os jobs se recuperem automaticamente. Saiba mais em recuperação automática de nós com Ray.

Detecção de jobs travados

Jobs de treinamento distribuído podem travar sem gerar erros. Um único pod falha — por partição de rede, problema de montagem de armazenamento ou falha de hardware — e todos os outros pods ficam bloqueados na próxima operação coletiva, aguardando indefinidamente. As GPUs permanecem alocadas com memória carregada, mas sem produzir computação útil. Como não há mensagem de erro ou crash, cientistas de dados frequentemente descobrem o problema horas depois ao verificar o progresso manualmente. Em escala, algumas horas de tempo travado não detectado em dezenas de GPUs representam custo desperdiçado significativo.

O SageMaker HyperPod EKS inclui agora um Job Monitoring Agent por nó que detecta essas condições automaticamente para workloads Ray Train, sem alterações no código. O agente monitora múltiplos sinais em nível de nó e de job para determinar quando um job travou, notificando o usuário pelo grupo de logs do Amazon CloudWatch do cluster e pelo dashboard Ray Train Grafana provisionado pelo HyperPod Observability add-on. Para regras de detecção customizadas, é possível usar a biblioteca toolkit-for-ray-on-sagemaker-ai para definir padrões de log e limites de timeout. Quando a ação configurada é cancel, o HyperPod encerra o processo worker travado e o FailureConfig integrado do Ray Train reinicia os workers a partir do último checkpoint. Saiba mais em HyperPod Hung Job Detection no Ray.

Checkpointing em camadas

A biblioteca amzn-sagemaker-checkpointing se integra ao checkpointing em camadas gerenciado pelo HyperPod, que grava checkpoints em disco local e os envia de forma assíncrona para o Amazon Simple Storage Service (S3). Quando um job é reiniciado, a biblioteca verifica primeiro o HyperPod Tiered Storage. Se o checkpoint ainda estiver disponível lá, a recuperação é mais rápida do que restaurar diretamente do S3. Para modelos grandes, isso pode reduzir significativamente o tempo de recuperação. Saiba mais em HyperPod Tiered Storage no Ray.

Inferência acelerada

O Ray Serve é o framework do Ray para implantar modelos de ML como endpoints escaláveis e prontos para produção, com suporte a composição de múltiplos modelos, autoscaling e engines de serving como vLLM. O SageMaker HyperPod suporta workloads Ray Serve no EKS, permitindo implantar e escalar endpoints de inferência no cluster HyperPod.

O SageMaker JumpStart oferece um catálogo de modelos pré-treinados. Com esse lançamento, a biblioteca toolkit-for-ray-on-sagemaker-ai passa a incluir um JumpStart model loader que baixa pesos de modelos do catálogo JumpStart e os implanta diretamente no Ray Serve no HyperPod — sem download manual de pesos, configuração de modelo ou setup de container.

Um desafio comum no serving de Modelos de Linguagem de Grande Escala (LLMs) é que a latência de inferência cresce com o comprimento do contexto. Cada novo token exige recalcular a atenção sobre todos os tokens anteriores, tornando documentos longos e conversas multi-turno lentos e custosos. O SageMaker HyperPod endereça isso com o Managed Tiered KV Cache. O cache em camadas armazena vetores de chave-valor de atenção na memória CPU de cada nó (L1) e no HyperPod Tiered Storage para compartilhamento entre instâncias (L2). Os deployments Ray Serve podem aproveitar essas capacidades com mudanças mínimas de código, reduzindo o tempo até o primeiro token em conversas multi-turno e workloads de documentos longos. Saiba mais em Inferência Ray Acelerada no SageMaker HyperPod.

Limpeza dos recursos

Para liberar capacidade de computação no cluster, basta excluir os clusters Ray criados durante os testes pela aba Tasks do SageMaker Studio. Se um cluster HyperPod foi criado, é possível excluir o cluster pelo console do Amazon SageMaker AI para evitar cobranças pelas instâncias de computação subjacentes. Os add-ons de pré-requisito instalados para testes podem ser desinstalados pelo console do EKS para liberar capacidade de computação.

Disponibilidade

Essa integração está disponível hoje em todas as regiões AWS onde o SageMaker HyperPod EKS é suportado. Tudo funciona com o KubeRay open-source e as APIs padrão do Ray, sem necessidade de modificar scripts existentes. Para começar, consulte a documentação do Amazon SageMaker HyperPod e o guia de introdução ao Ray no HyperPod. É possível também explorar o fluxo completo neste demo interativo. Para discutir como o Ray no HyperPod pode suportar seus workloads, entre em contato com um representante AWS.

Fonte

Introducing new Ray capabilities on SageMaker HyperPod (https://aws.amazon.com/blogs/machine-learning/introducing-new-ray-capabilities-on-sagemaker-hyperpod/)

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *