Make the code changes necessary to eventually dogfood v3. This includes: - removing `custom_sampling_context`