Migration: Solving the Hard Problems

Our core principle for the migration was to replace the execution engine — Kubernetes to SageMaker — while keeping our ML workflows completely unchanged. The actual ML code — the Python scripts that train models, process data, and run inference — had to work identically on both platforms. No modifications to model training logic, no changes to data preprocessing, no updates to inference code.

Forcing hundreds of users across dozens of teams to rewrite their business-critical ML workflows was not an option. The cost of such a disruption in terms of lost productivity and engineering effort would have made the migration untenable, which meant the burden of compatibility was entirely on our platform. The requirement of zero code changes transformed the project into a complex systems engineering challenge for the ML Platform team, but it was a necessary one. The real task wasn’t just running a container on a different platform — it was ensuring environmental parity.

During the transition, we solved numerous challenges across the stack. Here are a few of the most complex ones we solved to make this possible.

Replicating the Kubernetes Runtime Environment
Our Kubernetes environment provided automatic credential injection via webhooks, metrics collection through sidecars, and configuration management via ConfigMaps. SageMaker offered none of these primitives. We built a compatibility layer into cross-platform base Docker images to replicate this behavior:

• Credentials: In Kubernetes, credentials from our internal secret management solution, Confidant , were automatically injected at pod creation. SageMaker has no equivalent mechanism. We built a custom solution, as part of the container entrypoint script, that fetches credentials at job startup and exposes them exactly as Kubernetes did, ensuring user code worked identically on both platforms
• Environment Variables: SageMaker constrains the number of environment variables passed via its API. Similar to our credential solution, we moved most environment setup to runtime, fetching additional configuration at job startup.
• Metrics: Kubernetes workloads sent StatsD metrics to sidecar containers. SageMaker has no sidecar support, so we reconfigured the runtime and networking to connect directly to our metrics aggregation gateway. The user-facing API remained unchanged.
• Hyperparameters: In Kubernetes, hyperparameters were stored in ConfigMaps and mounted as files. SageMaker’s API has much stricter size limits than K8s, making direct parameter passing impossible for our use cases. We developed a solution to upload hyperparameters to AWS S3 before each job and have SageMaker automatically download them to its standard input path. This overcame the API limitation while still using SageMaker’s native capabilities.

These represent only a subset of the environmental differences we systematically solved across the migration.

Building for the Hybrid Architecture
We developed new SageMaker-compatible base images to replace our old LyftLearn images. The critical design requirement was that these images must work across our entire hybrid platform: in SageMaker (for training and batch processing) and in Kubernetes (for serving). This meant the same Docker image that trained a model would also serve it, guaranteeing consistency. These base images serve as a foundation that teams extend with their own dependencies.

We built SageMaker-compatible base images with different capabilities to match our workload diversity. Here are some of the most important ones:

• LyftLearn image: For traditional ML workloads
• LyftLearn Distributed image: Adds Spark ecosystem integration for distributed processing
• LyftLearn DL image: Adds GPU support and libraries for deep learning workloads

The Spark-compatible images presented the biggest challenge. They needed to maintain full compatibility with our existing Spark infrastructure — custom wrappers, executor configurations, and JAR (Java Archive) dependencies. But they also had to run correctly in three distinct execution contexts: SageMaker Jobs, SageMaker Studio notebooks, and Model serving in K8s.

These images detect their execution environment at runtime and adapt. They automatically configure different environment variables, use different users and permissions, and set up Spark appropriately for each context, all while preserving an identical core runtime.

Matching Kubernetes Job Launch Times
In Kubernetes, notebooks, training, and processing jobs could start quickly because nodes were warm due to a significant percentage of cluster resources sitting idle. SageMaker provisions instances on-demand — no idle waste, but slower startup.

For JupyterLab notebooks, we adopted SOCI (Seekable Open Container Initiative) indexes. SOCI enables lazy loading: SageMaker fetches only the filesystem layers needed immediately rather than pulling entire multi-gigabyte images. This cut notebook startup times by 40–50%.

For training and batch processing jobs, SOCI wasn’t available. We optimized our Docker image sizes, which were sufficient for most of our workloads. However, this wasn’t enough for our most latency-sensitive workflows. Some models retrain every 15 minutes, making slower startup times unacceptable. For this subset of jobs, we adopted SageMaker’s warm pools, which keep instances alive between runs.

These optimizations gave us Kubernetes-like startup times with fully serverless infrastructure.

Cross-Cluster Networking for Spark
Many of our ML Platform users rely heavily on the interactive Spark experience in JupyterLab notebooks. In Kubernetes, this was simple, as the driver and executors ran in the same cluster. The new architecture, however, required the Spark driver to run in a SageMaker Studio notebook while the executors remained on our EKS K8s cluster.

This hybrid model presented a major networking challenge, as shown in the diagram below. Spark client mode requires bidirectional communication:

• The driver (in SageMaker) must call the EKS API Server Endpoint to request executor pods.
• The executor pods must be able to establish inbound connections directly back to the driver’s SageMaker Instance Elastic Network Interface (ENI).

The default SageMaker Studio networking blocked these critical inbound connections, breaking Spark’s communication model. This issue was a fundamental blocker that could jeopardize the entire migration. Without a solution for interactive Spark, we could not move our users to SageMaker Studio. To resolve this, we partnered closely with the AWS team. As a result of this collaboration, they introduced networking changes to the Studio Domains in our account that enabled the required inbound traffic from our EKS cluster. Despite the cross-cluster setup, Spark performance remained the same, and the interactive experience for ML Platform users was identical to the original Kubernetes environment.

LyftLearn 2.0: The Hybrid Architecture

As a result of this architectural transformation, we arrived at the hybrid architecture we planned: SageMaker for LyftLearn Compute and Kubernetes for LyftLearn Serving.

As the diagram illustrates, the two systems are fully decoupled, each operating as a purpose-built stack:

LyftLearn Serving runs on Kubernetes, powering a distributed architecture for real-time inference. Dozens of ML teams deploy their own model serving services — each containing their team’s models with custom prediction handlers and configurations — handling production predictions for specific use cases (pricing, fraud, dispatch, ETA, etc.). The Model Registry Service coordinates model deployments across these services. (We detailed this serving architecture in our 2023 blog post: Powering Millions of Real-Time Decisions with LyftLearn Serving.)

LyftLearn Compute runs on SageMaker, where the SageMaker Manager Service orchestrates training, batch processing, Hyperparameter Optimization (HPO), and JupyterLab notebooks through AWS SDK calls. EventBridge and SQS provide event-driven state management, replacing our background watchers.

Integration happens through the Model Registry and S3. Training jobs in SageMaker generate model binaries and save them to S3. The Model Registry tracks these artifacts, and model serving services pull them for deployment. Docker images flow from CI/CD through ECR to both platforms. The LyftLearn database maintains job metadata and model configurations across both stacks.

Each LyftLearn product operates independently while maintaining seamless end-to-end ML workflows.

Putting It All Together

We rolled out changes repository by repository, running both infrastructures in parallel. Our approach was systematic: build a comprehensive compatibility layer that made SageMaker feel like Kubernetes to ML code, validate each workflow type thoroughly, then migrate teams incrementally. Each repository required minimal changes — typically updating configuration files and workflow APIs — while the actual ML code remained untouched.

For our users, the migration was nearly invisible. But behind the scenes, the operational improvements were substantial. We reduced ML training and batch processing compute costs by eliminating idle cluster resources and moving to on-demand provisioning. System reliability improved significantly, with infrastructure-related incidents becoming rare occurrences. Most importantly, this stability and the serverless nature of the new compute freed our team to focus on building platform capabilities rather than managing low-level infrastructure components.