Principal Machine Learning Engineer

Okta; San Francisco (CA); Hybrid; USD 238,000 - 326,000; Full time

Benefits offered by this job

Equity

401(k)

Flexible spending account

Paid leave

Job summary

Okta is seeking a Principal Machine Learning Engineer to evolve authorization for agent access policies using dynamic AI security mechanisms. You will design scalable ML and Generative AI systems with retrieval, inference, and evaluation pipelines, while mentoring engineers and collaborating with product and security teams to enhance safety and reliability of AI-powered access controls.

The role requires 10+ years of software development with Python, deep ML experience, and hands-on work with

Qualifications

• 10+ years of software development experience with Python.
• Hands-on experience with applied machine learning, including feature engineering, training, and fine-tuning models.
• Experience with modern Generative AI platforms such as AWS Bedrock, OpenAI, Anthropic.
• Deep understanding of retrieval-augmented generation (RAG), embeddings, and knowledge-base workflows.
• Hands-on experience with AI agent frameworks including LiteLLM, LangGraph, LangChain, LlamaIndex, MCP.

Responsibilities

• Build intent-based enforcement to confirm agent runtime requests match intended purpose.
• Apply LLM reasoning and prompt parsing to interpret prompts, tool payloads, and intent in real time.
• Integrate low-latency inference or semantic evaluation engines into the API gateway request path.
• Use embeddings, vector search, or zero-shot classification to score alignment between agent intent and executed actions.
• Design confidence-scored decision engines that feed semantic verification results into Cedar policy frameworks.
• Establish evaluation benchmarks, prompt injection defenses, and guardrails to prevent bypasses or false positives.
• Architect scalable ML and Generative AI systems with retrieval, inference, and evaluation pipelines.
• Optimize prompting, context retrieval, and RAG workflows for Claude-based systems.
• Build automated evaluation pipelines to measure model quality, correctness, groundedness, and safety in production.
• Implement schema validation, structured output enforcement, and guardrails for reliable, compliant AI outputs.
• Mentor and coach engineers to support team and community growth.

Skills

Python

Go / TypeScript

Machine Learning

Retrieval-Augmented Generation

Embeddings

API development

Education

Bachelor's or Master’s degree in Computer Science or related field

Tools

AWS Bedrock

OpenAI

Anthropic

LiteLLM

LangGraph

LangChain

LlamaIndex

MCP

FastAPI

PyTorch

TensorFlow

Spark ML

Airflow

Cedar

Claude

Job description

Okta is hiring a Principal Machine Learning Engineer (Okta Secures AI) to evolve authorization for agent access policies using dynamic AI security mechanisms.

Responsibilities

• Build intent-based enforcement to confirm agent runtime requests match intended purpose
• Apply LLM reasoning and prompt parsing to interpret prompts, tool payloads, and intent in real time
• Integrate low-latency inference or semantic evaluation engines into the API gateway request path
• Use embeddings, vector search, or zero-shot classification to score alignment between agent intent and executed actions
• Design confidence-scored decision engines that feed semantic verification results into policy frameworks such as Cedar
• Establish evaluation benchmarks, prompt injection defenses, and guardrails to prevent bypasses or false positives
• Architect scalable ML and Generative AI systems with retrieval, inference, and evaluation pipelines
• Optimize prompting, context retrieval, and RAG workflows for accuracy, safety, and efficiency in Claude-based systems
• Build automated evaluation pipelines to measure model quality, correctness, groundedness, and safety in production
• Implement schema validation, structured output enforcement, and guardrails for reliable, compliant AI outputs
• Mentor and coach engineers to support team and community growth

Requirements

• 10+ years of software development experience with strong Python expertise; Go or TypeScript familiarity is a plus
• Hands-on experience with applied machine learning, including feature engineering, training, and fine-tuning models
• Hands-on experience with modern Generative AI platforms such as AWS Bedrock, OpenAI, Anthropic, etc.
• Deep understanding of retrieval-augmented generation (RAG), embeddings, and knowledge-base workflows
• Hands-on experience with AI agent frameworks including LiteLLM, LangGraph, LangChain, LlamaIndex, MCP, or similar
• Familiarity with ML frameworks such as FastAPI, PyTorch, TensorFlow, and Spark ML, plus orchestration tools like Airflow
• Experience defining evaluation metrics, pipelines, and feedback loops for ML/GenAI systems
• Proven ability to collaborate with product and engineering teams to drive greenfield initiatives, manage unknowns, and iterate quickly
• Experience building tools or infrastructure for AI/ML applications with deep understanding of the developer lifecycle in an AI-native world

Technologies

• Python, Go, TypeScript
• AWS Bedrock, OpenAI, Anthropic
• LiteLLM, LangGraph, LangChain, LlamaIndex, MCP
• FastAPI, PyTorch, TensorFlow, Spark ML
• Airflow, Cedar
• Claude-based systems

Location & Salary

• San Francisco, CA (hybrid)
• USD 238,000 - 326,000 per yearly

Education

• Bachelor’s or Master’s degree in Computer Science or related field

Benefits

• Equity (where applicable), bonus, and benefits including health, dental and vision insurance
• 401(k)
• Flexible spending account
• Paid leave including PTO and parental leave

Extra Credit

• Experience integrating AI-driven systems with identity, authentication, or security products
• Exposure to ethical AI, model risk, or compliance frameworks
• Familiarity with evaluation datasets, synthetic data generation, or LLM-as-a-judge methods

Okta Experience

• Supporting Your Well-Being
• Driving Social Impact
• Developing Talent and Fostering Connection + Community