Staff / Principal ML Infrastructure Engineer

The Apple AI Platform team builds the foundational data storage, table format, and high-throughput data-loading systems powering Apple Foundation Models and large-scale ML training. We are designing and building an AI-native table format from the ground up: combining transactional snapshot isolation, concurrent updates, and multimodal schema evolution with built-in global indexing and sub-millisecond random access across multi-petabyte datasets on thousands of GPUs and TPUs.

As a Staff/Principal ML Infrastructure Engineer on the team, you will architect the on-disk format, indexing structures, transactional metadata layer, and Rust data plane and training data loader that Apple's largest training workloads run on.

Minimum Qualifications

• 10+ years of experience building storage engines, table or file formats, database internals, or high-performance distributed data and ML systems.
• Proven track record of designing and implementing table or file format internals (such as Apache Iceberg, Lance, Delta Lake, Parquet, Arrow, Vortex, Riegeli/ArrayRecord, RocksDB/LSM trees, or custom indexing and compaction engines), not just consuming them via high-level APIs.
• Production systems programming mastery in Rust or C++ (Rust strongly preferred for our core data plane), plus Python/PyO3 for ML framework bindings.
• Deep hands-on expertise in low-level I/O and memory performance engineering: mmap, zero-copy memory layouts, SIMD and vectorized decoding, RoaringBitmaps, async runtimes (tokio), and high-throughput S3 and GCS object-store client internals (range coalescing, connection pooling, tail-latency hedging).
• Strong understanding of transactional storage semantics: snapshot isolation, MVCC and optimistic concurrency control, manifest trees, and metadata scaling at millions of files and tens of billions of rows.
• Working knowledge of how large-scale distributed training (JAX/Grain, PyTorch) consumes data across GPU and TPU meshes (global shuffling, sequence packing, deterministic checkpoint and resume, host-to-device memory pipelines).
• Track record of setting technical direction across teams, shipping production systems at petabyte scale, and elevating senior engineers.
• B.S., M.S., or Ph.D. in Computer Science, Computer Engineering, or equivalent practical experience.

Preferred Qualifications

• Core contributor or committer experience on Apache Iceberg, Apache Arrow, Parquet, Lance, Vortex, DataFusion, Velox, DuckDB, Polars, Daft, or RocksDB.
• Experience designing or scaling ML training data loaders and formats (PyGrain, ArrayRecord/Riegeli, Megatron-LM .idx/.bin, Mosaic StreamingDataset, Ray Data, or NVIDIA DALI/AIStore) on multi-hundred-host GPU and TPU clusters.
• Experience designing multimodal dataset storage (combining pre-tokenized sequences, audio, image, and video blobs, and evolving feature or label columns without full dataset rewrites).
• Experience with formal correctness or high-rigor verification practices (differential parity testing, property-based fuzzing, or state-machine modeling for concurrent commits).