AI Infrastructure Engineer

Pokee AI; United States; Remote; USD 140,000 - 220,000; Full time

Job summary

Pokee AI seeks an experienced ML infrastructure engineer to design and maintain the core platform that trains and serves reinforcement-learning agents at production scale. You will own scalable training pipelines, low-latency inference, and robust infrastructure across cloud and edge environments.

Collaborate with research scientists to productionize new algorithms, optimize latency and throughput on AWS or GCP, and ensure security/compliance for enterprise workloads.

Qualifications

• 3+ years in ML infrastructure or ML platform engineering.
• Strong Python and systems-level language experience.
• Hands-on ML serving frameworks expertise.
• Experience with container orchestration and cloud infra (AWS/GCP).
• Familiarity with experiment tracking and pipeline tools.

Responsibilities

• Architect and operate training and inference infrastructure for RL models.
• Tune serving for latency, throughput, and cost across cloud and edge deployments.
• Build and maintain CI/CD, experiment tracking, and model versioning.
• Develop data pipelines for training data collection and preprocessing.
• Collaborate with research to productionize new algorithms.

Skills

Python

Rust

C++

Go

ML serving frameworks

Kubernetes

Docker

AWS/GCP

GPU computing

MLflow

Weights & Biases

Airflow

Tools

NVIDIA TensorRT

ONNX Runtime

vLLM

Triton

CUDA

Job description

Role overview

Design and maintain the core platform that trains and serves reinforcement-learning-based AI agents at production scale. The role spans scalable training pipelines, low-latency inference, and infrastructure that runs reliably across both public cloud and on-premise or edge environments for enterprise customers.

Responsibilities

• Architect and operate training and inference infrastructure for RL-driven agent models, balancing scalability and reliability.
• Tune model serving for latency, throughput, and cost across AWS or GCP and on-device deployments.
• Build and maintain CI/CD pipelines, experiment tracking, and model versioning systems.
• Develop efficient data pipelines for training data collection, preprocessing, and reward signal computation.
• Partner with research scientists to productionize new algorithms and architectures.
• Ensure infrastructure meets enterprise requirements around reliability, security, and compliance (e.g., SOC 2, data residency).

Requirements

• 3+ years in ML infrastructure, ML platform engineering, or a closely related systems role.
• Strong proficiency in Python and a systems-level language such as Rust, C++, or Go.
• Hands-on experience with ML serving frameworks like vLLM, TensorRT, Triton, or ONNX Runtime.
• Experience with container orchestration (Kubernetes, Docker) and cloud infrastructure on AWS or GCP.
• Solid grasp of GPU computing, distributed systems, and performance profiling.
• Familiarity with experiment tracking and pipeline orchestration tools such as MLflow, Weights & Biases, or Airflow.

Nice to have

• Edge or on-device inference optimization (GGUF quantization, TensorRT-LLM, CoreML, QNN).
• On-premise GPU deployment experience with systems like NVIDIA DGX, Dell PowerEdge, or Lenovo ThinkStation.
• Background supporting RL training loops or online learning in production.
• Enterprise software experience with security and compliance frameworks.
• Contributions to open-source ML infrastructure projects.